package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
)

const persistedJobVersion = 1

// persistedJob is deliberately independent of the public job JSON. Runtime
// state, caches, and locks never reach disk; only the data needed to restore a
// completed generation is retained.
type persistedJob struct {
	Version      int           `json:"version"`
	ID           string        `json:"id"`
	CreatedAt    int64         `json:"createdAt"`
	StartedAt    int64         `json:"startedAt,omitempty"`
	FinishedAt   int64         `json:"finishedAt,omitempty"`
	DurationMS   int64         `json:"durationMs,omitempty"`
	Quality      string        `json:"quality,omitempty"`
	Thumbnail    string        `json:"thumbnail,omitempty"`
	Pipeline     int           `json:"pipeline"`
	Seed         uint64        `json:"seed"`
	Steps        int           `json:"steps"`
	TextureSteps int           `json:"textureSteps"`
	Guidance     float32       `json:"guidance"`
	Frames       []frameMeta   `json:"frames,omitempty"`
	LivePreview  bool          `json:"livePreview,omitempty"`
	StageTimings []stageTiming `json:"stageTimings,omitempty"`
}

func persistedFrameName(i int) string {
	return filepath.Join("frames", fmt.Sprintf("%06d.bin", i))
}

// persistJob writes a complete job into a temporary sibling directory and then
// renames that directory into place. Startup therefore sees either the previous
// complete job or no job, never a half-written mesh/frame set.
func (s *server) persistJob(j *job) error {
	if s.storeDir == "" {
		return nil
	}

	j.mu.Lock()
	if j.mesh == nil {
		j.mu.Unlock()
		return fmt.Errorf("job %s has no mesh", j.ID)
	}
	manifest := persistedJob{
		Version: persistedJobVersion, ID: j.ID, CreatedAt: j.CreatedAt,
		StartedAt: j.StartedAt, FinishedAt: j.FinishedAt, DurationMS: j.DurationMS,
		Quality: j.Quality, Thumbnail: j.Thumbnail,
		Pipeline: j.pipeline, Seed: j.seed, Steps: j.steps,
		TextureSteps: j.textureSteps, Guidance: j.guidance,
		Frames: append([]frameMeta(nil), j.Frames...), LivePreview: j.LivePreview,
		StageTimings: append([]stageTiming(nil), j.StageTimings...),
	}
	mesh := j.mesh
	input := j.image
	source := j.source
	if len(source) == 0 {
		source = input
	}
	previews := append([][]byte(nil), j.previews...)
	j.mu.Unlock()

	if len(previews) != len(manifest.Frames) {
		return fmt.Errorf("job %s has %d previews but %d frame records",
			j.ID, len(previews), len(manifest.Frames))
	}
	if err := os.MkdirAll(s.storeDir, 0o755); err != nil {
		return fmt.Errorf("create job store: %w", err)
	}
	tmp, err := os.MkdirTemp(s.storeDir, "."+j.ID+".tmp-")
	if err != nil {
		return fmt.Errorf("create temporary job directory: %w", err)
	}
	defer os.RemoveAll(tmp)

	if err := writeMeshFile(filepath.Join(tmp, "mesh.t2mesh"), mesh); err != nil {
		return err
	}
	if len(source) > 0 {
		if err := os.WriteFile(filepath.Join(tmp, "source.img"), source, 0o644); err != nil {
			return fmt.Errorf("write source image: %w", err)
		}
	}
	if len(input) > 0 {
		if err := os.WriteFile(filepath.Join(tmp, "input.img"), input, 0o644); err != nil {
			return fmt.Errorf("write generation input: %w", err)
		}
	}
	if len(previews) > 0 {
		if err := os.Mkdir(filepath.Join(tmp, "frames"), 0o755); err != nil {
			return fmt.Errorf("create frame directory: %w", err)
		}
		for i, blob := range previews {
			if err := os.WriteFile(filepath.Join(tmp, persistedFrameName(i)), blob, 0o644); err != nil {
				return fmt.Errorf("write preview %d: %w", i, err)
			}
		}
	}
	data, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return fmt.Errorf("encode manifest: %w", err)
	}
	data = append(data, '\n')
	if err := os.WriteFile(filepath.Join(tmp, "manifest.json"), data, 0o644); err != nil {
		return fmt.Errorf("write manifest: %w", err)
	}

	finalDir := filepath.Join(s.storeDir, j.ID)
	if err := os.Rename(tmp, finalDir); err != nil {
		return fmt.Errorf("commit job %s: %w", j.ID, err)
	}

	// Frames are immutable and now durable, so release their duplicate in-memory
	// copies. The preview endpoint transparently reads them from finalDir.
	j.mu.Lock()
	j.persistDir = finalDir
	j.meshPath = filepath.Join(finalDir, "mesh.t2mesh")
	if len(input) > 0 {
		j.inputPath = filepath.Join(finalDir, "input.img")
	}
	if len(source) > 0 {
		j.sourcePath = filepath.Join(finalDir, "source.img")
	}
	j.previews = nil
	j.mu.Unlock()
	return nil
}

// restoreJobs scans only complete, versioned job directories. Temporary
// directories left by a killed write and malformed/corrupt entries are ignored
// with a log message; one bad asset must not prevent the server from starting.
func (s *server) restoreJobs() (int, error) {
	if s.storeDir == "" {
		return 0, nil
	}
	if err := os.MkdirAll(s.storeDir, 0o755); err != nil {
		return 0, fmt.Errorf("create job store: %w", err)
	}
	entries, err := os.ReadDir(s.storeDir)
	if err != nil {
		return 0, fmt.Errorf("read job store: %w", err)
	}
	restored := 0
	for _, entry := range entries {
		if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") {
			continue
		}
		dir := filepath.Join(s.storeDir, entry.Name())
		j, err := loadPersistedJob(dir)
		if err != nil {
			log.Printf("ignoring persisted job %s: %v", entry.Name(), err)
			continue
		}
		s.mu.Lock()
		if _, exists := s.jobs[j.ID]; !exists {
			s.jobs[j.ID] = j
			restored++
		}
		s.mu.Unlock()
	}
	return restored, nil
}

func loadPersistedJob(dir string) (*job, error) {
	data, err := os.ReadFile(filepath.Join(dir, "manifest.json"))
	if err != nil {
		return nil, fmt.Errorf("read manifest: %w", err)
	}
	var m persistedJob
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, fmt.Errorf("decode manifest: %w", err)
	}
	if m.Version != persistedJobVersion {
		return nil, fmt.Errorf("unsupported manifest version %d", m.Version)
	}
	if m.ID == "" || filepath.Base(m.ID) != m.ID || m.ID != filepath.Base(dir) {
		return nil, fmt.Errorf("invalid job id %q", m.ID)
	}
	meshPath := filepath.Join(dir, "mesh.t2mesh")
	if st, err := os.Stat(meshPath); err != nil || !st.Mode().IsRegular() {
		if err == nil {
			err = fmt.Errorf("not a regular file")
		}
		return nil, fmt.Errorf("mesh: %w", err)
	}
	sourcePath := filepath.Join(dir, "source.img")
	if st, err := os.Stat(sourcePath); err != nil || !st.Mode().IsRegular() {
		sourcePath = "" // optional for generations saved before source retention
	}
	inputPath := filepath.Join(dir, "input.img")
	if st, err := os.Stat(inputPath); err != nil || !st.Mode().IsRegular() {
		inputPath = "" // legacy jobs stored their processed input as source.img
	}
	for i := range m.Frames {
		if st, err := os.Stat(filepath.Join(dir, persistedFrameName(i))); err != nil || !st.Mode().IsRegular() {
			if err == nil {
				err = fmt.Errorf("not a regular file")
			}
			return nil, fmt.Errorf("preview %d: %w", i, err)
		}
	}
	return &job{
		ID: m.ID, State: "done", CreatedAt: m.CreatedAt, StartedAt: m.StartedAt,
		FinishedAt: m.FinishedAt, DurationMS: m.DurationMS,
		Quality: m.Quality, Thumbnail: m.Thumbnail,
		PreviewSeq: len(m.Frames), Frames: m.Frames,
		LivePreview: m.LivePreview || len(m.Frames) > 0, StageTimings: m.StageTimings,
		pipeline: m.Pipeline, seed: m.Seed, steps: m.Steps,
		textureSteps: m.TextureSteps, guidance: m.Guidance,
		persistDir: dir, meshPath: meshPath, inputPath: inputPath, sourcePath: sourcePath,
	}, nil
}

func (s *server) loadJobMesh(j *job) (*meshData, error) {
	j.mu.Lock()
	defer j.mu.Unlock()
	if j.mesh != nil {
		return j.mesh, nil
	}
	if j.meshPath == "" {
		return nil, fmt.Errorf("mesh not ready")
	}
	mesh, err := readMeshFile(j.meshPath)
	if err != nil {
		return nil, err
	}
	j.mesh = mesh
	return mesh, nil
}

func loadJobPreview(j *job, seq int) ([]byte, error) {
	j.mu.Lock()
	if seq < 0 || seq >= j.PreviewSeq {
		j.mu.Unlock()
		return nil, fmt.Errorf("no such preview frame")
	}
	if seq < len(j.previews) && j.previews[seq] != nil {
		blob := j.previews[seq]
		j.mu.Unlock()
		return blob, nil
	}
	dir := j.persistDir
	j.mu.Unlock()
	if dir == "" {
		return nil, fmt.Errorf("no such preview frame")
	}
	return os.ReadFile(filepath.Join(dir, persistedFrameName(seq)))
}

func writeMeshFile(path string, mesh *meshData) error {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o644)
	if err != nil {
		return fmt.Errorf("create mesh: %w", err)
	}
	if err := writeMeshBinary(f, mesh); err != nil {
		f.Close()
		return fmt.Errorf("write mesh: %w", err)
	}
	if err := f.Sync(); err != nil {
		f.Close()
		return fmt.Errorf("sync mesh: %w", err)
	}
	if err := f.Close(); err != nil {
		return fmt.Errorf("close mesh: %w", err)
	}
	return nil
}

func readMeshFile(path string) (*meshData, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open mesh: %w", err)
	}
	defer f.Close()
	var magic [8]byte
	if _, err := io.ReadFull(f, magic[:]); err != nil {
		return nil, fmt.Errorf("read mesh magic: %w", err)
	}
	pbrWidth := 0
	switch string(magic[:]) {
	case "T2MESH01":
	case "T2MESH02":
		pbrWidth = 5
	case "T2MESH03":
		pbrWidth = 6
	default:
		return nil, fmt.Errorf("bad mesh magic %q", magic)
	}
	var nv32, nt32 uint32
	if err := binary.Read(f, binary.LittleEndian, &nv32); err != nil {
		return nil, fmt.Errorf("read vertex count: %w", err)
	}
	if err := binary.Read(f, binary.LittleEndian, &nt32); err != nil {
		return nil, fmt.Errorf("read triangle count: %w", err)
	}
	if nv32 == 0 || nt32 == 0 || nv32 > 100_000_000 || nt32 > 100_000_000 {
		return nil, fmt.Errorf("invalid mesh size %d vertices, %d triangles", nv32, nt32)
	}
	nv, nt := int(nv32), int(nt32)
	st, err := f.Stat()
	if err != nil {
		return nil, fmt.Errorf("stat mesh: %w", err)
	}
	expected := int64(16) + int64(nv)*24 + int64(nv*pbrWidth)*4 + int64(nt)*12
	if st.Size() != expected {
		return nil, fmt.Errorf("mesh size is %d bytes, expected %d", st.Size(), expected)
	}
	m := &meshData{NVerts: nv, NTris: nt, Verts: make([]float32, 3*nv), Normals: make([]float32, 3*nv)}
	if err := binary.Read(f, binary.LittleEndian, m.Verts); err != nil {
		return nil, fmt.Errorf("read vertices: %w", err)
	}
	if err := binary.Read(f, binary.LittleEndian, m.Normals); err != nil {
		return nil, fmt.Errorf("read normals: %w", err)
	}
	if pbrWidth != 0 {
		stored := make([]float32, pbrWidth*nv)
		if err := binary.Read(f, binary.LittleEndian, stored); err != nil {
			return nil, fmt.Errorf("read PBR attributes: %w", err)
		}
		if pbrWidth == 6 {
			m.PBR = stored
		} else {
			m.PBR = make([]float32, 6*nv)
			for i := 0; i < nv; i++ {
				copy(m.PBR[6*i:6*i+5], stored[5*i:5*i+5])
				m.PBR[6*i+5] = 1
			}
		}
	}
	m.Tris = make([]int32, 3*nt)
	if err := binary.Read(f, binary.LittleEndian, m.Tris); err != nil {
		return nil, fmt.Errorf("read triangles: %w", err)
	}
	var extra [1]byte
	if n, err := f.Read(extra[:]); err != io.EOF || n != 0 {
		return nil, fmt.Errorf("mesh contains trailing data")
	}
	return m, nil
}