package protocol

import (
	"encoding/json"
	"fmt"
)

// Recognized Anthropic thinking blocks are omitted only from converted output.
func anthropicThinking(x object) (bool, error) {
	var names []string
	switch str(x["type"]) {
	case "thinking":
		names = []string{"thinking", "signature"}
	case "redacted_thinking":
		names = []string{"data"}
	default:
		return false, nil
	}
	if err := fields(x, append([]string{"type"}, names...)...); err != nil {
		return true, err
	}
	for _, name := range names {
		if _, ok := x[name].(string); !ok {
			return true, fmt.Errorf("invalid thinking %s", name)
		}
	}
	return true, nil
}

func anthropicOutput(v any) (any, error) {
	content, ok := v.([]any)
	if !ok {
		return nil, fmt.Errorf("invalid Anthropic content")
	}
	clean := []any{}
	for _, block := range content {
		thinking, err := anthropicThinking(obj(block))
		if err != nil {
			return nil, err
		}
		if !thinking {
			clean = append(clean, block)
		}
	}
	return clean, nil
}

type thinkingBlock struct {
	kind   string
	closed bool
}

func (s *streamState) thinkingEvent(typ string, m object) (bool, error) {
	if typ == "message_delta" || typ == "message_stop" {
		for _, b := range s.thinking {
			if !b.closed {
				return true, fmt.Errorf("unfinished thinking block")
			}
		}
		return false, nil
	}
	if typ != "content_block_start" && typ != "content_block_delta" && typ != "content_block_stop" {
		return false, nil
	}
	index, err := responseIndex(m, "index")
	if err != nil {
		return true, err
	}
	b := s.thinking[index]
	if typ == "content_block_start" {
		if b != nil || s.source[index] != nil {
			return true, fmt.Errorf("duplicate content index")
		}
		x := obj(m["content_block"])
		thinking, err := anthropicThinking(x)
		if err != nil || !thinking {
			return thinking, err
		}
		if !s.started || s.finished {
			return true, fmt.Errorf("thinking outside active message")
		}
		if s.thinking == nil {
			s.thinking = map[int]*thinkingBlock{}
		}
		s.thinking[index] = &thinkingBlock{kind: str(x["type"])}
		return true, nil
	}
	if b == nil {
		return false, nil
	}
	if b.closed || s.finished {
		return true, fmt.Errorf("event after thinking block closed")
	}
	if typ == "content_block_stop" {
		b.closed = true
		return true, nil
	}
	d := obj(m["delta"])
	field := ""
	switch str(d["type"]) {
	case "thinking_delta":
		field = "thinking"
	case "signature_delta":
		field = "signature"
	}
	if b.kind != "thinking" || field == "" {
		return true, fmt.Errorf("invalid thinking delta")
	}
	if err := fields(d, "type", field); err != nil {
		return true, err
	}
	if _, ok := d[field].(string); !ok {
		return true, fmt.Errorf("invalid thinking delta value")
	}
	return true, nil
}

// chatOutput omits output-only reasoning without exposing it as visible text.
func chatOutput(x object) (object, error) {
	if v := x["reasoning_content"]; v != nil {
		if _, ok := v.(string); !ok {
			return nil, fmt.Errorf("reasoning_content must be a string or null")
		}
	}
	m := object{}
	for k, v := range x {
		if k != "reasoning_content" {
			m[k] = v
		}
	}
	return m, nil
}

// Some Chat providers include streaming-style indices in complete tool calls.
func chatResponseMessage(x object) (object, error) {
	m, err := chatOutput(x)
	if err != nil {
		return nil, err
	}
	if calls, ok := x["tool_calls"].([]any); ok {
		clean := make([]any, len(calls))
		for i, v := range calls {
			call := obj(v)
			if _, exists := call["index"]; exists {
				index, err := responseIndex(call, "index")
				if err != nil || index != i {
					return nil, fmt.Errorf("invalid nonstream tool index")
				}
			}
			c := object{}
			for k, v := range call {
				if k != "index" {
					c[k] = v
				}
			}
			clean[i] = c
		}
		m["tool_calls"] = clean
	}
	return m, nil
}

// Phase describes output metadata, not message content or request semantics.
func responseMessage(x object) object {
	m := object{}
	for k, v := range x {
		if k != "phase" {
			m[k] = v
		}
	}
	if content, ok := x["content"].([]any); ok {
		clean := make([]any, len(content))
		for i, v := range content {
			p := obj(v)
			if logs, ok := p["logprobs"].([]any); ok && len(logs) == 0 && str(p["type"]) == "output_text" {
				q := object{}
				for k, v := range p {
					if k != "logprobs" {
						q[k] = v
					}
				}
				clean[i] = q
			} else {
				clean[i] = v
			}
		}
		m["content"] = clean
	}
	return m
}

func validateReasoning(x object) error {
	if err := fields(x, "type", "id", "status", "content", "summary", "encrypted_content"); err != nil {
		return err
	}
	for field, typ := range map[string]string{"content": "reasoning_text", "summary": "summary_text"} {
		if v, exists := x[field]; exists && v != nil {
			if _, ok := v.([]any); !ok {
				return fmt.Errorf("reasoning %s must be an array", field)
			}
		}
		for _, v := range arr(x[field]) {
			p := obj(v)
			if str(p["type"]) != typ {
				return fmt.Errorf("unsupported reasoning %s content", field)
			}
			if err := fields(p, "type", "text"); err != nil {
				return err
			}
		}
	}
	return nil
}

type reasoningItem struct {
	id     string
	closed bool
	parts  map[string]bool
}

func responseIndex(m object, field string) (int, error) {
	n, ok := m[field].(json.Number)
	i, err := n.Int64()
	if !ok || err != nil || i < 0 || int64(int(i)) != i {
		return 0, fmt.Errorf("invalid %s", field)
	}
	return int(i), nil
}

func (s *streamState) reasoningEvent(typ string, m object) (bool, error) {
	if typ == "response.completed" || typ == "response.incomplete" {
		seen := 0
		for i, v := range arr(obj(m["response"])["output"]) {
			x := obj(v)
			r := s.reasoning[i]
			if str(x["type"]) == "reasoning" {
				if r == nil || !r.closed || r.id != str(x["id"]) {
					return true, fmt.Errorf("final reasoning item disagrees with stream")
				}
				seen++
			} else if r != nil {
				return true, fmt.Errorf("final reasoning index changed type")
			}
		}
		if seen != len(s.reasoning) {
			return true, fmt.Errorf("final response missing reasoning item")
		}
		return false, nil
	}
	known := false
	switch typ {
	case "response.reasoning_text.delta", "response.reasoning_text.done",
		"response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
		"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
		known = true
	case "response.output_item.added", "response.output_item.done",
		"response.content_part.added", "response.content_part.done",
		"response.output_text.delta", "response.output_text.done",
		"response.function_call_arguments.delta", "response.function_call_arguments.done":
	default:
		return false, nil
	}
	key, err := responseIndex(m, "output_index")
	if err != nil {
		return true, err
	}
	r := s.reasoning[key]
	x := obj(m["item"])
	if typ == "response.output_item.added" {
		if r != nil || s.source[key] != nil {
			return true, fmt.Errorf("duplicate output index %d", key)
		}
		if str(x["type"]) != "reasoning" {
			return false, nil
		}
		if err := validateReasoning(x); err != nil {
			return true, err
		}
		if str(x["id"]) == "" {
			return true, fmt.Errorf("missing reasoning item id")
		}
		if s.reasoning == nil {
			s.reasoning = map[int]*reasoningItem{}
		}
		s.reasoning[key] = &reasoningItem{id: str(x["id"]), parts: map[string]bool{}}
		return true, nil
	}
	if r == nil {
		if known || str(x["type"]) == "reasoning" {
			return true, fmt.Errorf("reasoning event outside reasoning item")
		}
		return false, nil
	}
	if r.closed {
		return true, fmt.Errorf("event after reasoning item closed")
	}
	if typ == "response.output_item.done" {
		if str(x["type"]) != "reasoning" || str(x["id"]) != r.id {
			return true, fmt.Errorf("reasoning item identity mismatch")
		}
		if err := validateReasoning(x); err != nil {
			return true, err
		}
		for _, closed := range r.parts {
			if !closed {
				return true, fmt.Errorf("reasoning item has open parts")
			}
		}
		r.closed = true
		return true, nil
	}
	if str(m["item_id"]) != r.id {
		return true, fmt.Errorf("reasoning event item identity mismatch")
	}
	field, partType := "content_index", "reasoning_text"
	switch typ {
	case "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
		"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
		field, partType = "summary_index", "summary_text"
	case "response.reasoning_text.delta", "response.reasoning_text.done",
		"response.content_part.added", "response.content_part.done":
	default:
		return true, fmt.Errorf("non-reasoning event on reasoning item")
	}
	i, err := responseIndex(m, field)
	if err != nil {
		return true, err
	}
	partKey := fmt.Sprintf("%s:%d", field, i)
	closed, exists := r.parts[partKey]
	switch typ {
	case "response.content_part.added", "response.reasoning_summary_part.added",
		"response.content_part.done", "response.reasoning_summary_part.done":
		p := obj(m["part"])
		if str(p["type"]) != partType {
			return true, fmt.Errorf("unsupported reasoning part")
		}
		if err := fields(p, "type", "text"); err != nil {
			return true, err
		}
		added := typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added"
		if (added && exists) || (!added && (!exists || closed)) {
			return true, fmt.Errorf("invalid reasoning part lifecycle")
		}
		r.parts[partKey] = !added
	default:
		if !exists || closed {
			return true, fmt.Errorf("reasoning text outside active part")
		}
	}
	return true, nil
}