package agent

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
)

const (
	kimiACPProtocolVersion = 1
	kimiACPRequestTimeout  = 30 * time.Second
)

var errKimiACPTransportAmbiguous = errors.New("kimi ACP transport outcome is ambiguous")

// KimiACPDriver drives one long-lived Kimi Code session through the public
// Agent Client Protocol exposed by `kimi acp`.
type KimiACPDriver struct {
	mu      sync.Mutex
	writeMu sync.Mutex

	stdin  io.WriteCloser
	output *kimiTranslatedOutput

	nextID  int
	pending map[int]chan kimiACPEnvelope

	sessionID             string
	sessionResumeFallback *drivercontract.SessionResumeFallback
	resumeCapable         bool
	promptInFlight        bool
	promptRequest         int
	promptGate            chan struct{}
	promptAccepted        bool
	toolTitles            map[string]string

	requestFuncForTest func(context.Context, string, any) (json.RawMessage, error)
}

type kimiACPEnvelope struct {
	JSONRPC string                `json:"jsonrpc,omitempty"`
	ID      json.RawMessage       `json:"id,omitempty"`
	Method  string                `json:"method,omitempty"`
	Params  json.RawMessage       `json:"params,omitempty"`
	Result  json.RawMessage       `json:"result,omitempty"`
	Error   *kimiACPResponseError `json:"error,omitempty"`
}

type kimiACPResponseError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

func (e *kimiACPResponseError) Error() string {
	if e == nil {
		return "kimi ACP error"
	}
	if strings.TrimSpace(e.Message) != "" {
		return e.Message
	}
	return fmt.Sprintf("JSON-RPC error %d", e.Code)
}

func (d *KimiACPDriver) Name() string { return "kimi" }

func (d *KimiACPDriver) SupportsStdinNotification() bool { return true }

func (d *KimiACPDriver) MCPToolPrefix() string { return "mcp__agentteam__" }

func (d *KimiACPDriver) UsesExplicitTurnLifecycle() bool { return true }

func (d *KimiACPDriver) EncodeStdinMessage(string, string) []byte { return []byte("{}\n") }

func (d *KimiACPDriver) currentSessionID() string {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.sessionID
}

func (d *KimiACPDriver) Spawn(ctx context.Context, cfg drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	carrier, err := cfg.RequireCompiledPrompt()
	if err != nil {
		return drivercontract.SpawnResult{}, err
	}
	cfg.CompiledPrompt = carrier
	if err := materializeContractCarrier(cfg); err != nil {
		return drivercontract.SpawnResult{}, fmt.Errorf("kimi: materialize contract carrier: %w", err)
	}
	prepared, err := prepareKimiACPLaunch(cfg)
	if err != nil {
		return drivercontract.SpawnResult{}, err
	}

	translatedReader, translatedWriter := io.Pipe()
	output := newKimiTranslatedOutput(512)
	output.attachWriter(translatedWriter)
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		defer translatedWriter.Close()
		for line := range output.lines {
			if _, err := io.WriteString(translatedWriter, line+"\n"); err != nil {
				return
			}
		}
	}()

	cmd, stdin, stdout, stderr, terminate, wait, err := startKimiACPProcess(ctx, cfg, prepared)
	if err != nil {
		prepared.Cleanup()
		output.close()
		_ = translatedReader.Close()
		return drivercontract.SpawnResult{}, err
	}

	d.mu.Lock()
	d.sessionResumeFallback = nil
	d.stdin = stdin
	d.output = output
	d.pending = make(map[int]chan kimiACPEnvelope)
	d.toolTitles = make(map[string]string)
	d.sessionID = ""
	d.promptInFlight = false
	d.promptRequest = 0
	d.promptGate = nil
	d.promptAccepted = false
	d.mu.Unlock()

	cleanupOnce := sync.Once{}
	cleanup := func() {
		cleanupOnce.Do(func() {
			d.closeTransport()
			terminate()
			output.close()
		})
	}
	preparedCleanupOnce := sync.Once{}
	cleanupPrepared := func() { preparedCleanupOnce.Do(prepared.Cleanup) }
	go d.readLoop(stdout, cleanup)

	sessionID, resumed, err := d.initializeSession(ctx, cfg, prepared.WorkDir)
	if err != nil {
		cleanup()
		_ = wait()
		cleanupPrepared()
		return drivercontract.SpawnResult{}, err
	}
	initialPrompt := kimiInitialPrompt(cfg, resumed)
	if strings.TrimSpace(initialPrompt) != "" || len(cfg.PromptAttachments) > 0 {
		go func() {
			input := drivercontract.RuntimeInput{
				ID:             cfg.InitialDeliveryMessageID,
				Resume:         resumed,
				CompiledPrompt: carrier,
				Attachments:    cfg.PromptAttachments,
			}
			if err := d.DeliverRuntimeInput(ctx, input, sessionID); err != nil {
				slog.Error("kimi: initial ACP prompt failed", "agent_id", cfg.AgentID, "error", err)
			}
		}()
	}

	d.mu.Lock()
	freshSessionFallback := d.sessionResumeFallback
	d.sessionResumeFallback = nil
	d.mu.Unlock()

	return drivercontract.SpawnResult{
		Cmd:                  cmd,
		Stdin:                stdin,
		Stdout:               translatedReader,
		StderrTail:           stderr.Tail,
		SessionID:            sessionID,
		SessionResumed:       resumed,
		FreshSessionFallback: freshSessionFallback,
		Terminate:            cleanup,
		Cleanup:              cleanup,
		Wait: func() error {
			err := wait()
			cleanup()
			<-readerDone
			cleanupPrepared()
			return err
		},
	}, nil
}

func (d *KimiACPDriver) initializeSession(
	ctx context.Context,
	cfg drivercontract.LaunchConfig,
	workDir string,
) (string, bool, error) {
	initializeResult, err := d.request(ctx, "initialize", map[string]any{
		"protocolVersion":    kimiACPProtocolVersion,
		"clientCapabilities": map[string]any{},
		"clientInfo": map[string]any{
			"name":    "AgentTeam Desktop",
			"version": "1",
		},
	})
	if err != nil {
		return "", false, fmt.Errorf("kimi: initialize ACP: %w", err)
	}
	resumeCapable, err := validateKimiACPInitializeResult(initializeResult)
	if err != nil {
		return "", false, err
	}
	d.mu.Lock()
	d.resumeCapable = resumeCapable
	d.mu.Unlock()

	sessionID, resumed, err := d.openSession(ctx, cfg, workDir)
	if err != nil {
		return "", false, fmt.Errorf("kimi: open ACP session: %w", err)
	}
	d.mu.Lock()
	d.sessionID = sessionID
	d.mu.Unlock()
	if err := d.setFullAccessMode(ctx, sessionID, cfg); err != nil {
		return "", false, err
	}
	d.emit(translatedEvent{Kind: consts.ParsedEventKindSessionInit, SessionID: sessionID})
	return sessionID, resumed, nil
}

func startKimiACPProcess(
	ctx context.Context,
	cfg drivercontract.LaunchConfig,
	prepared kimiACPPreparedLaunch,
) (*exec.Cmd, io.WriteCloser, io.ReadCloser, *capturingPrefixWriter, func(), func() error, error) {
	binary := strings.TrimSpace(cfg.Binary)
	if binary == "" {
		binary = "kimi"
	}
	cmd := runtimeconfig.RuntimeCommandContext(ctx, runtimeconfig.TypeKimi, binary, prepared.Args)
	cmd.Dir = prepared.WorkDir
	cmd.Env = prepared.Env
	stderr := newCapturingPrefixWriter(fmt.Sprintf("[kimi %s] ", cfg.AgentID), 32*1024)
	cmd.Stderr = stderr
	configureRuntimeCommand(cmd)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, nil, nil, nil, nil, fmt.Errorf("kimi: get ACP stdout: %w", err)
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		_ = stdout.Close()
		return nil, nil, nil, nil, nil, nil, fmt.Errorf("kimi: get ACP stdin: %w", err)
	}
	if err := cmd.Start(); err != nil {
		_ = stdout.Close()
		_ = stdin.Close()
		return nil, nil, nil, nil, nil, nil, fmt.Errorf("kimi: start ACP process: %w", err)
	}
	terminateOnce := sync.Once{}
	terminate := func() {
		terminateOnce.Do(func() {
			_ = stdin.Close()
			_ = stdout.Close()
			terminateRuntimeCommand(cmd)
		})
	}
	return cmd, stdin, stdout, stderr, terminate, cmd.Wait, nil
}

func (d *KimiACPDriver) request(ctx context.Context, method string, params any) (json.RawMessage, error) {
	if d.requestFuncForTest != nil {
		return d.requestFuncForTest(ctx, method, params)
	}
	requestCtx, cancel := context.WithTimeout(ctx, kimiACPRequestTimeout)
	defer cancel()
	id, responseCh, err := d.beginRequest(method, params)
	if err != nil {
		return nil, err
	}
	select {
	case <-requestCtx.Done():
		d.forgetPending(id)
		return nil, requestCtx.Err()
	case response := <-responseCh:
		if response.Error != nil {
			return nil, response.Error
		}
		return response.Result, nil
	}
}

func (d *KimiACPDriver) beginRequest(method string, params any) (int, chan kimiACPEnvelope, error) {
	d.mu.Lock()
	if d.stdin == nil {
		d.mu.Unlock()
		return 0, nil, io.ErrClosedPipe
	}
	d.nextID++
	id := d.nextID
	responseCh := make(chan kimiACPEnvelope, 1)
	d.pending[id] = responseCh
	d.mu.Unlock()

	payload := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
	if params != nil {
		payload["params"] = params
	}
	data, err := json.Marshal(payload)
	if err != nil {
		d.forgetPending(id)
		return 0, nil, err
	}
	if err := d.writeFrame(append(data, '\n'), nil); err != nil {
		d.forgetPending(id)
		return 0, nil, err
	}
	return id, responseCh, nil
}

func (d *KimiACPDriver) writeFrame(data []byte, accepted func()) error {
	d.writeMu.Lock()
	d.mu.Lock()
	stdin := d.stdin
	d.mu.Unlock()
	if stdin == nil {
		d.writeMu.Unlock()
		return io.ErrClosedPipe
	}
	n, err := stdin.Write(data)
	if n == len(data) && err == nil {
		if accepted != nil {
			accepted()
		}
		d.writeMu.Unlock()
		return nil
	}
	d.writeMu.Unlock()
	if err == nil {
		err = io.ErrShortWrite
	}
	if n > 0 {
		d.closeTransport()
		return errors.Join(errKimiACPTransportAmbiguous, err)
	}
	return err
}

func (d *KimiACPDriver) writeNotification(method string, params any) error {
	payload := map[string]any{"jsonrpc": "2.0", "method": method, "params": params}
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	return d.writeFrame(append(data, '\n'), nil)
}

func (d *KimiACPDriver) forgetPending(id int) {
	d.mu.Lock()
	delete(d.pending, id)
	d.mu.Unlock()
}

func (d *KimiACPDriver) closeTransport() {
	d.writeMu.Lock()
	defer d.writeMu.Unlock()

	d.mu.Lock()
	stdin := d.stdin
	d.stdin = nil
	promptID := d.promptRequest
	promptSessionID := d.sessionID
	promptGate := d.promptGate
	promptAccepted := d.promptAccepted
	if promptID != 0 {
		delete(d.pending, promptID)
		d.promptInFlight = false
		d.promptRequest = 0
		d.promptGate = nil
		d.promptAccepted = false
	}
	d.mu.Unlock()
	if stdin != nil {
		_ = stdin.Close()
	}
	if promptGate != nil && !promptAccepted {
		close(promptGate)
	}
	if promptAccepted {
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindError,
			Message:   "kimi ACP transport closed",
			SessionID: promptSessionID,
			Metadata:  map[string]any{"will_retry": false},
		})
		d.emit(translatedEvent{
			Kind:       consts.ParsedEventKindTurnEnd,
			SessionID:  promptSessionID,
			StopReason: consts.ParsedEventStopReasonEndTurn,
		})
	}
}

func (d *KimiACPDriver) readLoop(stdout io.Reader, cleanup func()) {
	defer func() {
		d.closeTransport()
		d.failPending(errors.New("kimi ACP transport closed"))
		cleanup()
	}()
	scanner := bufio.NewScanner(stdout)
	scanner.Buffer(make([]byte, 64*1024), maxProcessLineBytes)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}
		var envelope kimiACPEnvelope
		if err := json.Unmarshal([]byte(line), &envelope); err != nil {
			slog.Warn("kimi: invalid ACP frame", "error", err)
			continue
		}
		d.routeEnvelope(envelope)
	}
	if err := scanner.Err(); err != nil && !errors.Is(err, os.ErrClosed) {
		slog.Warn("kimi: ACP read failed", "error", err)
	}
}

func (d *KimiACPDriver) routeEnvelope(envelope kimiACPEnvelope) {
	if len(envelope.ID) > 0 && envelope.Method == "" {
		id, ok := parseKimiACPID(envelope.ID)
		if !ok {
			return
		}
		d.mu.Lock()
		ch := d.pending[id]
		gate := d.promptGate
		isPrompt := id == d.promptRequest && gate != nil
		if !isPrompt {
			delete(d.pending, id)
		}
		d.mu.Unlock()
		if isPrompt {
			<-gate
			d.finishPromptResponse(id, envelope)
			return
		}
		if ch != nil {
			ch <- envelope
		}
		return
	}
	if envelope.Method != "" && len(envelope.ID) > 0 {
		d.handleServerRequest(envelope)
		return
	}
	if envelope.Method != "" {
		if envelope.Method == "session/update" {
			d.mu.Lock()
			gate := d.promptGate
			d.mu.Unlock()
			if gate != nil {
				<-gate
			}
		}
		d.handleACPNotification(envelope.Method, envelope.Params)
	}
}

func parseKimiACPID(raw json.RawMessage) (int, bool) {
	var number int
	if json.Unmarshal(raw, &number) == nil {
		return number, true
	}
	var stringID string
	if json.Unmarshal(raw, &stringID) == nil {
		value, err := strconv.Atoi(strings.TrimSpace(stringID))
		return value, err == nil
	}
	return 0, false
}

func (d *KimiACPDriver) failPending(err error) {
	d.mu.Lock()
	pending := d.pending
	d.pending = make(map[int]chan kimiACPEnvelope)
	d.mu.Unlock()
	for _, ch := range pending {
		ch <- kimiACPEnvelope{Error: &kimiACPResponseError{Code: -32099, Message: err.Error()}}
	}
}

func validateKimiACPInitializeResult(result json.RawMessage) (bool, error) {
	var decoded struct {
		ProtocolVersion   int `json:"protocolVersion"`
		AgentCapabilities struct {
			SessionCapabilities struct {
				Resume *struct{} `json:"resume"`
			} `json:"sessionCapabilities"`
		} `json:"agentCapabilities"`
	}
	if err := json.Unmarshal(result, &decoded); err != nil {
		return false, fmt.Errorf("kimi: decode ACP initialize response: %w", err)
	}
	if decoded.ProtocolVersion != kimiACPProtocolVersion {
		return false, fmt.Errorf("kimi: unsupported ACP protocol version %d; require %d", decoded.ProtocolVersion, kimiACPProtocolVersion)
	}
	return decoded.AgentCapabilities.SessionCapabilities.Resume != nil, nil
}

func (d *KimiACPDriver) emit(event translatedEvent) {
	d.mu.Lock()
	output := d.output
	d.mu.Unlock()
	if output == nil {
		return
	}
	data, err := json.Marshal(event)
	if err == nil {
		output.emit(string(data))
	}
}

func (d *KimiACPDriver) ParseLine(line string) []drivercontract.ParsedEvent {
	var event translatedEvent
	if json.Unmarshal([]byte(strings.TrimSpace(line)), &event) != nil {
		return nil
	}
	parsed := drivercontract.ParsedEvent{
		Kind:                event.Kind,
		Text:                event.Text,
		SectionID:           event.SectionID,
		ToolName:            event.ToolName,
		ToolCallID:          event.ToolCallID,
		ToolInput:           event.ToolInput,
		RuntimeInputID:      event.RuntimeInputID,
		SessionID:           event.SessionID,
		Message:             event.Message,
		Metadata:            event.Metadata,
		StopReason:          event.StopReason,
		InputTokens:         event.InputTokens,
		CachedInputTokens:   event.CachedInputTokens,
		OutputTokens:        event.OutputTokens,
		ActiveContextTokens: event.ActiveContextTokens,
		ContextWindowTokens: event.ContextWindowTokens,
		Model:               event.Model,
	}
	switch event.Kind {
	case consts.ParsedEventKindSessionInit:
		if parsed.SessionID == "" {
			return nil
		}
	case consts.ParsedEventKindText, consts.ParsedEventKindThinking:
		if parsed.Text == "" && !parsedBoolMetadata(parsed.Metadata, codexAuthoritativeTurnStartMetadataKey) {
			return nil
		}
	case consts.ParsedEventKindToolCall, consts.ParsedEventKindToolOutput,
		consts.ParsedEventKindTurnEnd, consts.ParsedEventKindError,
		consts.ParsedEventKindUsage:
	default:
		return nil
	}
	return []drivercontract.ParsedEvent{parsed}
}

var _ drivercontract.Driver = (*KimiACPDriver)(nil)
var _ drivercontract.BusyStdinMessageDriver = (*KimiACPDriver)(nil)
var _ drivercontract.RuntimeInputDeliveryDriver = (*KimiACPDriver)(nil)
var _ drivercontract.ActiveRunInterruptDriver = (*KimiACPDriver)(nil)
var _ drivercontract.ExplicitTurnLifecycleDriver = (*KimiACPDriver)(nil)