package main

// What `-p --output-format ...` prints, in the shapes `claude -p` prints.
//
// The reader is a program: a CI step that runs `jq -r .result`, a wrapper that
// reads stream-json line by line to show progress, an SDK that was written
// against Claude Code. None of them can be asked to change, so the shapes are
// Claude Code's — the same type and subtype names, the same field names, the
// assistant message as the Messages API returns it — and where Codey has
// nothing to put in a field Claude Code fills, the field is left out rather
// than filled with something made up.
//
// Nothing else may reach stdout in these modes. Notices go to stderr, and so
// do approval prompts; see cliRun.UIOut.

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/llm"
)

// Result subtypes, as Claude Code names them.
const (
	resultSuccess   = "success"
	resultMaxTurns  = "error_max_turns"
	resultExecution = "error_during_execution"
)

// printer writes one run in a Claude Code output format. It is the agent's
// observer for the length of the run, and the event log of it as far as the
// reader on stdout is concerned.
type printer struct {
	mu     sync.Mutex
	w      io.Writer
	notes  io.Writer
	format string
	// array is json with --verbose, which Claude Code prints as every line
	// stream-json would have printed, collected into one JSON array.
	array   bool
	records []any

	session string
	model   string
	start   time.Time
	// api is the time spent waiting on the model, measured from the start of
	// each turn to the usage that ends it. Compaction is a model call too and
	// falls inside that window, which is where it belongs.
	api       time.Duration
	turnStart time.Time
	done      bool

	// now and newID are the clock and the uuid source, replaceable so a test
	// can compare output byte for byte.
	now   func() time.Time
	newID func() string
}

func newPrinter(w, notes io.Writer, format string, verbose bool) *printer {
	return &printer{
		w: w, notes: notes, format: format,
		array: format == formatJSON && verbose,
		start: time.Now(), now: time.Now, newID: uuid4,
	}
}

// claudeOut is the printer of a Claude-shaped -p run, once there is one. It is
// a package variable for main, which reports a run that failed before it got
// as far as the model — no credentials, a config error — and whose reader is
// still waiting for a result object.
var claudeOut *printer

// reportFailure gives the reader on stdout its result object when the run
// ended before printing one.
func reportFailure(err error) {
	if claudeOut != nil {
		claudeOut.finish(outcome{Err: err})
	}
}

// begin prints the init line, which only stream-json has.
func (p *printer) begin(cwd, model, mode string, tools []string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.model = model
	p.start = p.now()
	if p.format != formatStreamJSON && !p.array {
		return
	}
	p.emit(initLine{
		Type: "system", Subtype: "init", CWD: cwd, SessionID: p.session,
		Tools: tools, MCPServers: mcpServers(tools), Model: model,
		PermissionMode: mode, UUID: p.newID(),
	})
}

// observe is the agent's Observe callback.
func (p *printer) observe(ev agent.Event) {
	p.mu.Lock()
	defer p.mu.Unlock()
	switch ev.Kind {
	case agent.EventTurnStart:
		p.turnStart = p.now()
	case agent.EventUsage:
		if !p.turnStart.IsZero() {
			p.api += p.now().Sub(p.turnStart)
			p.turnStart = time.Time{}
		}
	case agent.EventNotice:
		fmt.Fprintln(p.notes, "codey: "+strings.TrimSpace(ev.Text))
	}
	if p.format != formatStreamJSON && !p.array {
		return
	}
	switch ev.Kind {
	case agent.EventMessage:
		if r := ev.Response; r != nil {
			model := r.Model
			if model == "" {
				model = p.model
			}
			p.emit(messageLine{
				Type: "assistant",
				Message: apiMessage{
					ID: r.ID, Type: "message", Role: llm.RoleAssistant, Model: model,
					Content: nonNil(r.Content), StopReason: r.StopReason, Usage: r.Usage,
				},
				SessionID: p.session, UUID: p.newID(),
			})
		}
		if m := ev.Message; m != nil {
			p.emit(messageLine{
				Type:      "user",
				Message:   userMessage{Role: m.Role, Content: nonNil(m.Content)},
				SessionID: p.session, UUID: p.newID(),
			})
		}
	case agent.EventCompact:
		p.emit(systemLine{Type: "system", Subtype: "compact_boundary", SessionID: p.session, UUID: p.newID()})
	}
}

// outcome is how a run ended, for the result object.
type outcome struct {
	Text  string
	Err   error
	Turns int
	Usage llm.Usage
}

// finish prints the result, once. A second call — main reporting the error
// the result already carries — prints nothing.
func (p *printer) finish(o outcome) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.done {
		return
	}
	p.done = true

	if p.format == formatText {
		// The answer and nothing else, which is what a script that captures
		// `claude -p` wants. A failure has no answer; its message goes to
		// stderr with the exit status.
		if o.Err == nil {
			fmt.Fprintln(p.w, o.Text)
		}
		return
	}

	res := resultLine{
		Type: "result", Subtype: resultSuccess,
		DurationMS: p.now().Sub(p.start).Milliseconds(), DurationAPIMS: p.api.Milliseconds(),
		NumTurns: o.Turns, Result: o.Text, SessionID: p.session, Usage: o.Usage, UUID: p.newID(),
	}
	if cost, ok := o.Usage.Cost(p.model); ok {
		res.TotalCostUSD = cost
	}
	if o.Err != nil {
		res.IsError = true
		res.Subtype = resultExecution
		if errors.Is(o.Err, agent.ErrMaxTurns) {
			res.Subtype = resultMaxTurns
		}
		res.Errors = []string{o.Err.Error()}
		// A run that stopped part-way still hands back what the model last
		// said, as the line front end does; a run with nothing to hand back
		// says why, because `.result` is the field everybody prints.
		if res.Result == "" {
			res.Result = o.Err.Error()
		}
	}
	p.emit(res)

	if p.array {
		p.write(p.records)
	}
}

// emit prints one record, or holds it for the array.
func (p *printer) emit(v any) {
	if p.array {
		p.records = append(p.records, v)
		return
	}
	p.write(v)
}

// write prints one JSON value on a line of its own. HTML is not escaped,
// because the reader is not a browser and Claude Code does not escape it:
// a diff full of < would match nothing a script greps for.
func (p *printer) write(v any) {
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(v); err != nil {
		fmt.Fprintln(p.notes, "codey: encoding output:", err)
		return
	}
	p.w.Write(unescapeHTML(buf.Bytes()))
}

// unescapeHTML undoes the escapes json.Marshal writes for <, > and & — a
// backslash, a u, and the code point in four hex digits (003c, 003e, 0026).
// The encoder's own setting does not reach them: a content block marshals
// itself, with json.Marshal, before the encoder sees it. The scan steps over
// every escape as a pair, so an escaped backslash followed by the letters
// u003c is left alone.
func unescapeHTML(b []byte) []byte {
	if !bytes.Contains(b, []byte(`\u00`)) {
		return b
	}
	out := make([]byte, 0, len(b))
	for i := 0; i < len(b); i++ {
		if b[i] != '\\' || i+1 == len(b) {
			out = append(out, b[i])
			continue
		}
		if b[i+1] == 'u' && i+6 <= len(b) {
			switch string(b[i+2 : i+6]) {
			case "003c":
				out, i = append(out, '<'), i+5
				continue
			case "003e":
				out, i = append(out, '>'), i+5
				continue
			case "0026":
				out, i = append(out, '&'), i+5
				continue
			}
		}
		out = append(out, b[i], b[i+1])
		i++
	}
	return out
}

type initLine struct {
	Type           string      `json:"type"`
	Subtype        string      `json:"subtype"`
	CWD            string      `json:"cwd"`
	SessionID      string      `json:"session_id"`
	Tools          []string    `json:"tools"`
	MCPServers     []mcpServer `json:"mcp_servers"`
	Model          string      `json:"model"`
	PermissionMode string      `json:"permissionMode"`
	UUID           string      `json:"uuid"`
}

type mcpServer struct {
	Name   string `json:"name"`
	Status string `json:"status"`
}

type systemLine struct {
	Type      string `json:"type"`
	Subtype   string `json:"subtype"`
	SessionID string `json:"session_id"`
	UUID      string `json:"uuid"`
}

// messageLine is an assistant or a user line. parent_tool_use_id is always
// null: Codey's subagents report through a tool result, not by speaking in
// the parent's stream, so nothing printed here belongs to a tool call.
type messageLine struct {
	Type            string  `json:"type"`
	Message         any     `json:"message"`
	ParentToolUseID *string `json:"parent_tool_use_id"`
	SessionID       string  `json:"session_id"`
	UUID            string  `json:"uuid"`
}

// apiMessage is an assistant message as the Messages API returns it.
type apiMessage struct {
	ID           string             `json:"id"`
	Type         string             `json:"type"`
	Role         string             `json:"role"`
	Model        string             `json:"model"`
	Content      []llm.ContentBlock `json:"content"`
	StopReason   string             `json:"stop_reason"`
	StopSequence *string            `json:"stop_sequence"`
	Usage        llm.Usage          `json:"usage"`
}

type userMessage struct {
	Role    string             `json:"role"`
	Content []llm.ContentBlock `json:"content"`
}

type resultLine struct {
	Type          string    `json:"type"`
	Subtype       string    `json:"subtype"`
	IsError       bool      `json:"is_error"`
	DurationMS    int64     `json:"duration_ms"`
	DurationAPIMS int64     `json:"duration_api_ms"`
	NumTurns      int       `json:"num_turns"`
	Result        string    `json:"result"`
	SessionID     string    `json:"session_id"`
	TotalCostUSD  float64   `json:"total_cost_usd"`
	Usage         llm.Usage `json:"usage"`
	Errors        []string  `json:"errors,omitempty"`
	UUID          string    `json:"uuid"`
}

// mcpServers names the servers whose tools are on offer. A server that failed
// to start has none, and is left out rather than reported in a state this
// list cannot tell apart from not configured.
func mcpServers(tools []string) []mcpServer {
	seen := map[string]bool{}
	for _, t := range tools {
		rest, ok := strings.CutPrefix(t, "mcp__")
		if !ok {
			continue
		}
		if name, _, ok := strings.Cut(rest, "__"); ok {
			seen[name] = true
		}
	}
	out := []mcpServer{}
	for name := range seen {
		out = append(out, mcpServer{Name: name, Status: "connected"})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

// nonNil keeps an empty content list printing as [] rather than null.
func nonNil(blocks []llm.ContentBlock) []llm.ContentBlock {
	if blocks == nil {
		return []llm.ContentBlock{}
	}
	return blocks
}

// uuid4 is a random version-4 UUID, which is what Claude Code stamps on every
// line. Nothing reads them back; they are there so a reader that keys lines by
// uuid has something to key on.
func uuid4() string {
	var b [16]byte
	rand.Read(b[:])
	b[6] = b[6]&0x0f | 0x40
	b[8] = b[8]&0x3f | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}

// printOnce is once for a Claude-shaped run: the prompt, then the result in
// the chosen format. The usage and turns are this run's alone, so a resumed
// session reports what this invocation cost rather than everything before it.
func printOnce(ctx context.Context, a *agent.Agent, p *printer, sessionDir, id, root string, client *llm.Client, prompt string) error {
	before, turnsBefore := a.State().Usage()

	turnCtx, endTurn := beginTurn(ctx)
	text, runErr := a.Run(turnCtx, prompt)
	endTurn()

	after, turnsAfter := a.State().Usage()
	if err := a.Save(sessionDir, id, root, "", client.Model()); err != nil {
		fmt.Fprintln(os.Stderr, "codey: saving session:", err)
	}
	if errors.Is(runErr, context.Canceled) {
		// Unlike the line front end, an interrupted -p run is a failure: the
		// script that started it did not get its answer. The session is saved
		// all the same, so -c picks it up.
		runErr = errors.New("interrupted — the session is saved; continue it with -c")
	}
	p.finish(outcome{Text: text, Err: runErr, Turns: turnsAfter - turnsBefore, Usage: usageSince(after, before)})
	return runErr
}

func usageSince(after, before llm.Usage) llm.Usage {
	return llm.Usage{
		InputTokens:              after.InputTokens - before.InputTokens,
		OutputTokens:             after.OutputTokens - before.OutputTokens,
		CacheReadInputTokens:     after.CacheReadInputTokens - before.CacheReadInputTokens,
		CacheCreationInputTokens: after.CacheCreationInputTokens - before.CacheCreationInputTokens,
	}
}