package daemon

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
)

// wireChatEvent decodes tool arguments and optional display traces separately.
type wireChatEvent struct {
	CallID  string          `json:"callId"`
	Args    json.RawMessage `json:"args"`
	Elapsed json.RawMessage `json:"elapsedMs"`
	Trace   json.RawMessage `json:"trace"`
	StreamEvent
}

func decodeChatEvent(data json.RawMessage) (*wireChatEvent, error) {
	var header struct {
		Type EventType `json:"type"`
	}
	if err := json.Unmarshal(data, &header); err != nil {
		return nil, fmt.Errorf("invalid event: %w", err)
	}
	if header.Type == "" {
		return nil, errors.New("invalid event: missing or empty type")
	}
	switch header.Type {
	case EventText, EventReset, EventRetry, EventError, EventInterrupted, EventMessage, EventUser, EventThinking, EventNote, EventCompacted, EventToolProgress, EventTool, EventUsage, EventCommitted, EventTurnMembership, EventTurnCompleted, "arguments_delta", "turn_started":
	default:
		return nil, nil
	}
	var event wireChatEvent
	if err := json.Unmarshal(data, &event); err != nil {
		return nil, fmt.Errorf("invalid %s event: %w", header.Type, err)
	}
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(data, &fields); err != nil {
		return nil, err
	}
	require := func(names ...string) error {
		for _, name := range names {
			value, exists := fields[name]
			if !exists || string(value) == "null" {
				return fmt.Errorf("invalid %s event: missing %s", header.Type, name)
			}
		}
		return nil
	}
	var err error
	switch header.Type {
	case EventText, EventMessage, EventUser, EventThinking, EventNote, EventError:
		err = require("text")
		if header.Type == EventMessage {
			err = require("text", "role")
			if event.Role != "assistant" {
				return nil, errors.New("invalid message event: role must be assistant")
			}
		}
		if header.Type == EventUser {
			err = require("text", "source", "triggeredAt")
		}
	case "arguments_delta":
		err = require("callId", "name", "text")
		if event.CallID == "" || event.ToolName == "" {
			return nil, errors.New("invalid argument delta identity")
		}
	case EventTurnMembership:
		err = require("turnId", "submissionIds")
		if event.TurnID == "" || event.SubmissionIDs == nil {
			return nil, errors.New("invalid turn membership identity")
		}
		for _, id := range event.SubmissionIDs {
			if id == "" {
				return nil, errors.New("invalid empty submission identity")
			}
		}
	case EventTurnCompleted:
		err = require("turnId")
		if event.TurnID == "" {
			return nil, errors.New("invalid empty turn identity")
		}
	case EventReset:
		if raw, ok := fields["before"]; ok && string(raw) == "null" {
			return nil, errors.New("invalid reset before")
		}
		if raw, ok := fields["more"]; ok && string(raw) == "null" {
			return nil, errors.New("invalid reset more")
		}
	case EventCommitted:
		err = require("seq")
		if event.Seq <= 0 {
			err = errors.New("invalid committed event: seq must be positive")
		}
	case EventCompacted:
		err = require("evicted", "summary")
		if event.Evicted < 0 || event.Evicted > 1000000 || len(event.Summary) > 60000 {
			err = errors.New("invalid compacted event")
		}
	case EventTool:
		err = require("callId", "name", "args", "result")
		if event.CallID == "" || event.ToolName == "" {
			return nil, errors.New("invalid tool identity")
		}
		if len(event.Args) != 0 && string(event.Args) != "null" {
			var encoded string
			if decodeErr := json.Unmarshal(event.Args, &encoded); decodeErr != nil {
				return nil, fmt.Errorf("invalid tool args: %w", decodeErr)
			}
			// Tool arguments are model-generated text, not protocol-owned JSON.
			if json.Unmarshal([]byte(encoded), &event.ToolArgs) != nil {
				event.ToolArgs = map[string]any{}
			}
		}
	case EventUsage:
		var usage Usage
		if decodeErr := json.Unmarshal(data, &usage); decodeErr != nil {
			return nil, fmt.Errorf("invalid usage event: %w", decodeErr)
		}
		if fieldErr := require("model", "recordedAt"); fieldErr != nil {
			return nil, fieldErr
		}
		if usage.RecordedAt == nil || *usage.RecordedAt < 0 {
			return nil, errors.New("invalid usage timestamp")
		}
		for _, count := range []*int{usage.PromptTokens, usage.CachedPromptTokens, usage.CacheWriteTokens, usage.CompletionTokens, usage.TotalTokens} {
			if count != nil && *count < 0 {
				return nil, errors.New("invalid usage token count")
			}
		}
		if usage.ElapsedMs != nil && *usage.ElapsedMs < 0 || usage.TokensPerSecond != nil && *usage.TokensPerSecond < 0 {
			return nil, errors.New("invalid usage rate or duration")
		}
		for _, step := range usage.CacheFade {
			if step.At < 0 || step.Cached != nil && *step.Cached < 0 {
				return nil, errors.New("invalid cache fade")
			}
		}
		event.Usage = &usage
	case EventToolProgress:
		if fieldErr := validateToolProgress(fields["progress"]); fieldErr != nil {
			return nil, fieldErr
		}
		if progress := event.Progress; progress != nil {
			if progress.CallID == "" || progress.Name == "" || len(progress.CallID) > 200 || len(progress.Name) > 100 {
				return nil, errors.New("invalid tool progress event")
			}
			progress.Name = cleanLabel(progress.Name)
			if intent := progress.Intent; intent != nil {
				if (intent.Kind != "write" && intent.Kind != "edit" && intent.Kind != "read" && intent.Kind != "run") || len(intent.Target) > 300 {
					return nil, errors.New("invalid tool progress intent")
				}
			}
			if code := progress.Code; code != nil {
				if code.Offset < 0 || (progress.Phase == "generating" && len(code.Text) > 512) || len(code.Text) > 16000 {
					return nil, errors.New("invalid tool progress code")
				} else if progress.Phase == "running" {
					if progress.Intent == nil && progress.Name == "python" {
						progress.Intent = ParsePythonIntent(code.Text)
					}
					head, _, _ := strings.Cut(strings.TrimSpace(code.Text), "\n")
					progress.Code = &ToolCodePreview{Text: sanitizeControlRunes(head)}
				}
			}
		}
	}
	if err != nil {
		return nil, err
	}
	if event.Type != EventUsage && len(event.Elapsed) != 0 {
		if err := json.Unmarshal(event.Elapsed, &event.ElapsedMs); err != nil {
			return nil, fmt.Errorf("invalid elapsedMs: %w", err)
		}
	}
	if event.Timestamp != nil && *event.Timestamp < 0 || event.ElapsedMs < 0 || event.Before < 0 {
		return nil, errors.New("invalid event timestamp, elapsed time, or history cursor")
	}
	if event.Image != nil {
		event.Image = ParseImageMetadata(event.Image)
		if event.Image == nil {
			return nil, errors.New("invalid image metadata")
		}
	}
	if len(event.Trace) != 0 && string(event.Trace) != "null" {
		// Display-only evidence must not prevent loading the saved conversation.
		event.ToolTrace = ParseToolTrace(event.Trace)
	}
	return &event, nil
}

// ResetStream discards unfinished tool arguments and starts the next stream from
// a fresh snapshot. Call it after the subscription has stopped writing.
func (c *ChatClient) ResetStream() {
	c.mu.Lock()
	c.argumentsByCall = make(map[string]*strings.Builder)
	c.afterSeq = -1
	c.afterGeneration = ""
	c.mu.Unlock()
}

type streamBatch struct {
	Generation string            `json:"generation"`
	Cursor     *int64            `json:"cursor"`
	Events     []json.RawMessage `json:"events"`
}

func streamFailure(kind StreamFailureKind, cause error) error {
	return &StreamError{Kind: kind, Cause: cause}
}

func classifyStreamFailure(err error) error {
	if err == nil {
		return nil
	}
	if _, ok := errors.AsType[*StreamError](err); ok {
		return err
	}
	if apiErr, ok := errors.AsType[*APIError](err); ok {
		kind := StreamTerminal
		if apiErr.StatusCode == http.StatusRequestTimeout || apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= 500 {
			kind = StreamTransient
		}
		return streamFailure(kind, err)
	}
	if errors.Is(err, bufio.ErrTooLong) {
		return streamFailure(StreamProtocol, err)
	}
	if _, ok := errors.AsType[net.Error](err); ok {
		return streamFailure(StreamTransient, err)
	}
	if isConnectionError(err) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return streamFailure(StreamTransient, err)
	}
	return streamFailure(StreamTerminal, err)
}

// Stream uses ctx for cancellation; onEvent must also honor ctx if it blocks.
func (c *ChatClient) Stream(ctx context.Context, tail int, onEvent func(StreamEvent) error) error {
	return c.StreamWithProgress(ctx, tail, onEvent, nil)
}

// StreamWithProgress reports fully consumed batches, including ignored additive
// events, so subscribers can reset their reconnect backoff after progress.
func (c *ChatClient) StreamWithProgress(ctx context.Context, tail int, onEvent func(StreamEvent) error, onBatchConsumed func()) error {
	defer func() {
		// Cancellation also clears state when the request fails before reading.
		// Transient failures and EOF retain the cursor and unfinished arguments.
		if ctx.Err() != nil {
			c.ResetStream()
		}
	}()

	c.mu.Lock()
	afterSeq, afterGeneration := c.afterSeq, c.afterGeneration
	c.mu.Unlock()

	query := url.Values{}
	if afterGeneration != "" {
		query.Set("after_generation", afterGeneration)
		query.Set("after_seq", fmt.Sprint(afterSeq))
	}
	if tail > 0 {
		query.Set("tail", fmt.Sprint(tail))
	}

	route := "/stream?" + query.Encode()
	operation := operation{Name: "stream session", Method: http.MethodGet, Path: sessionPath(c.agentID, route), Policy: readRecovery}
	err := scanEventStream(ctx, c.conn, operation, streamLimits{requireSSE: true, lineBytes: 10 * 1024 * 1024, errorBytes: 64 * 1024}, func(scanner *bufio.Scanner) error {
		return c.readStream(ctx, scanner, onBatchConsumed, func(event StreamEvent) error {
			if err := onEvent(event); err != nil {
				return streamFailure(StreamTerminal, err)
			}
			return nil
		})
	})
	return classifyStreamFailure(err)
}

func (c *ChatClient) readStream(ctx context.Context, scanner *bufio.Scanner, onBatchConsumed func(), onEvent func(StreamEvent) error) error {
	reporter := NewToolProgressReporter(func(progress *ToolProgress) error {
		return onEvent(StreamEvent{Type: EventToolProgress, Progress: progress})
	})
	var eventType string
	for scanner.Scan() {
		if err := ctx.Err(); err != nil {
			return err
		}
		line := scanner.Text()
		if line == "" {
			eventType = ""
			continue
		}
		if strings.HasPrefix(line, "event:") {
			eventType = strings.TrimSpace(line[6:])
			continue
		}
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(line[5:])
		if payload == "" {
			continue
		}
		if eventType == "error" {
			var failure struct {
				Error string `json:"error"`
				Text  string `json:"text"`
			}
			if err := json.Unmarshal([]byte(payload), &failure); err != nil {
				return streamFailure(StreamProtocol, fmt.Errorf("invalid stream error: %w", err))
			}
			if failure.Error != "" {
				return streamFailure(StreamTerminal, errors.New(failure.Error))
			}
			if failure.Text != "" {
				return streamFailure(StreamTerminal, errors.New(failure.Text))
			}
			return streamFailure(StreamProtocol, errors.New("invalid stream error: missing error or text"))
		}
		var batch streamBatch
		if err := json.Unmarshal([]byte(payload), &batch); err != nil {
			return streamFailure(StreamProtocol, fmt.Errorf("invalid stream batch: %w", err))
		}
		if batch.Generation == "" || batch.Cursor == nil || *batch.Cursor < 0 || batch.Events == nil {
			return streamFailure(StreamProtocol, errors.New("invalid stream batch: missing generation, cursor, or events"))
		}
		// Validate the whole batch before delivering events or changing preview state.
		events := make([]wireChatEvent, 0, len(batch.Events))
		for index, raw := range batch.Events {
			event, err := decodeChatEvent(raw)
			if err != nil {
				return streamFailure(StreamProtocol, err)
			}
			if event != nil {
				if event.Type == EventReset && index != 0 {
					return streamFailure(StreamProtocol, errors.New("reset must begin its batch"))
				}
				events = append(events, *event)
			}
		}
		reset := len(events) > 0 && events[0].Type == EventReset
		c.mu.Lock()
		previous, generation := c.afterSeq, c.afterGeneration
		c.mu.Unlock()
		if generation != batch.Generation && !reset {
			return streamFailure(StreamProtocol, errors.New("initial or changed generation stream must begin with reset"))
		}
		if generation == batch.Generation && *batch.Cursor < previous && !reset {
			return streamFailure(StreamProtocol, errors.New("stream cursor regressed without reset"))
		}
		if err := c.validateArgumentBatch(events); err != nil {
			return streamFailure(StreamProtocol, err)
		}
		for _, event := range events {
			switch event.Type {
			case EventReset, EventRetry, "turn_started", EventMessage, EventInterrupted, EventError:
				c.mu.Lock()
				c.argumentsByCall = make(map[string]*strings.Builder)
				c.mu.Unlock()
				if err := reporter.Report(nil, "running"); err != nil {
					return err
				}
			}
			if event.Type == "turn_started" {
				continue
			}
			if event.Type == EventTool || (event.Type == EventToolProgress && event.Progress != nil && event.Progress.Phase == "running") {
				callID := event.CallID
				if event.Progress != nil {
					callID = event.Progress.CallID
				}
				c.mu.Lock()
				delete(c.argumentsByCall, callID)
				c.mu.Unlock()
				if err := reporter.Report(nil, "running"); err != nil {
					return err
				}
			}
			if event.Type == "arguments_delta" {
				if event.Text != "" {
					arguments := c.appendArguments(event.CallID, event.Text)
					call := &ToolCallAssembly{ID: event.CallID}
					call.Function.Name = event.ToolName
					call.Function.Arguments = arguments
					if err := reporter.Report(call, "generating"); err != nil {
						return err
					}
				}
				continue
			}
			event.Replayed = reset && event.Type != EventReset
			if err := onEvent(event.StreamEvent); err != nil {
				return err
			}
		}
		c.mu.Lock()
		c.afterSeq = *batch.Cursor
		c.afterGeneration = batch.Generation
		c.mu.Unlock()
		if onBatchConsumed != nil {
			onBatchConsumed()
		}
	}
	return nil
}

func (c *ChatClient) appendArguments(callID, text string) string {
	c.mu.Lock()
	defer c.mu.Unlock()
	arguments, exists := c.argumentsByCall[callID]
	if !exists {
		arguments = new(strings.Builder)
		c.argumentsByCall[callID] = arguments
	}
	arguments.WriteString(text)
	return arguments.String()
}

// validateArgumentBatch checks preview limits before delivering any batch event.
func (c *ChatClient) validateArgumentBatch(events []wireChatEvent) error {
	sizes := make(map[string]int)
	c.mu.Lock()
	for id, args := range c.argumentsByCall {
		sizes[id] = args.Len()
	}
	c.mu.Unlock()
	for _, event := range events {
		switch event.Type {
		case EventReset, EventRetry, "turn_started", EventMessage, EventInterrupted, EventError:
			clear(sizes)
		case EventTool:
			delete(sizes, event.CallID)
		case EventToolProgress:
			if event.Progress != nil && event.Progress.Phase == "running" {
				delete(sizes, event.Progress.CallID)
			}
		case "arguments_delta":
			if event.Text == "" {
				continue
			}
			sizes[event.CallID] += len(event.Text)
			retained := 0
			for id, size := range sizes {
				retained += len(id) + size
			}
			if len(sizes) > 32 || retained > 2000000 {
				return errors.New("too many tool argument previews for this client to display")
			}
		}
	}
	return nil
}