package claude

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"g.fc.run/ai/byre/internal/mcp"
)

type HookFunc func(ctx context.Context, input json.RawMessage) (map[string]any, error)

type HookMatcher struct {
	Matcher string
	Func    HookFunc
}

type Options struct {
	Cwd         string
	Model       string
	SessionID   string
	Resume      string
	ResumeAt    string
	ForkSession bool
	GitName     string
	GitEmail    string
	Env         []string
	Guidance    string
	Hooks       map[string][]HookMatcher
	MCPServers  []*mcp.Server
	StdioMCP    []mcp.Stdio
	Transcript  io.Writer
	OnBadLine   func(line []byte, reason string)
}

var knownTypes = map[string]bool{
	"assistant":              true,
	"user":                   true,
	"result":                 true,
	"system":                 true,
	"stream_event":           true,
	"tool_progress":          true,
	"rate_limit_event":       true,
	"control_request":        true,
	"control_response":       true,
	"control_cancel_request": true,
	"keep_alive":             true,
}

type Event struct {
	Type    string
	Subtype string
	Raw     json.RawMessage
	Read    time.Time
}

type Block struct {
	Type      string          `json:"type"`
	ID        string          `json:"id"`
	ToolUseID string          `json:"tool_use_id"`
	Text      string          `json:"text"`
	Thinking  string          `json:"thinking"`
	Name      string          `json:"name"`
	Input     json.RawMessage `json:"input"`
	Content   json.RawMessage `json:"content"`
	IsError   bool            `json:"is_error"`
}

type Result struct {
	Subtype      string                `json:"subtype"`
	IsError      bool                  `json:"is_error"`
	Result       string                `json:"result"`
	DurationMS   int                   `json:"duration_ms"`
	NumTurns     int                   `json:"num_turns"`
	TotalCostUSD float64               `json:"total_cost_usd"`
	ModelUsage   map[string]ModelUsage `json:"modelUsage"`
}

type ModelUsage struct {
	ContextWindow int `json:"contextWindow"`
}

type ContextUsage struct {
	Model  string
	Tokens int
}

func (e Event) ContextUsage() (ContextUsage, error) {
	msg := struct {
		Message struct {
			Model string `json:"model"`
			Usage struct {
				InputTokens              int `json:"input_tokens"`
				CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
				CacheReadInputTokens     int `json:"cache_read_input_tokens"`
				OutputTokens             int `json:"output_tokens"`
			} `json:"usage"`
		} `json:"message"`
	}{}
	err := json.Unmarshal(e.Raw, &msg)
	if err != nil {
		return ContextUsage{}, err
	}
	u := msg.Message.Usage
	return ContextUsage{
		Model:  msg.Message.Model,
		Tokens: u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens + u.OutputTokens,
	}, nil
}

func (e Event) Blocks() ([]Block, error) {
	msg := struct {
		Message struct {
			Content json.RawMessage `json:"content"`
		} `json:"message"`
	}{}
	err := json.Unmarshal(e.Raw, &msg)
	if err != nil {
		return nil, err
	}

	if len(msg.Message.Content) > 0 && msg.Message.Content[0] == '"' {
		text := ""
		err = json.Unmarshal(msg.Message.Content, &text)
		if err != nil {
			return nil, err
		}
		return []Block{{Type: "text", Text: text}}, nil
	}

	blocks := []Block{}
	err = json.Unmarshal(msg.Message.Content, &blocks)
	if err != nil {
		return nil, err
	}
	return blocks, nil
}

func (e Event) ParentToolUseID() (string, error) {
	msg := struct {
		ParentToolUseID string `json:"parent_tool_use_id"`
	}{}
	err := json.Unmarshal(e.Raw, &msg)
	return msg.ParentToolUseID, err
}

func (e Event) UUID() (string, error) {
	msg := struct {
		UUID string `json:"uuid"`
	}{}
	err := json.Unmarshal(e.Raw, &msg)
	return msg.UUID, err
}

func (e Event) TextDelta() (string, error) {
	msg := struct {
		ParentToolUseID *string `json:"parent_tool_use_id"`
		Event           struct {
			Type  string `json:"type"`
			Delta struct {
				Type string `json:"type"`
				Text string `json:"text"`
			} `json:"delta"`
		} `json:"event"`
	}{}
	err := json.Unmarshal(e.Raw, &msg)
	if err != nil {
		return "", err
	}
	if msg.ParentToolUseID != nil || msg.Event.Type != "content_block_delta" || msg.Event.Delta.Type != "text_delta" {
		return "", nil
	}
	return msg.Event.Delta.Text, nil
}

func (e Event) Result() (Result, error) {
	res := Result{}
	err := json.Unmarshal(e.Raw, &res)
	return res, err
}

type Session struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	events chan Event
	done   chan struct{}

	writeMu    sync.Mutex
	transcript io.Writer
	scrub      []byte

	pendingMu sync.Mutex
	pending   map[string]chan controlResult

	ctx        context.Context
	stop       context.CancelFunc
	inflightMu sync.Mutex
	inflight   map[string]context.CancelFunc

	hooks   map[string]HookFunc
	servers map[string]*mcp.Server

	idMu sync.Mutex
	id   string

	onBadLine func(line []byte, reason string)
	lastLine  atomic.Int64
}

type controlResult struct {
	response json.RawMessage
	err      error
}

type envelope struct {
	Type      string          `json:"type"`
	Subtype   string          `json:"subtype"`
	SessionID string          `json:"session_id"`
	RequestID string          `json:"request_id"`
	Request   json.RawMessage `json:"request"`
	Response  json.RawMessage `json:"response"`
}

func Start(ctx context.Context, opts Options) (*Session, error) {
	s := &Session{
		events:     make(chan Event, 256),
		done:       make(chan struct{}),
		transcript: opts.Transcript,
		pending:    map[string]chan controlResult{},
		inflight:   map[string]context.CancelFunc{},
		hooks:      map[string]HookFunc{},
		servers:    map[string]*mcp.Server{},
		id:         initialID(opts),
		onBadLine:  opts.OnBadLine,
	}
	s.lastLine.Store(time.Now().UnixNano())
	s.ctx, s.stop = context.WithCancel(context.Background())

	args, err := s.buildArgs(opts)
	if err != nil {
		return nil, err
	}

	s.cmd = exec.Command("claude", args...)
	s.cmd.Dir = opts.Cwd
	s.cmd.Stderr = os.Stderr
	s.cmd.Env = buildEnv(opts)
	s.cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	s.stdin, err = s.cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	stdout, err := s.cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

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

	go s.readLoop(stdout)

	_, err = s.Control(ctx, map[string]any{
		"subtype": "initialize",
		"hooks":   s.hookConfig(opts.Hooks),
	})
	if err != nil {
		s.Close()
		return nil, fmt.Errorf("initialize: %w", err)
	}

	return s, nil
}

func initialID(opts Options) string {
	if opts.SessionID != "" {
		return opts.SessionID
	}
	if opts.ForkSession {
		return ""
	}
	return opts.Resume
}

func (s *Session) buildArgs(opts Options) ([]string, error) {
	settings, err := json.Marshal(map[string]any{
		"attribution": map[string]any{"commit": "", "pr": ""},
	})
	if err != nil {
		return nil, err
	}

	servers := map[string]any{}
	for _, srv := range opts.MCPServers {
		s.servers[srv.Name] = srv
		servers[srv.Name] = map[string]any{"type": "sdk", "name": srv.Name, "alwaysLoad": true}
	}
	for _, srv := range opts.StdioMCP {
		servers[srv.Name] = map[string]any{"type": "stdio", "command": srv.Command, "args": srv.Args}
	}

	mcpConfig, err := json.Marshal(map[string]any{"mcpServers": servers})
	if err != nil {
		return nil, err
	}

	args := []string{
		"--output-format", "stream-json",
		"--input-format", "stream-json",
		"--verbose",
		"--include-partial-messages",
		"--permission-mode", "bypassPermissions",
		"--allow-dangerously-skip-permissions",
		"--settings", string(settings),
		"--mcp-config", string(mcpConfig),
	}

	if opts.Model != "" {
		args = append(args, "--model", opts.Model)
	}

	if opts.Guidance != "" {
		args = append(args, "--append-system-prompt", opts.Guidance)
	}

	if opts.SessionID != "" {
		args = append(args, "--session-id="+opts.SessionID)
	}

	if opts.Resume != "" {
		args = append(args, "--resume="+opts.Resume)
	}

	if opts.ResumeAt != "" {
		args = append(args, "--resume-session-at="+opts.ResumeAt)
	}

	if opts.ForkSession {
		args = append(args, "--fork-session")
	}

	return args, nil
}

func buildEnv(opts Options) []string {
	env := []string{}
	for _, kv := range os.Environ() {
		if len(kv) >= 11 && kv[:11] == "CLAUDECODE=" {
			continue
		}
		env = append(env, kv)
	}

	env = append(env,
		"CLAUDE_CODE_ENTRYPOINT=sdk-go",
		"CLAUDE_CODE_DISABLE_AUTOUPDATER=1",
		"GIT_AUTHOR_NAME="+opts.GitName,
		"GIT_AUTHOR_EMAIL="+opts.GitEmail,
		"GIT_COMMITTER_NAME="+opts.GitName,
		"GIT_COMMITTER_EMAIL="+opts.GitEmail,
	)
	return append(env, opts.Env...)
}

func (s *Session) hookConfig(hooks map[string][]HookMatcher) map[string]any {
	if len(hooks) == 0 {
		return nil
	}

	cfg := map[string]any{}
	n := 0
	for event, matchers := range hooks {
		list := []map[string]any{}
		for _, m := range matchers {
			id := fmt.Sprintf("hook_%d", n)
			n++
			s.hooks[id] = m.Func
			list = append(list, map[string]any{
				"matcher":         m.Matcher,
				"hookCallbackIds": []string{id},
			})
		}
		cfg[event] = list
	}

	return cfg
}

func (s *Session) Events() <-chan Event {
	return s.events
}

func (s *Session) Done() <-chan struct{} {
	return s.done
}

func (s *Session) Pid() int {
	return s.cmd.Process.Pid
}

func (s *Session) LastLine() time.Time {
	return time.Unix(0, s.lastLine.Load())
}

func (s *Session) Exited() bool {
	select {
	case <-s.done:
		return true
	default:
		return false
	}
}

func (s *Session) ID() string {
	s.idMu.Lock()
	defer s.idMu.Unlock()
	return s.id
}

type Image struct {
	MediaType string
	Data      []byte
}

func (s *Session) Send(text string, images ...Image) error {
	var content any = text
	if len(images) > 0 {
		blocks := []map[string]any{}
		for _, img := range images {
			blocks = append(blocks, map[string]any{
				"type": "image",
				"source": map[string]any{
					"type":       "base64",
					"media_type": img.MediaType,
					"data":       base64.StdEncoding.EncodeToString(img.Data),
				},
			})
		}
		if text != "" {
			blocks = append(blocks, map[string]any{"type": "text", "text": text})
		}
		content = blocks
	}

	return s.write(map[string]any{
		"type":               "user",
		"message":            map[string]any{"role": "user", "content": content},
		"parent_tool_use_id": nil,
		"session_id":         s.ID(),
	})
}

func (s *Session) Interrupt(ctx context.Context) error {
	_, err := s.Control(ctx, map[string]any{"subtype": "interrupt"})
	return err
}

func (s *Session) Control(ctx context.Context, req map[string]any) (json.RawMessage, error) {
	id := newRequestID()
	ch := make(chan controlResult, 1)

	s.pendingMu.Lock()
	s.pending[id] = ch
	s.pendingMu.Unlock()

	defer func() {
		s.pendingMu.Lock()
		delete(s.pending, id)
		s.pendingMu.Unlock()
	}()

	err := s.write(map[string]any{
		"type":       "control_request",
		"request_id": id,
		"request":    req,
	})
	if err != nil {
		return nil, err
	}

	select {
	case res := <-ch:
		return res.response, res.err
	case <-s.done:
		return nil, errors.New("claude exited")
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (s *Session) Close() error {
	s.writeMu.Lock()
	s.stdin.Close()
	s.writeMu.Unlock()

	select {
	case <-s.done:
	case <-time.After(5 * time.Second):
		syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
		<-s.done
	}

	return s.cmd.Wait()
}

func (s *Session) Kill() {
	syscall.Kill(-s.cmd.Process.Pid, syscall.SIGKILL)
	<-s.done
	s.cmd.Wait()
}

func (s *Session) write(v any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}

	s.writeMu.Lock()
	defer s.writeMu.Unlock()

	s.record("send", b)

	_, err = s.stdin.Write(append(b, '\n'))
	return err
}

func (s *Session) record(dir string, line []byte) {
	if s.transcript == nil {
		return
	}

	if len(s.scrub) > 0 {
		line = bytes.ReplaceAll(line, s.scrub, []byte("[redacted]"))
	}

	b, err := json.Marshal(map[string]any{"dir": dir, "msg": json.RawMessage(line)})
	if err != nil {
		slog.Error("record transcript", "err", err)
		return
	}

	s.transcript.Write(append(b, '\n'))
}

func (s *Session) readLoop(stdout io.Reader) {
	defer close(s.done)
	defer close(s.events)
	defer s.stop()

	r := bufio.NewReader(stdout)
	for {
		line, err := r.ReadBytes('\n')
		line = bytes.TrimRight(line, "\r\n")
		if len(line) > 0 {
			s.lastLine.Store(time.Now().UnixNano())
		}
		if len(line) > 0 && line[0] == '{' {
			s.handleLine(line)
		} else if len(line) > 0 {
			s.badLine(line, "not a JSON object")
		}
		if err != nil {
			if err != io.EOF {
				slog.Error("read claude stdout", "err", err)
			}
			return
		}
	}
}

func (s *Session) handleLine(line []byte) {
	read := time.Now()
	env := envelope{}
	err := json.Unmarshal(line, &env)
	if err != nil {
		s.badLine(line, err.Error())
		return
	}

	if !knownTypes[env.Type] {
		s.badLine(line, fmt.Sprintf("unrecognised type %q", env.Type))
	}

	s.writeMu.Lock()
	if env.Type == "control_response" {
		s.learnScrub(env.Response)
	}
	if env.Type != "stream_event" {
		s.record("recv", line)
	}
	s.writeMu.Unlock()

	switch env.Type {
	case "control_response":
		s.handleControlResponse(env.Response)
		return

	case "control_request":
		ctx, cancel := context.WithCancel(s.ctx)
		s.inflightMu.Lock()
		s.inflight[env.RequestID] = cancel
		s.inflightMu.Unlock()
		go s.handleControlRequest(ctx, cancel, env.RequestID, env.Request)
		return

	case "control_cancel_request":
		s.inflightMu.Lock()
		cancel := s.inflight[env.RequestID]
		s.inflightMu.Unlock()
		if cancel != nil {
			cancel()
		}
		return

	case "keep_alive":
		return
	}

	if env.Type == "system" && env.Subtype == "init" {
		s.idMu.Lock()
		s.id = env.SessionID
		s.idMu.Unlock()
	}

	s.events <- Event{Type: env.Type, Subtype: env.Subtype, Raw: append(json.RawMessage{}, line...), Read: read}
}

func (s *Session) badLine(line []byte, reason string) {
	slog.Warn("bad claude line", "reason", reason, "line", string(line))
	if s.onBadLine != nil {
		s.onBadLine(line, reason)
	}
}

func (s *Session) learnScrub(raw json.RawMessage) {
	resp := struct {
		Response struct {
			Account struct {
				Email string `json:"email"`
			} `json:"account"`
		} `json:"response"`
	}{}
	err := json.Unmarshal(raw, &resp)
	if err != nil || resp.Response.Account.Email == "" {
		return
	}
	s.scrub = []byte(resp.Response.Account.Email)
}

func (s *Session) handleControlResponse(raw json.RawMessage) {
	resp := struct {
		Subtype   string          `json:"subtype"`
		RequestID string          `json:"request_id"`
		Response  json.RawMessage `json:"response"`
		Error     string          `json:"error"`
	}{}
	err := json.Unmarshal(raw, &resp)
	if err != nil {
		slog.Error("parse control response", "err", err)
		return
	}

	s.pendingMu.Lock()
	ch := s.pending[resp.RequestID]
	s.pendingMu.Unlock()

	if ch == nil {
		slog.Warn("control response for unknown request", "request_id", resp.RequestID)
		return
	}

	if resp.Subtype == "error" {
		ch <- controlResult{err: errors.New(resp.Error)}
		return
	}

	ch <- controlResult{response: resp.Response}
}

func (s *Session) handleControlRequest(ctx context.Context, cancel context.CancelFunc, id string, raw json.RawMessage) {
	defer func() {
		s.inflightMu.Lock()
		delete(s.inflight, id)
		s.inflightMu.Unlock()
		cancel()
	}()

	result, err := s.dispatchControl(ctx, raw)
	if err != nil {
		slog.Error("control request failed", "request_id", id, "err", err)
		s.write(map[string]any{
			"type": "control_response",
			"response": map[string]any{
				"subtype":    "error",
				"request_id": id,
				"error":      err.Error(),
			},
		})
		return
	}

	s.write(map[string]any{
		"type": "control_response",
		"response": map[string]any{
			"subtype":    "success",
			"request_id": id,
			"response":   result,
		},
	})
}

func (s *Session) dispatchControl(ctx context.Context, raw json.RawMessage) (any, error) {
	req := struct {
		Subtype    string          `json:"subtype"`
		CallbackID string          `json:"callback_id"`
		Input      json.RawMessage `json:"input"`
		ServerName string          `json:"server_name"`
		Message    json.RawMessage `json:"message"`
	}{}
	err := json.Unmarshal(raw, &req)
	if err != nil {
		return nil, err
	}

	switch req.Subtype {
	case "hook_callback":
		fn := s.hooks[req.CallbackID]
		if fn == nil {
			return nil, fmt.Errorf("no hook callback %s", req.CallbackID)
		}
		return fn(ctx, req.Input)

	case "mcp_message":
		srv := s.servers[req.ServerName]
		if srv == nil {
			return nil, fmt.Errorf("no mcp server %s", req.ServerName)
		}
		resp, err := srv.Handle(ctx, req.Message)
		if err != nil {
			return nil, err
		}
		return map[string]any{"mcp_response": resp}, nil

	case "can_use_tool":
		return map[string]any{"behavior": "allow", "updatedInput": req.Input}, nil
	}

	return nil, fmt.Errorf("unsupported control request %s", req.Subtype)
}

func newRequestID() string {
	b := make([]byte, 8)
	rand.Read(b)
	return "req_" + hex.EncodeToString(b)
}