package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"

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

func (d *KimiACPDriver) recordSessionResumeFallback(reason, operation string) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.sessionResumeFallback != nil {
		return
	}
	d.sessionResumeFallback = newSessionResumeFallback(reason, operation)
}

const kimiSessionInvalidRPCCode = -32000

func isAuthoritativeKimiACPSessionInvalidError(method string, err error) bool {
	if strings.TrimSpace(method) != "session/resume" {
		return false
	}
	var responseErr *kimiACPResponseError
	if err == nil || !errors.As(err, &responseErr) {
		return false
	}
	// Validate ACP's typed server-error boundary before inspecting its data.
	// Invalid-params, transport, and provider/configuration failures cannot
	// authorize replacing a session even when their text contains a session id.
	if responseErr.Code != kimiSessionInvalidRPCCode {
		return false
	}
	if kimiACPErrorDataSessionInvalid(responseErr.Data) {
		return true
	}
	switch normalizeKimiACPErrorToken(responseErr.Message) {
	case "session not found", "session does not exist", "invalid session", "invalid session id", "unknown session":
		return true
	default:
		return false
	}
}

func kimiACPErrorDataSessionInvalid(data json.RawMessage) bool {
	if len(data) == 0 || strings.TrimSpace(string(data)) == "null" {
		return false
	}
	// ACP provider/model errors may expose reason/status/error/message fields;
	// only the direct structured discriminator fields below are session
	// authority, and only for the validated outer RPC code above.
	var decoded struct {
		Code string `json:"code"`
		Kind string `json:"kind"`
		Type string `json:"type"`
	}
	if json.Unmarshal(data, &decoded) != nil {
		return false
	}
	for _, value := range []string{decoded.Code, decoded.Kind, decoded.Type} {
		switch normalizeKimiACPErrorToken(value) {
		case "session not found", "invalid session id":
			return true
		}
	}
	return false
}

func normalizeKimiACPErrorToken(value string) string {
	value = strings.ToLower(strings.TrimSpace(value))
	value = strings.NewReplacer("_", " ", "-", " ").Replace(value)
	return strings.Join(strings.Fields(value), " ")
}

func (d *KimiACPDriver) openSession(ctx context.Context, cfg drivercontract.LaunchConfig, workDir string) (string, bool, error) {
	servers, err := kimiACPMCPServers(cfg.MCPs)
	if err != nil {
		return "", false, err
	}
	requested := strings.TrimSpace(cfg.SessionID)
	if requested != "" {
		d.mu.Lock()
		resumeCapable := d.resumeCapable
		d.mu.Unlock()
		if !resumeCapable {
			return "", false, fmt.Errorf("kimi: ACP runtime does not advertise session/resume")
		}
		_, err := d.request(ctx, "session/resume", map[string]any{
			"sessionId":  requested,
			"cwd":        workDir,
			"mcpServers": servers,
		})
		if err == nil {
			return requested, true, nil
		}
		if !isAuthoritativeKimiACPSessionInvalidError("session/resume", err) {
			return "", false, fmt.Errorf("resume %q: %w", requested, err)
		}
		d.recordSessionResumeFallback("invalid_or_not_found_session", "session/resume")
	}

	result, err := d.request(ctx, "session/new", map[string]any{
		"cwd":        workDir,
		"mcpServers": servers,
	})
	if err != nil {
		return "", false, err
	}
	var decoded struct {
		SessionID string `json:"sessionId"`
	}
	if err := json.Unmarshal(result, &decoded); err != nil {
		return "", false, fmt.Errorf("decode session/new response: %w", err)
	}
	if strings.TrimSpace(decoded.SessionID) == "" {
		return "", false, fmt.Errorf("session/new returned no sessionId")
	}
	return strings.TrimSpace(decoded.SessionID), false, nil
}

func (d *KimiACPDriver) setFullAccessMode(ctx context.Context, sessionID string, cfg drivercontract.LaunchConfig) error {
	if _, err := cfg.EffectivePermissionMode(); err != nil {
		return err
	}
	if _, err := d.request(ctx, "session/set_mode", map[string]any{
		"sessionId": sessionID,
		"modeId":    "yolo",
	}); err != nil {
		return fmt.Errorf("kimi: set full-access ACP mode: %w", err)
	}
	return nil
}

// DeliverBusyStdinInput exposes compiler-produced ACP prompt delivery to
// Process's protocol routing gate.
func (d *KimiACPDriver) DeliverBusyStdinInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	return d.DeliverRuntimeInput(ctx, input, sessionID)
}

// DeliverRuntimeInput writes one complete session/prompt frame and returns at
// the transport-acceptance boundary. The potentially long turn completes
// asynchronously through session/update notifications and the prompt response.
func (d *KimiACPDriver) DeliverRuntimeInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	if input.Empty() {
		return nil
	}
	current, err := d.validateSession(sessionID)
	if err != nil {
		return errors.Join(errRuntimeInputMutationNotAttempted, err)
	}
	blocks, err := kimiACPPromptBlocks(input)
	if err != nil {
		return errors.Join(errRuntimeInputMutationNotAttempted, err)
	}

	d.mu.Lock()
	if d.promptInFlight {
		d.mu.Unlock()
		return errors.Join(errRuntimeInputMutationNotAttempted, fmt.Errorf("kimi: prompt already in flight"))
	}
	if d.stdin == nil {
		d.mu.Unlock()
		return errors.Join(errRuntimeInputMutationNotAttempted, io.ErrClosedPipe)
	}
	d.promptInFlight = true
	d.nextID++
	id := d.nextID
	d.pending[id] = nil
	gate := make(chan struct{})
	d.promptRequest = id
	d.promptGate = gate
	d.promptAccepted = false
	d.mu.Unlock()

	payload := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"method":  "session/prompt",
		"params": map[string]any{
			"sessionId": current,
			"prompt":    blocks,
		},
	}
	data, err := json.Marshal(payload)
	if err != nil {
		d.abortPromptStart(id, gate)
		return errors.Join(errRuntimeInputMutationNotAttempted, err)
	}
	if err := d.writeFrame(append(data, '\n'), func() {
		d.mu.Lock()
		if d.promptRequest == id {
			d.promptAccepted = true
		}
		d.mu.Unlock()
		// ACP has no distinct applied-input acknowledgement. A complete request
		// frame is therefore the acceptance boundary. Publish the authoritative
		// turn start before allowing streamed updates or the response through.
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindThinking,
			SessionID: current,
			Metadata:  map[string]any{codexAuthoritativeTurnStartMetadataKey: true},
		})
		close(gate)
	}); err != nil {
		d.abortPromptStart(id, gate)
		if errors.Is(err, errKimiACPTransportAmbiguous) {
			return err
		}
		return errors.Join(errRuntimeInputMutationNotAttempted, err)
	}

	return nil
}

func (d *KimiACPDriver) abortPromptStart(id int, gate chan struct{}) {
	d.mu.Lock()
	delete(d.pending, id)
	closeGate := false
	if d.promptRequest == id {
		d.promptInFlight = false
		d.promptRequest = 0
		d.promptGate = nil
		d.promptAccepted = false
		closeGate = true
	}
	d.mu.Unlock()
	if closeGate {
		close(gate)
	}
}

func (d *KimiACPDriver) finishPromptResponse(id int, response kimiACPEnvelope) {
	stopReason := consts.ParsedEventStopReasonEndTurn
	message := ""
	if response.Error != nil {
		message = response.Error.Error()
	} else {
		stopReason = kimiACPStopReason(response.Result)
	}

	// Serialize the response terminal boundary against closeTransport so a
	// valid final response wins before a following stdout EOF is observed.
	d.writeMu.Lock()
	defer d.writeMu.Unlock()
	d.mu.Lock()
	if d.promptRequest != id {
		d.mu.Unlock()
		return
	}
	sessionID := d.sessionID
	delete(d.pending, id)
	d.promptInFlight = false
	d.promptRequest = 0
	d.promptGate = nil
	d.promptAccepted = false
	d.mu.Unlock()

	if message != "" {
		d.emit(translatedEvent{
			Kind:      consts.ParsedEventKindError,
			Message:   message,
			SessionID: sessionID,
			Metadata:  map[string]any{"will_retry": false},
		})
	}
	d.emit(translatedEvent{Kind: consts.ParsedEventKindTurnEnd, SessionID: sessionID, StopReason: stopReason})
}

func kimiACPStopReason(result json.RawMessage) string {
	var decoded struct {
		StopReason string `json:"stopReason"`
	}
	if json.Unmarshal(result, &decoded) != nil {
		return consts.ParsedEventStopReasonEndTurn
	}
	switch strings.TrimSpace(decoded.StopReason) {
	case "max_tokens":
		return consts.ParsedEventStopReasonMaxTokens
	default:
		return consts.ParsedEventStopReasonEndTurn
	}
}

func (d *KimiACPDriver) validateSession(sessionID string) (string, error) {
	current := d.currentSessionID()
	sessionID = strings.TrimSpace(sessionID)
	if current == "" {
		return "", fmt.Errorf("kimi: session is not initialized")
	}
	if sessionID == "" {
		sessionID = current
	}
	if sessionID != current {
		return "", fmt.Errorf("kimi: session mismatch: got %q, want %q", sessionID, current)
	}
	return current, nil
}

func (d *KimiACPDriver) InterruptActiveRun(ctx context.Context, sessionID string) error {
	current, err := d.validateSession(sessionID)
	if err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := d.writeNotification("session/cancel", map[string]any{"sessionId": current}); err != nil {
		return fmt.Errorf("kimi: cancel ACP turn: %w", err)
	}
	return nil
}