package llm

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

// StreamEventKind is what a streaming callback is being told about.
type StreamEventKind string

const (
	StreamText      StreamEventKind = "text"
	StreamThinking  StreamEventKind = "thinking"
	StreamToolStart StreamEventKind = "tool_start"
	StreamToolInput StreamEventKind = "tool_input"
	StreamBlockStop StreamEventKind = "block_stop"
)

// StreamEvent is one incremental update.
type StreamEvent struct {
	Kind     StreamEventKind
	Index    int
	Text     string // the delta, not the accumulated value
	ToolID   string
	ToolName string
}

// StreamMessage sends a streaming request, calls onEvent for each delta, and
// returns the assembled response. The callback may be nil.
func (c *Client) StreamMessage(ctx context.Context, req Request, onEvent func(StreamEvent)) (*Response, error) {
	if !c.HasCredentials() {
		return nil, errNoCredentials
	}
	a, err := c.resolveAuth(ctx)
	if err != nil {
		return nil, err
	}
	if req.Model == "" {
		req.Model = c.Model()
	}
	if req.MaxTokens == 0 {
		req.MaxTokens = c.MaxTokens
	}
	req.Stream = true

	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("llm: encoding request: %w", err)
	}

	var lastErr error
	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
		if attempt > 0 {
			if err := sleep(ctx, c.backoff(attempt, lastErr)); err != nil {
				return nil, err
			}
		}
		resp, err := c.streamAttempt(ctx, body, a, onEvent)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		var apiErr *APIError
		if ok := asAPIError(err, &apiErr); ok && !apiErr.Retryable() {
			return nil, err
		}
		// A stream that broke partway is not safely retryable either: the
		// caller may already have seen deltas.
		if _, isAPI := err.(*APIError); !isAPI {
			return nil, err
		}
	}
	return nil, lastErr
}

func (c *Client) streamAttempt(ctx context.Context, body []byte, a auth, onEvent func(StreamEvent)) (*Response, error) {
	httpReq, err := c.newRequest(ctx, http.MethodPost, "/v1/messages", body, a)
	if err != nil {
		return nil, err
	}
	// The only thing that distinguishes a streaming request from a buffered
	// one on the wire. Everything else, including the credential, comes from
	// newRequest so the two routes cannot drift apart.
	httpReq.Header.Set("accept", "text/event-stream")

	httpResp, err := c.HTTP.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("llm: %w", err)
	}
	defer httpResp.Body.Close()

	if httpResp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(httpResp.Body)
		return nil, parseAPIError(httpResp.StatusCode, httpResp.Header.Get("retry-after"), raw)
	}
	return readStream(httpResp.Body, onEvent)
}

// blockBuilder accumulates one content block across its deltas. Tool input
// arrives as a series of JSON fragments that are only valid once concatenated,
// which is the fiddly part of streaming.
type blockBuilder struct {
	block ContentBlock
	json  strings.Builder
}

func readStream(r io.Reader, onEvent func(StreamEvent)) (*Response, error) {
	emit := func(ev StreamEvent) {
		if onEvent != nil {
			onEvent(ev)
		}
	}

	resp := &Response{Role: RoleAssistant}
	blocks := map[int]*blockBuilder{}
	var order []int

	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data:") {
			continue // "event:" lines are redundant; the payload names its type
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "" || payload == "[DONE]" {
			continue
		}

		var ev struct {
			Type    string `json:"type"`
			Index   int    `json:"index"`
			Message *struct {
				ID    string `json:"id"`
				Model string `json:"model"`
				Usage Usage  `json:"usage"`
			} `json:"message"`
			ContentBlock *ContentBlock `json:"content_block"`
			Delta        *struct {
				Type        string       `json:"type"`
				Text        string       `json:"text"`
				PartialJSON string       `json:"partial_json"`
				Thinking    string       `json:"thinking"`
				Signature   string       `json:"signature"`
				StopReason  string       `json:"stop_reason"`
				StopDetails *StopDetails `json:"stop_details"`
			} `json:"delta"`
			Usage *Usage `json:"usage"`
			Error *struct {
				Type    string `json:"type"`
				Message string `json:"message"`
			} `json:"error"`
		}
		if err := json.Unmarshal([]byte(payload), &ev); err != nil {
			return nil, fmt.Errorf("llm: malformed stream event: %w", err)
		}

		switch ev.Type {
		case "error":
			if ev.Error != nil {
				return nil, &APIError{StatusCode: 200, Type: ev.Error.Type, Message: ev.Error.Message}
			}
			return nil, fmt.Errorf("llm: stream error")

		case "message_start":
			if ev.Message != nil {
				resp.ID, resp.Model, resp.Usage = ev.Message.ID, ev.Message.Model, ev.Message.Usage
			}

		case "content_block_start":
			if ev.ContentBlock == nil {
				continue
			}
			blocks[ev.Index] = &blockBuilder{block: *ev.ContentBlock}
			order = append(order, ev.Index)
			if ev.ContentBlock.Type == BlockToolUse {
				emit(StreamEvent{Kind: StreamToolStart, Index: ev.Index,
					ToolID: ev.ContentBlock.ID, ToolName: ev.ContentBlock.Name})
			}

		case "content_block_delta":
			b, ok := blocks[ev.Index]
			if !ok || ev.Delta == nil {
				continue
			}
			switch ev.Delta.Type {
			case "text_delta":
				b.block.Text += ev.Delta.Text
				emit(StreamEvent{Kind: StreamText, Index: ev.Index, Text: ev.Delta.Text})
			case "thinking_delta":
				b.block.Thinking += ev.Delta.Thinking
				emit(StreamEvent{Kind: StreamThinking, Index: ev.Index, Text: ev.Delta.Thinking})
			case "signature_delta":
				b.block.Signature += ev.Delta.Signature
			case "input_json_delta":
				// Fragments are not individually parseable; buffer and parse
				// the whole thing at content_block_stop.
				b.json.WriteString(ev.Delta.PartialJSON)
				emit(StreamEvent{Kind: StreamToolInput, Index: ev.Index, Text: ev.Delta.PartialJSON})
			}

		case "content_block_stop":
			b, ok := blocks[ev.Index]
			if !ok {
				continue
			}
			if s := b.json.String(); s != "" {
				if !json.Valid([]byte(s)) {
					return nil, fmt.Errorf("llm: tool %q streamed invalid JSON arguments: %s", b.block.Name, s)
				}
				b.block.Input = json.RawMessage(s)
			}
			emit(StreamEvent{Kind: StreamBlockStop, Index: ev.Index})

		case "message_delta":
			if ev.Delta != nil {
				resp.StopReason = ev.Delta.StopReason
				if ev.Delta.StopDetails != nil {
					resp.StopDetails = ev.Delta.StopDetails
				}
			}
			if ev.Usage != nil {
				// message_delta carries the final output count; the input
				// counts arrived with message_start.
				resp.Usage.OutputTokens = ev.Usage.OutputTokens
				if ev.Usage.InputTokens > 0 {
					resp.Usage.InputTokens = ev.Usage.InputTokens
				}
			}
		}
	}
	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("llm: reading stream: %w", err)
	}

	for _, i := range order {
		if b := blocks[i]; b != nil {
			// A tool_use with no input at all still needs a valid object.
			if b.block.Type == BlockToolUse && len(b.block.Input) == 0 {
				b.block.Input = json.RawMessage("{}")
			}
			resp.Content = append(resp.Content, b.block)
		}
	}
	if resp.StopReason == "" {
		return nil, fmt.Errorf("llm: stream ended without a stop_reason")
	}
	return resp, nil
}

// parseAPIError builds an APIError from a non-2xx body.
func parseAPIError(status int, retryAfterHeader string, raw []byte) *APIError {
	apiErr := &APIError{StatusCode: status, Body: string(raw)}
	var envelope struct {
		Error struct {
			Type    string `json:"type"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if json.Unmarshal(raw, &envelope) == nil {
		apiErr.Type = envelope.Error.Type
		apiErr.Message = envelope.Error.Message
	}
	if retryAfterHeader != "" {
		apiErr.Message += " (retry-after " + retryAfterHeader + "s)"
	}
	return apiErr
}