//go:build darwin && arm64

package ml

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"time"

	"forge.lthn.ai/core/go-ai/ml"
	"forge.lthn.ai/core/go/pkg/cli"
	"gopkg.in/yaml.v3"
)

var sequenceCmd = &cli.Command{
	Use:   "sequence",
	Short: "Run a training sequence of multiple lessons",
	Long: `Runs an ordered sequence of lessons defined in a YAML file.

Sequence YAML format:
  id: lek-full
  title: "LEK Full Training Sequence"
  mode: vertical
  model-path: /path/to/model
  lessons:
    - sovereignty.yaml
    - privacy.yaml
    - censorship.yaml

Mode:
  vertical   Run lessons strictly in order (default)
  horizontal Run all lessons, order doesn't matter

State is tracked per-sequence so runs can be resumed.`,
	RunE: runSequence,
}

var (
	sequenceFile      string
	sequenceModelPath string
	sequenceOutput    string
	sequenceMaxTokens int
	sequenceTemp      float64
	sequenceMemLimit  int
)

func init() {
	sequenceCmd.Flags().StringVar(&sequenceFile, "file", "", "Sequence YAML file (required)")
	sequenceCmd.Flags().StringVar(&sequenceModelPath, "model-path", "", "Path to model directory (required)")
	sequenceCmd.Flags().StringVar(&sequenceOutput, "output", "", "Output JSONL file (default: <sequence-id>.jsonl)")
	sequenceCmd.Flags().IntVar(&sequenceMaxTokens, "max-tokens", 1024, "Max tokens per response")
	sequenceCmd.Flags().Float64Var(&sequenceTemp, "temperature", 0.4, "Sampling temperature")
	sequenceCmd.Flags().IntVar(&sequenceMemLimit, "memory-limit", 24, "Metal memory limit in GB")
	sequenceCmd.MarkFlagRequired("file")
	sequenceCmd.MarkFlagRequired("model-path")
}

// sequenceDef is a YAML sequence definition.
type sequenceDef struct {
	ID        string   `yaml:"id"`
	Title     string   `yaml:"title"`
	Mode      string   `yaml:"mode"` // "vertical" (default) or "horizontal"
	ModelPath string   `yaml:"model-path"`
	Lessons   []string `yaml:"lessons"` // Relative paths to lesson YAML files
}

// sequenceState tracks progress through a sequence.
type sequenceState struct {
	SequenceID string            `json:"sequence_id"`
	Completed  map[string]bool   `json:"completed"` // lesson ID → done
	Current    string            `json:"current"`
	UpdatedAt  string            `json:"updated_at"`
}

func runSequence(cmd *cli.Command, args []string) error {
	start := time.Now()

	// Load sequence YAML
	data, err := os.ReadFile(sequenceFile)
	if err != nil {
		return fmt.Errorf("read sequence: %w", err)
	}

	var seq sequenceDef
	if err := yaml.Unmarshal(data, &seq); err != nil {
		return fmt.Errorf("parse sequence: %w", err)
	}

	if seq.ID == "" {
		seq.ID = strings.TrimSuffix(filepath.Base(sequenceFile), filepath.Ext(sequenceFile))
	}
	if seq.Mode == "" {
		seq.Mode = "vertical"
	}

	// Model path from sequence or flag
	modelPath := sequenceModelPath
	if modelPath == "" && seq.ModelPath != "" {
		modelPath = seq.ModelPath
	}
	if modelPath == "" {
		return fmt.Errorf("model-path is required (flag or sequence YAML)")
	}

	// Resolve output
	if sequenceOutput == "" {
		sequenceOutput = seq.ID + ".jsonl"
	}

	slog.Info("sequence: loaded",
		"id", seq.ID,
		"title", seq.Title,
		"mode", seq.Mode,
		"lessons", len(seq.Lessons),
	)

	// Load state
	stateFile := seq.ID + ".sequence-state.json"
	state := loadSequenceState(stateFile)
	if state.SequenceID == "" {
		state.SequenceID = seq.ID
		state.Completed = make(map[string]bool)
	}

	// Load model once for all lessons
	slog.Info("sequence: loading model", "path", modelPath)
	backend, err := ml.NewMLXBackend(modelPath)
	if err != nil {
		return fmt.Errorf("load model: %w", err)
	}

	opts := ml.GenOpts{
		Temperature: sequenceTemp,
		MaxTokens:   sequenceMaxTokens,
	}

	// Open output file
	outFile, err := os.OpenFile(sequenceOutput, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return fmt.Errorf("create output: %w", err)
	}
	defer outFile.Close()
	encoder := json.NewEncoder(outFile)

	baseDir := filepath.Dir(sequenceFile)
	totalGenerated := 0

	for i, lessonPath := range seq.Lessons {
		// Resolve lesson path
		if !filepath.IsAbs(lessonPath) {
			lessonPath = filepath.Join(baseDir, lessonPath)
		}

		// Load lesson
		lessonData, err := os.ReadFile(lessonPath)
		if err != nil {
			slog.Error("sequence: failed to read lesson",
				"path", lessonPath,
				"error", err,
			)
			if seq.Mode == "vertical" {
				return fmt.Errorf("vertical sequence halted: %w", err)
			}
			continue
		}

		var lesson lessonDef
		if err := yaml.Unmarshal(lessonData, &lesson); err != nil {
			slog.Error("sequence: failed to parse lesson",
				"path", lessonPath,
				"error", err,
			)
			if seq.Mode == "vertical" {
				return fmt.Errorf("vertical sequence halted: %w", err)
			}
			continue
		}

		if lesson.ID == "" {
			lesson.ID = strings.TrimSuffix(filepath.Base(lessonPath), filepath.Ext(lessonPath))
		}

		// Skip completed lessons
		if state.Completed[lesson.ID] {
			slog.Info("sequence: skipping completed lesson",
				"lesson", fmt.Sprintf("%d/%d", i+1, len(seq.Lessons)),
				"id", lesson.ID,
			)
			continue
		}

		state.Current = lesson.ID

		slog.Info("sequence: starting lesson",
			"lesson", fmt.Sprintf("%d/%d", i+1, len(seq.Lessons)),
			"id", lesson.ID,
			"title", lesson.Title,
			"prompts", len(lesson.Prompts),
		)

		// Load sandwich files for this lesson
		var kbText, kernelText string
		hasSandwich := false
		if lesson.Sandwich != nil {
			lessonDir := filepath.Dir(lessonPath)
			if lesson.Sandwich.KB != "" {
				kbPath := lesson.Sandwich.KB
				if !filepath.IsAbs(kbPath) {
					kbPath = filepath.Join(lessonDir, kbPath)
				}
				d, err := os.ReadFile(kbPath)
				if err != nil {
					slog.Error("sequence: failed to read KB", "error", err)
				} else {
					kbText = string(d)
				}
			}
			if lesson.Sandwich.Kernel != "" {
				kernelPath := lesson.Sandwich.Kernel
				if !filepath.IsAbs(kernelPath) {
					kernelPath = filepath.Join(lessonDir, kernelPath)
				}
				d, err := os.ReadFile(kernelPath)
				if err != nil {
					slog.Error("sequence: failed to read kernel", "error", err)
				} else {
					kernelText = string(d)
				}
			}
			hasSandwich = kbText != "" && kernelText != ""
		}

		// Run each prompt in the lesson
		generated := 0
		for j, prompt := range lesson.Prompts {
			var messages []ml.Message
			if lesson.System != "" {
				messages = append(messages, ml.Message{Role: "system", Content: lesson.System})
			}

			userContent := prompt.Prompt
			if hasSandwich {
				userContent = buildSandwich(kbText, prompt.Prompt, kernelText)
			}
			messages = append(messages, ml.Message{Role: "user", Content: userContent})

			slog.Info("sequence: generating",
				"lesson", lesson.ID,
				"prompt", fmt.Sprintf("%d/%d", j+1, len(lesson.Prompts)),
				"id", prompt.ID,
			)

			response, err := backend.Chat(cmd.Context(), messages, opts)
			if err != nil {
				slog.Error("sequence: generation failed",
					"lesson", lesson.ID,
					"prompt", prompt.ID,
					"error", err,
				)
				continue
			}

			record := struct {
				Messages []ml.Message `json:"messages"`
			}{
				Messages: []ml.Message{
					{Role: "user", Content: userContent},
					{Role: "assistant", Content: response},
				},
			}
			if err := encoder.Encode(record); err != nil {
				return fmt.Errorf("write record: %w", err)
			}

			generated++
			totalGenerated++
		}

		// Mark lesson complete
		state.Completed[lesson.ID] = true
		state.UpdatedAt = time.Now().Format(time.RFC3339)
		saveSequenceState(stateFile, state)

		slog.Info("sequence: lesson complete",
			"id", lesson.ID,
			"generated", generated,
			"total", len(lesson.Prompts),
		)
	}

	state.Current = ""
	state.UpdatedAt = time.Now().Format(time.RFC3339)
	saveSequenceState(stateFile, state)

	slog.Info("sequence: complete",
		"id", seq.ID,
		"output", sequenceOutput,
		"total_generated", totalGenerated,
		"lessons_completed", len(state.Completed),
		"duration", time.Since(start).Round(time.Second),
	)

	return nil
}

func loadSequenceState(path string) sequenceState {
	data, err := os.ReadFile(path)
	if err != nil {
		return sequenceState{}
	}
	var state sequenceState
	json.Unmarshal(data, &state)
	return state
}

func saveSequenceState(path string, state sequenceState) {
	data, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return
	}
	os.WriteFile(path, data, 0644)
}