package pipeline

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"strings"
	"time"

	"textmachine/backend/internal/config"
	"textmachine/backend/internal/ledger"
	"textmachine/backend/internal/llm"
	"textmachine/backend/internal/obs"
	"textmachine/backend/internal/store"
)

// runner.go: the Phase-0 mini-runner — one chunk through the configured stage
// list (C1: draft→edit) with full money discipline: reserve → call → settle+
// checkpoint (одна транзакция) → request_log. Resume: найденный чекпоинт =
// вызов не повторяется и не оплачивается. Гейты/эскалация/циклы по главам —
// Фаза 1, сюда же (код раннера).

// Runner executes a book's pipeline.
type Runner struct {
	Book     *config.Book
	Models   *config.Models
	Pipeline *config.Pipeline
	Store    *store.Store
	Pricer   *ledger.Pricer
	Log      *slog.Logger

	clients   map[string]llm.LLMClient
	templates map[string]*PromptTemplate
}

// NewRunner loads all three configs rooted at book.yaml and opens the project
// store. Caller owns Close.
func NewRunner(bookPath string, logger *slog.Logger) (*Runner, error) {
	book, err := config.LoadBook(bookPath)
	if err != nil {
		return nil, err
	}
	models, err := config.LoadModels(book.ModelsFile)
	if err != nil {
		return nil, err
	}
	pipe, err := config.LoadPipeline(book.Pipeline, models)
	if err != nil {
		return nil, err
	}
	pricer, err := models.Prices()
	if err != nil {
		return nil, err
	}
	st, err := store.Open(book.ProjectDB)
	if err != nil {
		return nil, err
	}
	r := &Runner{
		Book:      book,
		Models:    models,
		Pipeline:  pipe,
		Store:     st,
		Pricer:    pricer,
		Log:       logger,
		clients:   map[string]llm.LLMClient{},
		templates: map[string]*PromptTemplate{},
	}
	if err := r.loadTemplates(); err != nil {
		st.Close()
		return nil, err
	}
	return r, nil
}

func (r *Runner) Close() error { return r.Store.Close() }

func (r *Runner) loadTemplates() error {
	for _, st := range r.Pipeline.Stages {
		tpl, err := LoadPromptTemplate(st.Prompt)
		if err != nil {
			return err
		}
		r.templates[st.Name] = tpl
	}
	return nil
}

func (r *Runner) client(model string) (llm.LLMClient, error) {
	if c, ok := r.clients[model]; ok {
		return c, nil
	}
	c, err := BuildClient(r.Models, model, r.Log)
	if err != nil {
		return nil, err
	}
	r.clients[model] = c
	return c, nil
}

// snapshotID materializes the job-context snapshot (§3.2): brief_hash, chunker
// version and the full stage plan (model, prompt version + CONTENT hash,
// sampling). Payload is rendered with a fixed field order — the id is a
// content hash, so identical context re-upserts idempotently.
func (r *Runner) snapshotID() (id, payload string, err error) {
	type stageSnap struct {
		Name          string  `json:"name"`
		Role          string  `json:"role"`
		Model         string  `json:"model"`
		PromptVersion string  `json:"prompt_version"`
		PromptSHA256  string  `json:"prompt_sha256"`
		Temperature   float64 `json:"temperature"`
		Reasoning     string  `json:"reasoning"`
	}
	snap := struct {
		BriefHash      string      `json:"brief_hash"`
		ChunkerVersion string      `json:"chunker_version"`
		PipelineCore   string      `json:"pipeline_core"`
		Stages         []stageSnap `json:"stages"`
	}{
		BriefHash:      r.Book.BriefHash(),
		ChunkerVersion: chunkerVersion,
		PipelineCore:   r.Pipeline.Core,
	}
	for _, st := range r.Pipeline.Stages {
		snap.Stages = append(snap.Stages, stageSnap{
			Name: st.Name, Role: st.Role, Model: st.Model,
			PromptVersion: st.PromptVersion, PromptSHA256: r.templates[st.Name].SHA256,
			Temperature: st.Temperature, Reasoning: st.Reasoning,
		})
	}
	data, err := json.Marshal(snap)
	if err != nil {
		return "", "", err
	}
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:]), string(data), nil
}

// StageResult reports one executed stage.
type StageResult struct {
	Stage        string
	Role         string
	Model        string // фактически ответившая модель
	FromResume   bool   // served from a checkpoint, no call made
	Usage        llm.Usage
	CostUSD      float64
	LatencyMS    int
	FinishReason string
	Text         string
}

// RunResult is the whole chunk's outcome.
type RunResult struct {
	BookID    string
	Chapter   int
	ChunkIdx  int
	Stages    []StageResult
	FinalText string
	TotalUSD  float64
}

// TranslateOneChunk runs the configured stages over the book's source file as
// a single chunk (Фаза 0). Chapter/chunk are fixed at 1/0 until the chunker
// lands in Phase 1.
func (r *Runner) TranslateOneChunk(ctx context.Context) (*RunResult, error) {
	srcRaw, err := os.ReadFile(r.Book.SourceFile)
	if err != nil {
		return nil, fmt.Errorf("pipeline: read source: %w", err)
	}
	source := strings.TrimSpace(string(srcRaw))
	if source == "" {
		return nil, fmt.Errorf("pipeline: source file %s is empty", r.Book.SourceFile)
	}

	snapID, snapPayload, err := r.snapshotID()
	if err != nil {
		return nil, err
	}
	if err := r.Store.UpsertSnapshot(snapID, r.Book.BriefHash(), snapPayload); err != nil {
		return nil, err
	}

	const chapter, chunkIdx = 1, 0
	res := &RunResult{BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx}
	prev := ""

	for _, st := range r.Pipeline.Stages {
		sr, err := r.runStage(ctx, st, snapID, chapter, chunkIdx, source, prev)
		if err != nil {
			return res, err
		}
		res.Stages = append(res.Stages, *sr)
		res.TotalUSD += sr.CostUSD
		prev = sr.Text
	}
	res.FinalText = prev
	return res, nil
}

func (r *Runner) runStage(ctx context.Context, st config.Stage, snapID string, chapter, chunkIdx int, source, prev string) (*StageResult, error) {
	job, err := r.Store.EnsureJob(r.Book.BookID, chapter, st.Name, snapID)
	if err != nil {
		return nil, err
	}
	if err := r.Store.SetJobStatus(job.ID, "running"); err != nil {
		return nil, err
	}

	msgs, err := Messages(r.templates[st.Name], RenderVars{Book: r.Book, Text: source, Draft: prev})
	if err != nil {
		return nil, err
	}

	// max_tokens: калибровка полигона — русский выход ≈1.9× входа в токенах;
	// ratio из конфига, floor защищает короткие чанки. Входит в request-hash
	// (меняет вызов), поэтому считается детерминированно от исходника.
	inputEst := EstimateTokens(source)
	maxTokens := int(float64(inputEst) * r.Pipeline.Defaults.MaxOutputRatio)
	if maxTokens < r.Pipeline.Defaults.MinMaxTokens {
		maxTokens = r.Pipeline.Defaults.MinMaxTokens
	}

	reqHash := RequestHash(r.Book.BookID, chapter, chunkIdx, st.Name, st.Role, st.Model,
		st.Temperature, st.Reasoning, false, maxTokens, snapID, msgs)

	ctx = obs.WithReqInfo(ctx, obs.ReqInfo{
		TraceID: obs.TraceFromContext(ctx), Book: r.Book.BookID,
		Chapter: chapter, Chunk: chunkIdx, Stage: st.Name, Role: st.Role,
	})

	// Resume path: a checkpoint means the call already happened and was paid —
	// serve it, log the hit, never re-bill (приёмка: kill -9 теряет ≤1 вызов).
	if cp, err := r.Store.GetCheckpoint(reqHash); err != nil {
		return nil, err
	} else if cp != nil {
		var usage llm.Usage
		_ = json.Unmarshal([]byte(cp.UsageJSON), &usage)
		r.Store.LogRequest(ctx, r.Log, store.RequestLog{
			BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
			Stage: st.Name, Role: st.Role, ModelRequested: st.Model, ModelActual: cp.ModelActual,
			RequestHash: reqHash, TMHit: true, OK: true, FinishReason: cp.FinishReason,
		})
		if err := r.Store.SetJobStatus(job.ID, "done"); err != nil {
			return nil, err
		}
		r.Log.InfoContext(ctx, "stage served from checkpoint", "stage", st.Name, "hash", reqHash[:12])
		return &StageResult{Stage: st.Name, Role: st.Role, Model: cp.ModelActual,
			FromResume: true, Usage: usage, CostUSD: 0, FinishReason: cp.FinishReason, Text: cp.ResponseText}, nil
	}

	price := r.Pricer.PriceFor(st.Model)
	promptEst := 0
	for _, m := range msgs {
		promptEst += EstimateTokens(m.Content)
	}
	estimate := ledger.EstimateUSD(price, promptEst, maxTokens)

	resv, verdict, err := r.Store.Reserve(r.Book.BookID, estimate, store.Ceilings{
		BookUSD: r.Book.Ceilings.BookUSD, DayUSD: r.Book.Ceilings.DayUSD,
	})
	if err != nil {
		return nil, err
	}
	switch verdict {
	case store.ReserveDeniedBook:
		return nil, fmt.Errorf("pipeline: book USD ceiling reached (%.2f$) — raise ceilings.book_usd or stop", r.Book.Ceilings.BookUSD)
	case store.ReserveDeniedDay:
		return nil, fmt.Errorf("pipeline: daily USD ceiling reached (%.2f$)", r.Book.Ceilings.DayUSD)
	}

	client, err := r.client(st.Model)
	if err != nil {
		_ = r.Store.Release(resv)
		return nil, err
	}

	start := time.Now()
	resp, err := client.Complete(ctx, llm.LLMRequest{
		Model:           st.Model,
		Messages:        msgs,
		MaxTokens:       maxTokens,
		Temperature:     st.Temperature,
		ReasoningEffort: st.Reasoning,
	})
	latency := int(time.Since(start).Milliseconds())
	if err != nil {
		// No 2xx ever arrived: nothing was billed by the provider — release the
		// reservation and surface the failure with a telemetry row.
		if rerr := r.Store.Release(resv); rerr != nil {
			r.Log.ErrorContext(ctx, "release after failed call also failed", "err", rerr)
		}
		r.Store.LogRequest(ctx, r.Log, store.RequestLog{
			BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
			Stage: st.Name, Role: st.Role, ModelRequested: st.Model,
			RequestHash: reqHash, LatencyMS: latency, Err: err.Error(), OK: false,
		})
		_ = r.Store.SetJobStatus(job.ID, "failed")
		return nil, fmt.Errorf("pipeline: stage %s call: %w", st.Name, err)
	}

	modelActual := resp.Model
	if modelActual == "" {
		modelActual = st.Model
	}
	cost := ledger.CostUSD(r.Pricer.PriceFor(modelActual), resp.Usage)
	usageJSON, err := json.Marshal(resp.Usage)
	if err != nil {
		return nil, err
	}

	// Деньги: settle + сырой ответ — одна транзакция (§3.3).
	if err := r.Store.SettleWithCheckpoint(resv, cost, store.Checkpoint{
		RequestHash: reqHash, JobID: job.ID, ChunkIdx: chunkIdx, Stage: st.Name, Role: st.Role,
		ModelRequested: st.Model, ModelActual: modelActual, ResponseText: resp.Text,
		UsageJSON: string(usageJSON), CostUSD: cost, FinishReason: resp.FinishReason,
		ProviderRequestID: resp.ProviderRequestID,
	}); err != nil {
		// The provider HAS billed this 2xx; failing to persist means the money
		// state is behind reality — fail loud, never continue on top.
		return nil, fmt.Errorf("pipeline: settle stage %s: %w", st.Name, err)
	}

	r.Store.LogRequest(ctx, r.Log, store.RequestLog{
		BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
		Stage: st.Name, Role: st.Role, ModelRequested: st.Model, ModelActual: modelActual,
		RequestHash: reqHash,
		PromptTokens: resp.Usage.PromptTokens, CachedTokens: resp.Usage.CachedTokens,
		CacheCreationTokens: resp.Usage.CacheCreationTokens,
		CompletionTokens:    resp.Usage.CompletionTokens, ReasoningTokens: resp.Usage.ReasoningTokens,
		CostUSD: cost, LatencyMS: latency, FinishReason: resp.FinishReason, OK: true,
	})
	if err := r.Store.SetJobStatus(job.ID, "done"); err != nil {
		return nil, err
	}
	r.Log.InfoContext(ctx, "stage completed", "stage", st.Name, "model", modelActual,
		"cost_usd", fmt.Sprintf("%.6f", cost), "latency_ms", latency,
		"prompt_tokens", resp.Usage.PromptTokens, "completion_tokens", resp.Usage.CompletionTokens)

	return &StageResult{Stage: st.Name, Role: st.Role, Model: modelActual,
		Usage: resp.Usage, CostUSD: cost, LatencyMS: latency,
		FinishReason: resp.FinishReason, Text: resp.Text}, nil
}