package fetch

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
)

// ModelFile describes one pinned OCR model asset on Hugging Face.
type ModelFile struct {
	Name   string
	URL    string
	SHA256 string
	Size   int64
}

// Pinned model files for the Linux OCR pipeline:
// - Meiki text detector (D-FINE): detects horizontal and vertical text lines
// - Meiki text recognizer: recognized horizontal dialogue / subtitles / UI
// - Manga-OCR (kha-white base ONNX): recognized vertical speech bubbles and whole images
var PinnedModels = []ModelFile{
	{
		Name:   "meiki.text.detect.v0.1.960x544.onnx",
		URL:    "https://huggingface.co/rtr46/meiki.text.detect.v0/resolve/a9cffa4f60cbf72ddb87edf19c6f98a01cd042e6/meiki.text.detect.v0.1.960x544.onnx",
		SHA256: "40b6a016667745cae7d3055929ae3b8b1e7716aac795f5904cd3c2c7c3b8404b",
		Size:   14503825,
	},
	{
		Name:   "meiki.text.rec.v0.960x32.onnx",
		URL:    "https://huggingface.co/rtr46/meiki.txt.recognition.v0/resolve/a28cf5874dc2438ebb1c86336be26bcec51e3375/meiki.text.rec.v0.960x32.onnx",
		SHA256: "3e96bc772fbee9717e536a6353032bb944c3382dd2f6960ef4890decda43b000",
		Size:   18593254,
	},
	{
		Name:   "encoder_model.onnx",
		URL:    "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/encoder_model.onnx",
		SHA256: "15fa8155fe9bc1a7d25d9bb353debaa4def033d0174e907dbd2dd6d995def85f",
		Size:   343454249,
	},
	{
		Name:   "decoder_model.onnx",
		URL:    "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/decoder_model.onnx",
		SHA256: "ef7765261e9d1cdc34d89356986c2bbc2a082897f753a89605ae80fdfa61f5e8",
		Size:   117480262,
	},
	{
		Name:   "vocab.txt",
		URL:    "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/vocab.txt",
		SHA256: "5cb5c5586d98a2f331d9f8828e4586479b0611bfba5d8c3b6dadffc84d6a36a3",
		Size:   30216,
	},
}

// VerifyFile returns true if path exists, matches expectedSize, and has expectedSHA256.
func VerifyFile(path string, expectedSHA string, expectedSize int64) bool {
	info, err := os.Stat(path)
	if err != nil || info.Size() != expectedSize {
		return false
	}
	f, err := os.Open(path)
	if err != nil {
		return false
	}
	defer f.Close()

	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return false
	}
	return hex.EncodeToString(h.Sum(nil)) == expectedSHA
}

// EnsureModels verifies all required model files in dir, downloading any missing
// or corrupted files atomically into dir.
func EnsureModels(ctx context.Context, dir string) error {
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("create models dir %s: %w", dir, err)
	}

	for _, m := range PinnedModels {
		dest := filepath.Join(dir, m.Name)
		if VerifyFile(dest, m.SHA256, m.Size) {
			continue
		}
		if err := downloadFile(ctx, dir, m); err != nil {
			return fmt.Errorf("download %s: %w", m.Name, err)
		}
	}
	return nil
}

func downloadFile(ctx context.Context, dir string, m ModelFile) error {
	tmpFile, err := os.CreateTemp(dir, ".tmp-"+m.Name+"-*")
	if err != nil {
		return fmt.Errorf("create temp file: %w", err)
	}
	tmpPath := tmpFile.Name()
	defer func() {
		tmpFile.Close()
		os.Remove(tmpPath)
	}()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, m.URL, nil)
	if err != nil {
		return fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("User-Agent", "chirami-ocr/1.0")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("http get: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("unexpected status %s from %s", resp.Status, m.URL)
	}

	h := sha256.New()
	tee := io.TeeReader(resp.Body, h)

	n, err := io.Copy(tmpFile, tee)
	if err != nil {
		return fmt.Errorf("copying body: %w", err)
	}
	if n != m.Size {
		return fmt.Errorf("expected %d bytes, got %d", m.Size, n)
	}

	actualSHA := hex.EncodeToString(h.Sum(nil))
	if actualSHA != m.SHA256 {
		return fmt.Errorf("sha256 mismatch: got %s, want %s", actualSHA, m.SHA256)
	}

	if err := tmpFile.Sync(); err != nil {
		return fmt.Errorf("sync temp file: %w", err)
	}
	if err := tmpFile.Close(); err != nil {
		return fmt.Errorf("close temp file: %w", err)
	}

	destPath := filepath.Join(dir, m.Name)
	if err := os.Rename(tmpPath, destPath); err != nil {
		return fmt.Errorf("atomic rename: %w", err)
	}
	return nil
}

// Missing returns any model files from PinnedModels that are not present and valid in dir.
func Missing(dir string) []ModelFile {
	var missing []ModelFile
	for _, m := range PinnedModels {
		dest := filepath.Join(dir, m.Name)
		if !VerifyFile(dest, m.SHA256, m.Size) {
			missing = append(missing, m)
		}
	}
	return missing
}

// Absent returns the model files that are missing from dir or have the wrong
// size, without hashing them.
func Absent(dir string) []ModelFile {
	var absent []ModelFile
	for _, m := range PinnedModels {
		if st, err := os.Stat(filepath.Join(dir, m.Name)); err != nil || st.Size() != m.Size {
			absent = append(absent, m)
		}
	}
	return absent
}