package comicanimator

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"html/template"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"git.michelsen.id/phill/preface-tools/internal/auth"
	"git.michelsen.id/phill/preface-tools/internal/tools"
	"git.michelsen.id/phill/preface-tools/internal/tools/comicanimator/openrouter"
)

type ApprovalVerifier interface{ VerifyInstructorPIN(string) bool }
type Tool struct {
	cfg                 Config
	log                 *slog.Logger
	approval            ApprovalVerifier
	store               *store
	queue               chan string
	client              *openrouter.Client
	approvalLimiter     *auth.Limiter
	student, instructor http.Handler
}

func New(cfg Config, log *slog.Logger, approval ApprovalVerifier) (*Tool, error) {
	if err := os.MkdirAll(cfg.UploadDirectory, 0750); err != nil {
		return nil, err
	}
	if err := os.MkdirAll(cfg.OutputDirectory, 0750); err != nil {
		return nil, err
	}
	t := &Tool{cfg: cfg, log: log, approval: approval, store: newStore(), queue: make(chan string, cfg.QueueCapacity)}
	if _, err := t.systemPrompt(); err != nil {
		return nil, fmt.Errorf("load Comic Animator system prompt: %w", err)
	}
	t.approvalLimiter = auth.NewLimiter(5, 5*time.Minute)
	t.client = &openrouter.Client{BaseURL: cfg.OpenRouterBaseURL, APIKey: cfg.OpenRouterAPIKey, SiteURL: cfg.OpenRouterSiteURL, AppName: cfg.OpenRouterAppName, HTTP: &http.Client{Timeout: cfg.HTTPTimeout}}
	sm := http.NewServeMux()
	sm.HandleFunc("GET /", t.page)
	sm.HandleFunc("POST /uploads", t.upload)
	sm.HandleFunc("GET /uploads/{id}/preview", t.preview)
	sm.HandleFunc("POST /prompt", t.prompt)
	sm.HandleFunc("POST /generations", t.submit)
	sm.HandleFunc("GET /generations", t.generations)
	sm.HandleFunc("GET /generations/{id}/status", t.status)
	sm.HandleFunc("GET /generations/{id}/video", t.video)
	sm.HandleFunc("GET /generations/{id}/download", t.download)
	sm.HandleFunc("GET /provider-media/{id}", t.providerMedia)
	t.student = sm
	im := http.NewServeMux()
	im.HandleFunc("GET /", t.instructorPage)
	im.HandleFunc("GET /generations", t.instructorGenerations)
	im.HandleFunc("GET /generations/{id}/video", t.instructorVideo)
	im.HandleFunc("GET /generations/{id}/download", t.instructorDownload)
	im.HandleFunc("GET /outputs/{filename}/download", t.outputDownload)
	t.instructor = im
	return t, nil
}
func (t *Tool) Definition() tools.Definition {
	return tools.Definition{Key: "comic-animator", Name: "Comic Animator", Description: "Bring a comic page gently to life."}
}
func (t *Tool) StudentHandler() http.Handler    { return t.student }
func (t *Tool) InstructorHandler() http.Handler { return t.instructor }
func id(prefix string) string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return prefix + hex.EncodeToString(b)
}
func claims(r *http.Request) auth.Claims { c, _ := auth.ClaimsFrom(r); return c }

const maxSystemPromptBytes = 64 << 10

func (t *Tool) systemPrompt() (string, error) {
	b, err := os.ReadFile(t.cfg.PromptFile)
	if err != nil {
		return "", err
	}
	if len(b) == 0 || len(b) > maxSystemPromptBytes {
		return "", fmt.Errorf("system prompt must contain between 1 and %d bytes", maxSystemPromptBytes)
	}
	prompt := strings.TrimSpace(string(b))
	if prompt == "" {
		return "", fmt.Errorf("system prompt is empty")
	}
	return prompt, nil
}

func (t *Tool) page(w http.ResponseWriter, r *http.Request) {
	render(w, toolPage, map[string]any{"CSRF": claims(r).CSRFToken, "Duration": t.cfg.VideoDuration})
}
func (t *Tool) upload(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, t.cfg.MaxUploadBytes+1<<20)
	if err := r.ParseMultipartForm(t.cfg.MaxUploadBytes); err != nil {
		fragmentError(w, http.StatusRequestEntityTooLarge, "Image is too large.")
		return
	}
	f, h, err := r.FormFile("image")
	if err != nil {
		fragmentError(w, 400, "Choose an image.")
		return
	}
	defer f.Close()
	tmp, err := os.CreateTemp(t.cfg.UploadDirectory, ".upload-*")
	if err != nil {
		fragmentError(w, 500, "Could not store image.")
		return
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)
	n, err := io.Copy(tmp, io.LimitReader(f, t.cfg.MaxUploadBytes+1))
	tmp.Close()
	if err != nil || n == 0 || n > t.cfg.MaxUploadBytes {
		fragmentError(w, 413, "Invalid image size.")
		return
	}
	rf, err := os.Open(tmpName)
	if err != nil {
		fragmentError(w, 500, "Could not inspect image.")
		return
	}
	head := make([]byte, 512)
	hn, _ := rf.Read(head)
	_, seekErr := rf.Seek(0, 0)
	var decodeErr error
	if http.DetectContentType(head[:hn]) == "image/webp" {
		decodeErr = validateWebP(head[:hn])
	} else {
		_, _, decodeErr = image.DecodeConfig(rf)
	}
	rf.Close()
	if seekErr != nil || decodeErr != nil {
		fragmentError(w, 415, "The image is corrupt or unsupported.")
		return
	}
	mt := http.DetectContentType(head[:hn])
	declared := strings.Split(h.Header.Get("Content-Type"), ";")[0]
	if declared != "" && declared != "application/octet-stream" && declared != mt {
		fragmentError(w, http.StatusUnsupportedMediaType, "The image type does not match its content.")
		return
	}
	ext := map[string]string{"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}[mt]
	if ext == "" {
		fragmentError(w, 415, "Use PNG, JPEG, or WebP.")
		return
	}
	uid := id("upl_")
	path := filepath.Join(t.cfg.UploadDirectory, uid+ext)
	if err = os.Rename(tmpName, path); err != nil {
		fragmentError(w, 500, "Could not store image.")
		return
	}
	u := Upload{ID: uid, SessionID: claims(r).SessionID, Path: path, Name: filepath.Base(h.Filename), MIMEType: mt, Size: n, CreatedAt: time.Now().UTC()}
	t.store.putUpload(u)
	fmt.Fprintf(w, `<div id="upload-preview" class="preview"><img src="uploads/%s/preview" alt="Uploaded comic preview"><input type="hidden" name="upload_id" value="%s"><p title="%s">%s</p></div>`, template.HTMLEscapeString(uid), template.HTMLEscapeString(uid), template.HTMLEscapeString(u.Name), template.HTMLEscapeString(u.Name))
}

func validateWebP(head []byte) error {
	if len(head) < 20 || string(head[:4]) != "RIFF" || string(head[8:12]) != "WEBP" {
		return fmt.Errorf("invalid WebP header")
	}
	chunk := string(head[12:16])
	if chunk != "VP8 " && chunk != "VP8L" && chunk != "VP8X" {
		return fmt.Errorf("invalid WebP chunk")
	}
	return nil
}
func (t *Tool) ownedUpload(r *http.Request, id string) (Upload, error) {
	u, ok := t.store.upload(id)
	if !ok {
		return u, os.ErrNotExist
	}
	if u.SessionID != claims(r).SessionID {
		return u, errForbidden
	}
	return u, nil
}
func (t *Tool) preview(w http.ResponseWriter, r *http.Request) {
	u, err := t.ownedUpload(r, r.PathValue("id"))
	if err != nil {
		http.NotFound(w, r)
		return
	}
	serveFile(w, r, u.Path, u.MIMEType, false)
}
func (t *Tool) prompt(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 128<<10)
	u, err := t.ownedUpload(r, r.FormValue("upload_id"))
	desc := strings.TrimSpace(r.FormValue("description"))
	if err != nil || desc == "" {
		fragmentError(w, 400, "Upload an image and describe the movement first.")
		return
	}
	if len(desc) > 12000 {
		fragmentError(w, 400, "Description is too long.")
		return
	}
	b, err := os.ReadFile(u.Path)
	if err != nil {
		fragmentError(w, 500, "Could not read image.")
		return
	}
	dataURL := "data:" + u.MIMEType + ";base64," + base64.StdEncoding.EncodeToString(b)
	systemPrompt, err := t.systemPrompt()
	if err != nil {
		t.log.Error("could not load system prompt", "path", t.cfg.PromptFile, "error", err)
		fragmentError(w, 500, "The prompt configuration could not be loaded.")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), t.cfg.HTTPTimeout)
	defer cancel()
	result, err := t.client.Prompt(ctx, t.cfg.PromptModel, desc, dataURL, systemPrompt)
	if err != nil {
		t.log.Error("prompt generation failed", "error", err)
		fragmentError(w, 502, "Prompt generation is temporarily unavailable.")
		return
	}
	if len(result) > 12000 {
		fragmentError(w, 502, "Generated prompt was too long.")
		return
	}
	fmt.Fprintf(w, `<label for="reviewed_prompt">Video prompt</label><textarea class="textarea textarea-bordered" id="reviewed_prompt" name="reviewed_prompt" form="comic-form" rows="10" maxlength="12000" required>%s</textarea><small>%d characters</small>`, template.HTMLEscapeString(result), len([]rune(result)))
}
func (t *Tool) submit(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 128<<10)
	key := claims(r).SessionID
	if !t.approvalLimiter.Allow(key) || !t.approval.VerifyInstructorPIN(r.FormValue("instructor_pin")) {
		t.approvalLimiter.Fail(key)
		w.Header().Set("HX-Retarget", "#approval-error")
		w.Header().Set("HX-Reswap", "innerHTML")
		fragmentError(w, 403, "Approval was not accepted.")
		return
	}
	t.approvalLimiter.Success(key)
	u, err := t.ownedUpload(r, r.FormValue("upload_id"))
	prompt := strings.TrimSpace(r.FormValue("reviewed_prompt"))
	duration, errDuration := strconv.Atoi(r.FormValue("duration"))
	if err != nil || prompt == "" || len(prompt) > 12000 || errDuration != nil || duration != t.cfg.VideoDuration {
		fragmentError(w, 400, "Check the image, prompt, and settings.")
		return
	}
	now := time.Now().UTC()
	g := Generation{ID: id("gen_"), SessionID: claims(r).SessionID, UploadID: u.ID, OriginalName: u.Name, ReviewedPrompt: prompt, Status: Queued, Duration: duration, Resolution: t.cfg.VideoResolution, CreatedAt: now, UpdatedAt: now}
	t.store.putGeneration(g)
	select {
	case t.queue <- g.ID:
		w.Header().Set("HX-Trigger", `{"generationQueued":{"id":"`+g.ID+`"}}`)
		t.renderCard(w, g, false)
	default:
		t.store.update(g.ID, func(x *Generation) {
			x.Status = Failed
			x.ErrorCode = "internal_error"
			x.ErrorMessage = "Generation queue is full."
		})
		fragmentError(w, 503, "Generation queue is full.")
	}
}
func (t *Tool) generations(w http.ResponseWriter, r *http.Request) {
	for _, g := range t.store.list(claims(r).SessionID) {
		t.renderCard(w, g, false)
	}
}
func (t *Tool) status(w http.ResponseWriter, r *http.Request) {
	g, ok := t.store.generation(r.PathValue("id"))
	if !ok || g.SessionID != claims(r).SessionID {
		http.NotFound(w, r)
		return
	}
	t.renderCard(w, g, false)
}
func (t *Tool) video(w http.ResponseWriter, r *http.Request)    { t.generationFile(w, r, false, false) }
func (t *Tool) download(w http.ResponseWriter, r *http.Request) { t.generationFile(w, r, false, true) }
func (t *Tool) generationFile(w http.ResponseWriter, r *http.Request, instructor, download bool) {
	g, ok := t.store.generation(r.PathValue("id"))
	if !ok || g.Status != Completed || (!instructor && g.SessionID != claims(r).SessionID) {
		http.NotFound(w, r)
		return
	}
	serveFile(w, r, g.OutputPath, g.OutputMIMEType, download)
}
func serveFile(w http.ResponseWriter, r *http.Request, path, mt string, download bool) {
	w.Header().Set("Content-Type", mt)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	if download {
		w.Header().Set("Content-Disposition", `attachment; filename="`+strings.ReplaceAll(filepath.Base(path), `"`, "")+`"`)
	}
	http.ServeFile(w, r, path)
}
func (t *Tool) signedURL(uploadID string) string {
	exp := time.Now().Add(t.cfg.SignedURLTTL).Unix()
	raw := uploadID + "\n" + strconv.FormatInt(exp, 10)
	m := hmac.New(sha256.New, []byte(t.cfg.SigningSecret))
	m.Write([]byte(raw))
	return strings.TrimRight(t.cfg.PublicBaseURL, "/") + "/tools/comic-animator/provider-media/" + url.PathEscape(uploadID) + "?expires=" + strconv.FormatInt(exp, 10) + "&signature=" + hex.EncodeToString(m.Sum(nil))
}
func (t *Tool) providerMedia(w http.ResponseWriter, r *http.Request) {
	exp, err := strconv.ParseInt(r.URL.Query().Get("expires"), 10, 64)
	if err != nil || time.Now().Unix() > exp {
		http.Error(w, "expired", 403)
		return
	}
	raw := r.PathValue("id") + "\n" + strconv.FormatInt(exp, 10)
	m := hmac.New(sha256.New, []byte(t.cfg.SigningSecret))
	m.Write([]byte(raw))
	got, err := hex.DecodeString(r.URL.Query().Get("signature"))
	if err != nil || !hmac.Equal(got, m.Sum(nil)) {
		http.Error(w, "forbidden", 403)
		return
	}
	u, ok := t.store.upload(r.PathValue("id"))
	if !ok {
		http.NotFound(w, r)
		return
	}
	serveFile(w, r, u.Path, u.MIMEType, false)
}