package artifact

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"neuralhunt/internal/settings"
)

type Worker struct {
	db         *sql.DB
	dir        string
	publicBase string
	settings   *settings.Manager
	http       *http.Client
}

func New(db *sql.DB, dir string, sm *settings.Manager) (*Worker, error) {
	if dir == "" {
		dir = "./data/artifacts"
	}
	abs, err := filepath.Abs(dir)
	if err != nil {
		return nil, err
	}
	if err := os.MkdirAll(abs, 0o750); err != nil {
		return nil, err
	}
	return &Worker{
		db:         db,
		dir:        abs,
		publicBase: strings.TrimRight(os.Getenv("ARTIFACT_PUBLIC_BASE_URL"), "/"),
		settings:   sm,
		http:       &http.Client{Timeout: envDuration("ARTIFACT_HTTP_TIMEOUT", 4*time.Minute)},
	}, nil
}

func envDuration(k string, d time.Duration) time.Duration {
	if raw := strings.TrimSpace(os.Getenv(k)); raw != "" {
		if v, err := time.ParseDuration(raw); err == nil {
			return v
		}
	}
	return d
}

type win struct {
	ID, Seed, Winner, Signature, Guess string
	DisplayName                        string
	RangeBits                          int
	Completed                          time.Time
	PublicJWK                          json.RawMessage
	PromptInstructions                 string
	NegativePrompt                     string
	StyleReference                     string
}

func (w *Worker) Run(ctx context.Context) {
	t := time.NewTicker(3 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			_ = w.one(ctx)
		}
	}
}

func (w *Worker) claim(ctx context.Context) (win, error) {
	tx, err := w.db.BeginTx(ctx, nil)
	if err != nil {
		return win{}, err
	}
	defer tx.Rollback()
	var x win
	var completedMS int64
	var raw string
	err = tx.QueryRowContext(ctx, `SELECT t.id,t.public_seed,t.winner_client_id,t.winner_signature,t.winning_guess,t.display_name,t.range_bits,t.completed_at,c.public_jwk,t.nft_prompt_instructions,t.nft_negative_prompt,t.nft_style_reference
		FROM tasks t JOIN clients c ON c.id=t.winner_client_id
		WHERE t.artifact_status='pending' ORDER BY t.completed_at LIMIT 1`).Scan(&x.ID, &x.Seed, &x.Winner, &x.Signature, &x.Guess, &x.DisplayName, &x.RangeBits, &completedMS, &raw, &x.PromptInstructions, &x.NegativePrompt, &x.StyleReference)
	if err != nil {
		return win{}, err
	}
	x.Completed = time.UnixMilli(completedMS).UTC()
	x.PublicJWK = json.RawMessage(raw)
	if _, err = tx.ExecContext(ctx, `UPDATE tasks SET artifact_status='generating',artifact_error=NULL WHERE id=? AND artifact_status='pending'`, x.ID); err != nil {
		return win{}, err
	}
	if err = tx.Commit(); err != nil {
		return win{}, err
	}
	return x, nil
}

type imageResult struct {
	Bytes            []byte
	Ext              string
	Provider         string
	Meta             map[string]any
	Usage            *imageUsage
	EstimatedCostUSD *float64
	PricingBasis     string
}

func (w *Worker) one(ctx context.Context) error {
	x, err := w.claim(ctx)
	if errors.Is(err, sql.ErrNoRows) {
		return nil
	}
	if err != nil {
		return err
	}
	cfg := w.settings.Get()
	preset := strings.ToLower(strings.TrimSpace(cfg.ArtifactPreset))
	traits := deriveCollectionTraits(x)
	prompt := ""
	negativePrompt := ""
	if preset == collectionPresetRaccoon {
		prompt = buildCollectionPrompt(x, traits)
		negativePrompt = collectionNegativePrompt(x)
	} else {
		prompt = buildPrompt(cfg.ArtifactPrompt, x.PromptInstructions, x)
		negativePrompt = joinPrompt(cfg.ArtifactNegativePrompt, x.NegativePrompt)
	}
	img, err := w.generate(ctx, cfg, x, prompt, negativePrompt)
	if err != nil {
		w.fail(ctx, x.ID, err)
		return err
	}
	if len(img.Bytes) == 0 {
		err := errors.New("image provider returned empty output")
		w.fail(ctx, x.ID, err)
		return err
	}
	if img.Ext == "" {
		img.Ext = "png"
	}

	idSum := sha256.Sum256([]byte(x.ID + "|" + x.Winner + "|" + x.Signature))
	artifactID := "artifact_" + hex.EncodeToString(idSum[:12])
	promptSum := sha256.Sum256([]byte(prompt))
	rawArtSum := sha256.Sum256(img.Bytes)

	finalBytes := img.Bytes
	finalExt := strings.TrimPrefix(strings.ToLower(img.Ext), ".")
	if preset == collectionPresetRaccoon {
		finalBytes = renderCardSVG(img.Bytes, img.Ext, x, traits)
		finalExt = "svg"
	}
	finalSum := sha256.Sum256(finalBytes)

	manifest := map[string]any{
		"artifact_id":              artifactID,
		"artifact_preset":          preset,
		"task_id":                  x.ID,
		"task_display_name":        x.DisplayName,
		"task_range_bits":          x.RangeBits,
		"winner_client_id":         x.Winner,
		"winner_public_jwk":        json.RawMessage(x.PublicJWK),
		"winning_guess":            x.Guess,
		"winner_guess_signature":   x.Signature,
		"completed_at":             x.Completed,
		"image_sha256":             hex.EncodeToString(finalSum[:]),
		"raw_art_sha256":           hex.EncodeToString(rawArtSum[:]),
		"prompt_sha256":            hex.EncodeToString(promptSum[:]),
		"task_prompt_instructions": x.PromptInstructions,
		"task_style_reference":     x.StyleReference,
		"provider":                 img.Provider,
		"provider_meta":            img.Meta,
		"note":                     "winner_guess_signature authenticates the winning guess; image_sha256 binds the final programmatically laid-out collectible card into the server manifest",
	}
	if preset == collectionPresetRaccoon {
		manifest["collection_character"] = "RIFT"
		manifest["collection_traits"] = traits
		manifest["layout"] = map[string]any{
			"format": "svg",
			"width":  1024,
			"height": 1536,
			"mode":   "programmatic-full-art-card-v1",
		}
	}
	mb, _ := json.MarshalIndent(manifest, "", "  ")

	relDir := filepath.Join("artifacts", artifactID)
	outDir := filepath.Join(w.dir, artifactID)
	if err := os.MkdirAll(outDir, 0o750); err != nil {
		w.fail(ctx, x.ID, err)
		return err
	}
	if preset == collectionPresetRaccoon {
		rawName := "art." + strings.TrimPrefix(strings.ToLower(img.Ext), ".")
		if err := atomicWrite(filepath.Join(outDir, rawName), img.Bytes, 0o640); err != nil {
			w.fail(ctx, x.ID, err)
			return err
		}
	}
	imageName := "image." + finalExt
	if err := atomicWrite(filepath.Join(outDir, imageName), finalBytes, 0o640); err != nil {
		w.fail(ctx, x.ID, err)
		return err
	}
	if err := atomicWrite(filepath.Join(outDir, "manifest.json"), mb, 0o640); err != nil {
		w.fail(ctx, x.ID, err)
		return err
	}
	imgURI := w.uri(filepath.ToSlash(filepath.Join(relDir, imageName)))
	manURI := w.uri(filepath.ToSlash(filepath.Join(relDir, "manifest.json")))
	_, err = w.db.ExecContext(ctx, `UPDATE tasks SET artifact_status='ready',artifact_uri=?,artifact_manifest_uri=?,artifact_error=NULL WHERE id=?`, imgURI, manURI, x.ID)
	return err
}

func joinPrompt(parts ...string) string {
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return strings.Join(out, "\n")
}

func buildPrompt(base, taskInstructions string, x win) string {
	winner := x.Winner
	if len(winner) > 24 {
		winner = winner[:24]
	}
	creative := joinPrompt(base, taskInstructions)
	return creative + fmt.Sprintf("\nTask fingerprint: %s. Winner fingerprint: %s. Difficulty: %d-bit probability space. Public seed fingerprint: %s. Treat these values only as deterministic creative seeds; do not render them as readable text.", x.ID, winner, x.RangeBits, shortHash(x.Seed))
}

func shortHash(s string) string {
	h := sha256.Sum256([]byte(s))
	return hex.EncodeToString(h[:8])
}

func (w *Worker) generate(ctx context.Context, cfg settings.Runtime, x win, prompt, negativePrompt string) (imageResult, error) {
	provider := strings.ToLower(strings.TrimSpace(cfg.ArtifactProvider))
	switch provider {
	case "local":
		return imageResult{Bytes: procedural(x), Ext: "svg", Provider: "local-procedural", Meta: map[string]any{"model": "deterministic-svg"}}, nil
	case "openai":
		return w.openAI(ctx, cfg, x, prompt)
	case "comfyui":
		return w.comfyUI(ctx, cfg, x, prompt, negativePrompt)
	case "a1111":
		return w.a1111(ctx, cfg, x, prompt, negativePrompt)
	case "auto":
		var errs []string
		if strings.TrimSpace(os.Getenv("OPENAI_API_KEY")) != "" {
			if r, err := w.openAI(ctx, cfg, x, prompt); err == nil {
				return r, nil
			} else {
				errs = append(errs, "openai: "+err.Error())
			}
		}
		if strings.TrimSpace(os.Getenv("COMFYUI_URL")) != "" && strings.TrimSpace(os.Getenv("COMFYUI_WORKFLOW_PATH")) != "" {
			if r, err := w.comfyUI(ctx, cfg, x, prompt, negativePrompt); err == nil {
				return r, nil
			} else {
				errs = append(errs, "comfyui: "+err.Error())
			}
		}
		if strings.TrimSpace(os.Getenv("A1111_URL")) != "" {
			if r, err := w.a1111(ctx, cfg, x, prompt, negativePrompt); err == nil {
				return r, nil
			} else {
				errs = append(errs, "a1111: "+err.Error())
			}
		}
		meta := map[string]any{"model": "deterministic-svg"}
		if len(errs) > 0 {
			meta["fallback_errors"] = errs
		}
		return imageResult{Bytes: procedural(x), Ext: "svg", Provider: "local-procedural-fallback", Meta: meta}, nil
	default:
		return imageResult{}, fmt.Errorf("unknown artifact provider %q", provider)
	}
}

func (w *Worker) comfyUI(ctx context.Context, cfg settings.Runtime, x win, prompt, negativePrompt string) (imageResult, error) {
	base := strings.TrimRight(strings.TrimSpace(os.Getenv("COMFYUI_URL")), "/")
	workflowPath := strings.TrimSpace(os.Getenv("COMFYUI_WORKFLOW_PATH"))
	if base == "" || workflowPath == "" {
		return imageResult{}, errors.New("COMFYUI_URL and COMFYUI_WORKFLOW_PATH are required")
	}
	raw, err := os.ReadFile(workflowPath)
	if err != nil {
		return imageResult{}, fmt.Errorf("read ComfyUI workflow: %w", err)
	}
	var workflow any
	if err := json.Unmarshal(raw, &workflow); err != nil {
		return imageResult{}, fmt.Errorf("parse ComfyUI workflow: %w", err)
	}
	seed := deterministicSeed(x)
	replacements := map[string]string{
		"{{PROMPT}}":          prompt,
		"{{NEGATIVE_PROMPT}}": negativePrompt,
		"{{SEED}}":            strconv.FormatInt(seed, 10),
		"{{WIDTH}}":           strconv.Itoa(cfg.ArtifactWidth),
		"{{HEIGHT}}":          strconv.Itoa(cfg.ArtifactHeight),
		"{{STEPS}}":           strconv.Itoa(cfg.ArtifactSteps),
		"{{MODEL}}":           cfg.ArtifactModel,
	}
	workflow = replaceJSON(workflow, replacements)
	clientID := "neuralhunt-" + shortHash(x.ID)
	payload := map[string]any{"prompt": workflow, "client_id": clientID}
	var queued struct {
		PromptID string         `json:"prompt_id"`
		Error    any            `json:"error"`
		Nodes    map[string]any `json:"node_errors"`
	}
	if err := w.doJSON(ctx, http.MethodPost, base+"/prompt", payload, &queued, "", ""); err != nil {
		return imageResult{}, err
	}
	if queued.PromptID == "" {
		return imageResult{}, fmt.Errorf("ComfyUI rejected workflow: error=%v node_errors=%v", queued.Error, queued.Nodes)
	}
	deadline := time.Now().Add(envDuration("COMFYUI_POLL_TIMEOUT", 4*time.Minute))
	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return imageResult{}, ctx.Err()
		case <-time.After(900 * time.Millisecond):
		}
		resp, err := w.http.Get(base + "/history/" + url.PathEscape(queued.PromptID))
		if err != nil {
			continue
		}
		b, readErr := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
		resp.Body.Close()
		if readErr != nil || resp.StatusCode/100 != 2 {
			continue
		}
		var history map[string]any
		if json.Unmarshal(b, &history) != nil {
			continue
		}
		entry, ok := history[queued.PromptID].(map[string]any)
		if !ok {
			continue
		}
		ref, ok := findComfyImage(entry["outputs"])
		if !ok {
			continue
		}
		q := url.Values{}
		q.Set("filename", ref.Filename)
		q.Set("subfolder", ref.Subfolder)
		q.Set("type", ref.Type)
		imgResp, err := w.http.Get(base + "/view?" + q.Encode())
		if err != nil {
			return imageResult{}, err
		}
		img, err := io.ReadAll(io.LimitReader(imgResp.Body, 64<<20))
		imgResp.Body.Close()
		if err != nil {
			return imageResult{}, err
		}
		if imgResp.StatusCode/100 != 2 {
			return imageResult{}, fmt.Errorf("ComfyUI /view HTTP %d", imgResp.StatusCode)
		}
		ext := strings.TrimPrefix(strings.ToLower(filepath.Ext(ref.Filename)), ".")
		if ext == "" {
			ext = "png"
		}
		return imageResult{Bytes: img, Ext: ext, Provider: "comfyui", Meta: map[string]any{"prompt_id": queued.PromptID, "model": cfg.ArtifactModel, "seed": seed}}, nil
	}
	return imageResult{}, errors.New("ComfyUI generation timed out")
}

type comfyImageRef struct{ Filename, Subfolder, Type string }

func findComfyImage(v any) (comfyImageRef, bool) {
	switch x := v.(type) {
	case map[string]any:
		if fn, ok := x["filename"].(string); ok && fn != "" {
			sub, _ := x["subfolder"].(string)
			typ, _ := x["type"].(string)
			if typ == "" {
				typ = "output"
			}
			return comfyImageRef{fn, sub, typ}, true
		}
		for _, child := range x {
			if r, ok := findComfyImage(child); ok {
				return r, true
			}
		}
	case []any:
		for _, child := range x {
			if r, ok := findComfyImage(child); ok {
				return r, true
			}
		}
	}
	return comfyImageRef{}, false
}

func replaceJSON(v any, repl map[string]string) any {
	switch x := v.(type) {
	case map[string]any:
		for k, child := range x {
			x[k] = replaceJSON(child, repl)
		}
		return x
	case []any:
		for i := range x {
			x[i] = replaceJSON(x[i], repl)
		}
		return x
	case string:
		original := x
		for from, to := range repl {
			x = strings.ReplaceAll(x, from, to)
		}
		// Exact numeric placeholders become JSON numbers where possible.
		if original == "{{SEED}}" || original == "{{WIDTH}}" || original == "{{HEIGHT}}" || original == "{{STEPS}}" {
			if n, err := strconv.ParseInt(x, 10, 64); err == nil {
				return n
			}
		}
		return x
	default:
		return v
	}
}

func (w *Worker) a1111(ctx context.Context, cfg settings.Runtime, x win, prompt, negativePrompt string) (imageResult, error) {
	base := strings.TrimRight(strings.TrimSpace(os.Getenv("A1111_URL")), "/")
	if base == "" {
		return imageResult{}, errors.New("A1111_URL is not configured")
	}
	seed := deterministicSeed(x)
	payload := map[string]any{
		"prompt":          prompt,
		"negative_prompt": negativePrompt,
		"steps":           cfg.ArtifactSteps,
		"width":           cfg.ArtifactWidth,
		"height":          cfg.ArtifactHeight,
		"seed":            seed,
		"cfg_scale":       envFloat("A1111_CFG_SCALE", 7.0),
	}
	if sampler := strings.TrimSpace(os.Getenv("A1111_SAMPLER")); sampler != "" {
		payload["sampler_name"] = sampler
	}
	if model := strings.TrimSpace(cfg.ArtifactModel); model != "" && !strings.HasPrefix(strings.ToLower(model), "gpt-image") {
		payload["override_settings"] = map[string]any{"sd_model_checkpoint": model}
	}
	var out struct {
		Images []string `json:"images"`
		Info   string   `json:"info"`
	}
	if err := w.doJSON(ctx, http.MethodPost, base+"/sdapi/v1/txt2img", payload, &out, os.Getenv("A1111_USER"), os.Getenv("A1111_PASSWORD")); err != nil {
		return imageResult{}, err
	}
	if len(out.Images) == 0 {
		return imageResult{}, errors.New("A1111 returned no images")
	}
	b64 := out.Images[0]
	if i := strings.Index(b64, ","); strings.HasPrefix(b64, "data:") && i >= 0 {
		b64 = b64[i+1:]
	}
	img, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return imageResult{}, err
	}
	return imageResult{Bytes: img, Ext: "png", Provider: "a1111", Meta: map[string]any{"seed": seed, "steps": cfg.ArtifactSteps, "model": cfg.ArtifactModel}}, nil
}

func (w *Worker) doJSON(ctx context.Context, method, endpoint string, payload any, out any, user, pass string) error {
	body, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, method, endpoint, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if strings.TrimSpace(user) != "" {
		req.SetBasicAuth(user, pass)
	}
	resp, err := w.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
	if err != nil {
		return err
	}
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("HTTP %d from %s: %s", resp.StatusCode, endpoint, truncate(string(raw), 1200))
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(raw, out); err != nil {
		return fmt.Errorf("decode %s: %w", endpoint, err)
	}
	return nil
}

func envFloat(k string, d float64) float64 {
	if raw := strings.TrimSpace(os.Getenv(k)); raw != "" {
		if v, err := strconv.ParseFloat(raw, 64); err == nil {
			return v
		}
	}
	return d
}

func deterministicSeed(x win) int64 {
	h := sha256.Sum256([]byte(x.ID + "|" + x.Winner + "|" + x.Seed))
	var n uint64
	for i := 0; i < 8; i++ {
		n = (n << 8) | uint64(h[i])
	}
	return int64(n & 0x7fffffffffffffff)
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}

func (w *Worker) uri(rel string) string {
	p := "/" + strings.TrimLeft(rel, "/")
	if w.publicBase == "" {
		return p
	}
	return w.publicBase + p
}

func (w *Worker) fail(ctx context.Context, taskID string, err error) {
	_, _ = w.db.ExecContext(ctx, `UPDATE tasks SET artifact_status='error',artifact_error=? WHERE id=?`, err.Error(), taskID)
}

func atomicWrite(path string, data []byte, mode os.FileMode) error {
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, mode); err != nil {
		return err
	}
	_ = os.Remove(path)
	return os.Rename(tmp, path)
}

func procedural(x win) []byte {
	h := sha256.Sum256([]byte(x.ID + x.Winner + x.Seed))
	a := int(h[0]) % 360
	b := int(h[1]) % 360
	c := int(h[2]) % 360
	svg := fmt.Sprintf(`<svg xmlns="http://www.w3.org/2000/svg" width="1024" height="1024" viewBox="0 0 1024 1024"><defs><radialGradient id="g"><stop stop-color="hsl(%d 100%% 70%%)"/><stop offset=".55" stop-color="hsl(%d 85%% 45%%)"/><stop offset="1" stop-color="#04111f"/></radialGradient><filter id="gl"><feGaussianBlur stdDeviation="8"/></filter></defs><rect width="1024" height="1024" fill="#04111f"/><circle cx="512" cy="512" r="390" fill="none" stroke="#6574ff" opacity=".18" stroke-width="2"/><circle cx="512" cy="512" r="290" fill="url(#g)" opacity=".22" filter="url(#gl)"/><g fill="hsl(%d 100%% 75%%)">`, a, b, c)
	for i := 0; i < 480; i++ {
		v := int(h[i%32])
		xv := 512 + ((i*73 + v*11) % 620) - 310
		yv := 512 + ((i*97 + v*7) % 620) - 310
		r := 1 + (v % 5)
		svg += fmt.Sprintf(`<circle cx="%d" cy="%d" r="%d" opacity=".65"/>`, xv, yv, r)
	}
	svg += fmt.Sprintf(`</g><text x="52" y="900" fill="#d9f7ff" font-family="monospace" font-size="28">%s</text><text x="52" y="940" fill="#7fe7ff" font-family="monospace" font-size="18">winner %s</text></svg>`, x.ID, x.Winner[:min(20, len(x.Winner))])
	return []byte(svg)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}