package runtime

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/krow/krow-backend/go-api/internal/repo"
)

// PostgresSink writes trajectories to agent_runs.
//
// Hand-written rather than built through repo.Repo's descriptor machinery, and
// deliberately so: that layer exists to serve the generic CRUD contract in
// docs/api-contract.md — filters, sorts, pagination, resource descriptors —
// and a trajectory has none of those. It is written once, whole, by the
// runtime, and read back by run id. One INSERT with explicit bind parameters
// is the honest shape for that, and inventing a resource descriptor to reach
// it would add a layer that only ever gets used one way.
type PostgresSink struct {
	db repo.Querier
}

// NewPostgresSink builds a sink over a pool or transaction.
func NewPostgresSink(db repo.Querier) *PostgresSink {
	return &PostgresSink{db: db}
}

var _ Sink = (*PostgresSink)(nil)

const insertRunSQL = `
INSERT INTO agent_runs (
  run_id, parent_run_id, org_id, user_id,
  agent_id, agent_version, tier, model,
  started_at, ended_at, termination, entries,
  input_tokens, output_tokens, cached_tokens, total_tokens, model_calls
) VALUES (
  $1, $2, $3::uuid, $4::uuid,
  $5, $6, $7, $8,
  $9, $10, $11, $12::jsonb,
  $13, $14, $15, $16, $17
)
ON CONFLICT (run_id) DO NOTHING`

// Save writes one finished trajectory.
//
// ON CONFLICT DO NOTHING because a run id is generated once and written once:
// a conflict means a retry of a save that already landed, and the first write
// is the authoritative one. Failing the second attempt would turn a harmless
// duplicate into a lost answer, since the caller treats a save error as
// something to report.
func (s *PostgresSink) Save(ctx context.Context, t *Trajectory) error {
	if t == nil {
		return fmt.Errorf("runtime: no trajectory to save")
	}
	if t.RunID == "" {
		return fmt.Errorf("runtime: a trajectory needs a run id")
	}
	if !t.Termination.Valid() {
		return fmt.Errorf("runtime: %q is not a termination reason", t.Termination)
	}
	// The table's tenancy column is NOT NULL, and a run with no organization is
	// a bug upstream rather than a row to write. Caught here so the failure
	// names the cause instead of surfacing as a constraint violation.
	if t.OrgID == "" {
		return fmt.Errorf("runtime: a trajectory needs an org id")
	}

	entries, err := json.Marshal(t.Entries)
	if err != nil {
		return fmt.Errorf("runtime: encoding trajectory entries: %w", err)
	}
	// A nil slice marshals to "null", which the jsonb_typeof CHECK refuses.
	// A run that recorded nothing is still a run worth keeping.
	if len(t.Entries) == 0 {
		entries = []byte("[]")
	}

	_, err = s.db.Exec(ctx, insertRunSQL,
		t.RunID,
		nullIfEmpty(t.ParentRunID),
		t.OrgID,
		nullIfEmpty(t.UserID),
		t.AgentID,
		t.AgentVersion,
		t.Tier,
		t.Model,
		t.StartedAt,
		t.EndedAt,
		string(t.Termination),
		string(entries),
		t.Usage.InputTokens,
		t.Usage.OutputTokens,
		t.Usage.CachedTokens,
		t.Usage.TotalTokens,
		t.Usage.ModelCalls,
	)
	if err != nil {
		return fmt.Errorf("runtime: saving trajectory %s: %w", t.RunID, err)
	}
	return nil
}

// nullIfEmpty keeps an empty optional out of a uuid column, where "" is not a
// value the type accepts.
func nullIfEmpty(s string) any {
	if s == "" {
		return nil
	}
	return s
}