package modelstore

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"rk3588sys/agent/internal/files"
)

type Item struct {
	Name    string `json:"name"`
	Sha256  string `json:"sha256"`
	Path    string `json:"path"`
	Size    int64  `json:"size"`
	MtimeMS int64  `json:"mtime_ms"`
}

type Manifest struct {
	Items []Item `json:"items"`
}

type Store struct {
	ModelsDir      string
	MaxUploadBytes int64
}

func New(modelsDir string, maxUploadMB int) *Store {
	max := int64(maxUploadMB) * 1024 * 1024
	return &Store{ModelsDir: modelsDir, MaxUploadBytes: max}
}

func (s *Store) ManifestPath() string {
	return filepath.Join(s.ModelsDir, "manifest.json")
}

func (s *Store) FilesDir() string {
	return filepath.Join(s.ModelsDir, "files")
}

func (s *Store) Upload(name string, r io.Reader, contentLength int64, expectedSha256 string) (Item, error) {
	if contentLength <= 0 {
		return Item{}, errors.New("missing Content-Length")
	}
	if s.MaxUploadBytes > 0 && contentLength > s.MaxUploadBytes {
		return Item{}, ErrPayloadTooLarge
	}
	filesDir := s.FilesDir()
	if err := files.EnsureDir(filesDir, 0o755); err != nil {
		return Item{}, err
	}

	tmpFile, err := os.CreateTemp(filesDir, ".tmp-*")
	if err != nil {
		return Item{}, fmt.Errorf("create temp: %w", err)
	}
	tmp := tmpFile.Name()
	ok := false
	defer func() {
		_ = tmpFile.Close()
		if !ok {
			_ = os.Remove(tmp)
		}
	}()

	h := sha256.New()
	mw := io.MultiWriter(tmpFile, h)
	if _, err := io.CopyN(mw, r, contentLength); err != nil {
		return Item{}, fmt.Errorf("read body: %w", err)
	}
	if err := tmpFile.Chmod(0o644); err != nil {
		return Item{}, fmt.Errorf("chmod temp: %w", err)
	}
	if err := tmpFile.Sync(); err != nil {
		return Item{}, fmt.Errorf("fsync temp: %w", err)
	}
	if err := tmpFile.Close(); err != nil {
		return Item{}, fmt.Errorf("close temp: %w", err)
	}

	sha := hex.EncodeToString(h.Sum(nil))
	if expectedSha256 != "" && !strings.EqualFold(expectedSha256, sha) {
		return Item{}, fmt.Errorf("sha256 mismatch")
	}

	finalName := fmt.Sprintf("%s__%s.rknn", name, sha)
	finalPath := filepath.Join(filesDir, finalName)
	if _, err := os.Stat(finalPath); err == nil {
		_ = os.Remove(tmp)
	} else {
		if err := files.ReplaceFile(tmp, finalPath); err != nil {
			return Item{}, err
		}
	}

	stat, err := os.Stat(finalPath)
	if err != nil {
		return Item{}, fmt.Errorf("stat final: %w", err)
	}

	item := Item{
		Name:    name,
		Sha256:  sha,
		Path:    filepath.ToSlash(finalPath),
		Size:    stat.Size(),
		MtimeMS: stat.ModTime().UnixMilli(),
	}
	if err := s.upsertManifest(item); err != nil {
		return Item{}, err
	}
	ok = true
	return item, nil
}

func (s *Store) List() (Manifest, error) {
	var m Manifest
	b, err := os.ReadFile(s.ManifestPath())
	if err != nil {
		if os.IsNotExist(err) {
			return Manifest{Items: []Item{}}, nil
		}
		return Manifest{}, fmt.Errorf("read manifest: %w", err)
	}
	if err := json.Unmarshal(b, &m); err != nil {
		return Manifest{}, fmt.Errorf("parse manifest: %w", err)
	}
	if m.Items == nil {
		m.Items = []Item{}
	}
	return m, nil
}

func (s *Store) upsertManifest(item Item) error {
	m, err := s.List()
	if err != nil {
		return err
	}
	found := false
	for i := range m.Items {
		if m.Items[i].Name == item.Name {
			m.Items[i] = item
			found = true
			break
		}
	}
	if !found {
		m.Items = append(m.Items, item)
	}
	b, err := json.Marshal(m)
	if err != nil {
		return fmt.Errorf("marshal manifest: %w", err)
	}
	if err := files.WriteFileAtomic(s.ManifestPath(), append(b, '\n'), 0o644); err != nil {
		return fmt.Errorf("write manifest: %w", err)
	}
	return nil
}

var ErrPayloadTooLarge = errors.New("payload too large")