package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
)

// CLIEvent represents a parsed NDJSON event from claude CLI stdout.
type CLIEvent struct {
	Type    string `json:"type"`
	Subtype string `json:"subtype,omitempty"`

	// For system init event
	SessionID string `json:"session_id,omitempty"`

	// For assistant message events
	Message *CLIMessage `json:"message,omitempty"`

	// For stream_event wrapper
	Event *StreamEvent `json:"event,omitempty"`

	// For result events — can be object or string, use RawMessage
	RawResult json.RawMessage `json:"result,omitempty"`
	Result    *CLIResult      `json:"-"`

	// Top-level fields on result events
	TotalCostUSD float64 `json:"total_cost_usd,omitempty"`
	DurationMS   float64 `json:"duration_ms,omitempty"`
	NumTurns     int     `json:"num_turns,omitempty"`
}

// StreamEvent is the inner event inside a stream_event wrapper.
type StreamEvent struct {
	Type         string        `json:"type"`
	Index        int           `json:"index,omitempty"`
	Delta        *StreamDelta  `json:"delta,omitempty"`
	ContentBlock *ContentBlock `json:"content_block,omitempty"`
}

type ContentBlock struct {
	Type string `json:"type"`
	Text string `json:"text,omitempty"`
}

type StreamDelta struct {
	Type       string `json:"type,omitempty"`
	Text       string `json:"text,omitempty"`
	StopReason string `json:"stop_reason,omitempty"`
}

type CLIMessage struct {
	Role    string       `json:"role,omitempty"`
	Content []CLIContent `json:"content,omitempty"`
}

type CLIContent struct {
	Type  string `json:"type"`
	Text  string `json:"text,omitempty"`
	Name  string `json:"name,omitempty"`
	Input any    `json:"input,omitempty"`
	ID    string `json:"id,omitempty"`
}

type CLIResult struct {
	Duration  float64 `json:"duration_ms,omitempty"`
	NumTurns  int     `json:"num_turns,omitempty"`
	CostUSD   float64 `json:"total_cost_usd,omitempty"`
	SessionID string  `json:"session_id,omitempty"`
}

// CLIInputMessage is the JSON format for sending messages via stream-json input.
type CLIInputMessage struct {
	Type    string          `json:"type"`
	Message CLIInputContent `json:"message"`
}

type CLIInputContent struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// CLIProcess manages a running claude CLI process.
type CLIProcess struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	stdout io.ReadCloser
	stderr io.ReadCloser

	Events       chan CLIEvent
	Errors       chan error
	CLISessionID string // session_id from init event, for resume

	done chan struct{}
	mu   sync.Mutex
}

// SpawnCLI starts a new claude CLI process for the given project directory.
// Uses --input-format stream-json for multi-turn conversation support.
func SpawnCLI(projectDir string) (*CLIProcess, error) {
	args := []string{
		"-p",
		"--input-format", "stream-json",
		"--output-format", "stream-json",
		"--verbose",
		"--include-partial-messages",
	}

	cmd := exec.Command("claude", args...)
	cmd.Dir = projectDir

	// Filter all Claude-related env vars to prevent nested session detection
	env := filterEnvMulti(os.Environ(), []string{
		"CLAUDECODE",
		"CLAUDE_CODE_ENTRYPOINT",
	})
	cmd.Env = env

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdin pipe: %w", err)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}

	cp := &CLIProcess{
		cmd:    cmd,
		stdin:  stdin,
		stdout: stdout,
		stderr: stderr,
		Events: make(chan CLIEvent, 100),
		Errors: make(chan error, 10),
		done:   make(chan struct{}),
	}

	go cp.readOutput()
	go cp.readErrors()

	return cp, nil
}

// Send writes a user message to the claude CLI process stdin as stream-json.
func (cp *CLIProcess) Send(message string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	msg := CLIInputMessage{
		Type: "user",
		Message: CLIInputContent{
			Role:    "user",
			Content: strings.TrimSpace(message),
		},
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal message: %w", err)
	}

	_, err = cp.stdin.Write(append(data, '\n'))
	return err
}

// Close terminates the claude CLI process.
func (cp *CLIProcess) Close() error {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	cp.stdin.Close()
	if cp.cmd.Process != nil {
		return cp.cmd.Process.Kill()
	}
	return nil
}

// Done returns a channel that's closed when the process exits.
func (cp *CLIProcess) Done() <-chan struct{} {
	return cp.done
}

func (cp *CLIProcess) readOutput() {
	defer close(cp.done)
	defer close(cp.Events)

	scanner := bufio.NewScanner(cp.stdout)
	scanner.Buffer(make([]byte, 0, 1024*1024), 1024*1024) // 1MB buffer

	for scanner.Scan() {
		line := scanner.Text()
		if line == "" {
			continue
		}

		var event CLIEvent
		if err := json.Unmarshal([]byte(line), &event); err != nil {
			cp.Errors <- fmt.Errorf("parse event: %w (line: %s)", err, truncate(line, 200))
			continue
		}

		// Parse result field — can be string or object
		if event.Type == "result" {
			if len(event.RawResult) > 0 {
				var result CLIResult
				if err := json.Unmarshal(event.RawResult, &result); err == nil {
					event.Result = &result
				}
			}
			// Always build Result from top-level fields (they're always present on result events)
			if event.Result == nil {
				event.Result = &CLIResult{}
			}
			if event.TotalCostUSD > 0 {
				event.Result.CostUSD = event.TotalCostUSD
			}
			if event.DurationMS > 0 {
				event.Result.Duration = event.DurationMS
			}
			if event.NumTurns > 0 {
				event.Result.NumTurns = event.NumTurns
			}
		}

		// Capture session_id from init event
		if event.Type == "system" && event.Subtype == "init" && event.SessionID != "" {
			cp.CLISessionID = event.SessionID
		}

		cp.Events <- event
	}

	if err := scanner.Err(); err != nil {
		cp.Errors <- fmt.Errorf("scanner: %w", err)
	}
}

func (cp *CLIProcess) readErrors() {
	scanner := bufio.NewScanner(cp.stderr)
	for scanner.Scan() {
		line := scanner.Text()
		if line != "" {
			cp.Errors <- fmt.Errorf("stderr: %s", line)
		}
	}
}

// filterEnvMulti returns a copy of env with all named variables removed.
func filterEnvMulti(env []string, names []string) []string {
	prefixes := make([]string, len(names))
	for i, n := range names {
		prefixes[i] = n + "="
	}

	result := make([]string, 0, len(env))
	for _, e := range env {
		skip := false
		for _, p := range prefixes {
			if strings.HasPrefix(e, p) {
				skip = true
				break
			}
		}
		if !skip {
			result = append(result, e)
		}
	}
	return result
}

// filterEnv returns a copy of env with the named variable removed.
func filterEnv(env []string, name string) []string {
	return filterEnvMulti(env, []string{name})
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}