from ...agent.orchestrator import Orchestrator
from ...entities import (
    ReasoningToken,
    Token,
    TokenDetail,
    ToolCall,
    ToolCallDiagnostic,
    ToolCallError,
    ToolCallResult,
    ToolCallToken,
)
from ...event import Event, EventType
from ...server.entities import ResponsesRequest
from ...utils import (
    to_json,
    tool_call_diagnostic_payload,
    tool_call_error_payload,
)
from .. import di_get_logger, di_get_orchestrator
from ..sse import sse_headers, sse_message
from . import orchestrate, resolve_model_id

from enum import Enum, auto
from logging import Logger
from typing import Any, AsyncIterator, cast

from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse


class ResponseState(Enum):
    REASONING = auto()
    TOOL_CALLING = auto()
    ANSWERING = auto()


router = APIRouter(tags=["responses"])


@router.post("/responses", response_model=None)
async def create_response(
    request: ResponsesRequest,
    logger: Logger = Depends(di_get_logger),
    orchestrator: Orchestrator = Depends(di_get_orchestrator),
) -> dict[str, Any] | StreamingResponse:
    assert orchestrator and isinstance(orchestrator, Orchestrator)
    assert logger and isinstance(logger, Logger)
    assert request and request.messages
    model_id = resolve_model_id(orchestrator, request.model)

    response, response_id, timestamp = await orchestrate(
        request, logger, orchestrator
    )

    if request.stream:

        async def generate() -> AsyncIterator[str]:
            seq = 0

            yield sse_message(
                to_json(
                    {
                        "type": "response.created",
                        "response": {
                            "id": str(response_id),
                            "created_at": timestamp,
                            "model": model_id,
                            "type": "response",
                            "status": "in_progress",
                        },
                    }
                ),
                event="response.created",
            )

            state: ResponseState | None = None
            tool_call_id: str | None = None

            async for token in response:
                call_id: str | None = None
                if isinstance(token, Event):
                    if token.type not in (
                        EventType.TOOL_DIAGNOSTIC,
                        EventType.TOOL_PROCESS,
                        EventType.TOOL_RESULT,
                    ):
                        continue
                    item = _tool_call_event_item(token)
                    if item is None:
                        continue
                    item_id = item.get("id")
                    call_id = item_id if isinstance(item_id, str) else None
                elif (
                    isinstance(token, ToolCallToken) and token.call is not None
                ):
                    call_id = str(token.call.id)

                new_state = _new_state(token)
                events = _switch_state(state, new_state, tool_call_id, call_id)
                state = new_state
                if state is ResponseState.TOOL_CALLING:
                    if call_id is not None:
                        tool_call_id = call_id
                else:
                    tool_call_id = None
                for event in events:
                    yield event

                for ev in _token_to_sse(token, seq):
                    yield ev

                seq += 1

            events = _switch_state(state, None, tool_call_id, None)
            for event in events:
                yield event

            yield sse_message(
                to_json({"type": "response.completed"}),
                event="response.completed",
            )

            yield sse_message("{}", event="done")
            await orchestrator.sync_messages()

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers=sse_headers(),
        )

    text = await response.to_str()
    body = {
        "id": str(response_id),
        "created": timestamp,
        "model": model_id,
        "type": "response",
        "output": [{"content": [{"type": "output_text", "text": text}]}],
        "usage": {
            "input_text_tokens": response.input_token_count,
            "output_text_tokens": response.output_token_count,
            "total_tokens": (
                response.input_token_count + response.output_token_count
            ),
        },
    }
    await orchestrator.sync_messages()
    return body


def _token_to_sse(
    token: ReasoningToken | ToolCallToken | Token | TokenDetail | Event | str,
    seq: int,
) -> list[str]:
    events: list[str] = []

    if isinstance(token, ReasoningToken):
        events.append(
            sse_message(
                to_json(
                    {
                        "type": "response.reasoning_text.delta",
                        "delta": token.token,
                        "output_index": 0,
                        "content_index": 0,
                        "sequence_number": seq,
                    }
                ),
                event="response.reasoning_text.delta",
            )
        )
    elif isinstance(token, Event) and token.type in (
        EventType.TOOL_DIAGNOSTIC,
        EventType.TOOL_PROCESS,
        EventType.TOOL_RESULT,
    ):
        item = _tool_call_event_item(token)
        if item is None:
            return events
        delta_obj = {
            "id": item["id"],
            "name": item["name"],
            "arguments": item.get("arguments"),
        }
        if "diagnostic" in item:
            delta_obj["diagnostic"] = item["diagnostic"]
        elif token.type is EventType.TOOL_RESULT:
            if "error" in item:
                delta_obj["error"] = item.get("error")
            else:
                result = item.get("result", None)
                delta_obj["result"] = to_json(result) if result else None

        event_type = (
            "response.tool_call_diagnostic.delta"
            if "diagnostic" in item
            else "response.function_call_arguments.delta"
        )
        events.append(
            sse_message(
                to_json(
                    {
                        **delta_obj,
                        "type": event_type,
                        "delta": to_json(delta_obj),
                        "id": item["id"],
                        "output_index": 0,
                        "content_index": 0,
                        "sequence_number": seq,
                    }
                ),
                event=event_type,
            )
        )
    elif isinstance(token, ToolCallToken):
        if token.call is not None:
            delta_obj = {
                "id": str(token.call.id),
                "name": token.call.name,
                "arguments": token.call.arguments,
            }
            events.append(
                sse_message(
                    to_json(
                        {
                            "type": "response.function_call_arguments.delta",
                            "delta": to_json(delta_obj),
                            "id": str(token.call.id),
                            "output_index": 0,
                            "content_index": 0,
                            "sequence_number": seq,
                        }
                    ),
                    event="response.function_call_arguments.delta",
                )
            )
        else:
            events.append(
                sse_message(
                    to_json(
                        {
                            "type": "response.custom_tool_call_input.delta",
                            "delta": token.token,
                            "output_index": 0,
                            "content_index": 0,
                            "sequence_number": seq,
                        }
                    ),
                    event="response.custom_tool_call_input.delta",
                )
            )
    else:
        events.append(
            sse_message(
                to_json(
                    {
                        "type": "response.output_text.delta",
                        "delta": (
                            token.token
                            if isinstance(token, Token)
                            else str(token)
                        ),
                        "output_index": 0,
                        "content_index": 0,
                        "sequence_number": seq,
                    }
                ),
                event="response.output_text.delta",
            )
        )
    return events


def _switch_state(
    state: ResponseState | None,
    new_state: ResponseState | None,
    current_tool_call_id: str | None,
    new_tool_call_id: str | None,
) -> list[str]:
    events: list[str] = []
    changed = state is not new_state or (
        state is ResponseState.TOOL_CALLING
        and new_state is ResponseState.TOOL_CALLING
        and new_tool_call_id is not None
    )
    if changed:
        if state is ResponseState.REASONING:
            events.append(_reasoning_text_done())
            events.append(_content_part_done())
            events.append(_output_item_done())
        elif state is ResponseState.TOOL_CALLING:
            events.append(_custom_tool_call_input_done(current_tool_call_id))
            events.append(_content_part_done(current_tool_call_id))
            events.append(_output_item_done(current_tool_call_id))
        elif state is ResponseState.ANSWERING:
            events.append(_output_text_done())
            events.append(_content_part_done())
            events.append(_output_item_done())

        if new_state is ResponseState.REASONING:
            events.append(_output_item_added(new_state))
            events.append(_content_part_added("reasoning_text"))
        elif new_state is ResponseState.TOOL_CALLING:
            events.append(_output_item_added(new_state, new_tool_call_id))
            events.append(_content_part_added("input_text", new_tool_call_id))
        elif new_state is ResponseState.ANSWERING:
            events.append(_output_item_added(new_state))
            events.append(_content_part_added("output_text"))

    return events


def _new_state(
    token: (
        ReasoningToken
        | ToolCallToken
        | Token
        | TokenDetail
        | Event
        | str
        | None
    ),
) -> ResponseState | None:
    if isinstance(token, ReasoningToken):
        new_state = ResponseState.REASONING
    elif isinstance(token, (ToolCallToken, Event)) and (
        not isinstance(token, Event)
        or token.type
        in (
            EventType.TOOL_DIAGNOSTIC,
            EventType.TOOL_PROCESS,
            EventType.TOOL_RESULT,
        )
    ):
        new_state = ResponseState.TOOL_CALLING
    elif token is not None:
        new_state = ResponseState.ANSWERING
    else:
        new_state = None
    return new_state


def _output_item_added(state: ResponseState, id: str | None = None) -> str:
    item_types = {
        ResponseState.REASONING: "reasoning_text",
        ResponseState.TOOL_CALLING: "custom_tool_call_input",
        ResponseState.ANSWERING: "output_text",
    }
    item = {"type": item_types[state]}
    if id is not None:
        item["id"] = id
    return sse_message(
        to_json(
            {
                "type": "response.output_item.added",
                "output_index": 0,
                "item": item,
            }
        ),
        event="response.output_item.added",
    )


def _output_item_done(id: str | None = None) -> str:
    data = {"type": "response.output_item.done", "output_index": 0}
    if id is not None:
        data["item"] = {"id": id}
    return sse_message(to_json(data), event="response.output_item.done")


def _reasoning_text_done() -> str:
    return sse_message(
        to_json(
            {
                "type": "response.reasoning_text.done",
                "output_index": 0,
                "content_index": 0,
            }
        ),
        event="response.reasoning_text.done",
    )


def _custom_tool_call_input_done(id: str | None = None) -> str:
    data = {
        "type": "response.custom_tool_call_input.done",
        "output_index": 0,
        "content_index": 0,
    }
    if id is not None:
        data["id"] = id
    return sse_message(
        to_json(data),
        event="response.custom_tool_call_input.done",
    )


def _output_text_done() -> str:
    return sse_message(
        to_json(
            {
                "type": "response.output_text.done",
                "output_index": 0,
                "content_index": 0,
            }
        ),
        event="response.output_text.done",
    )


def _content_part_added(part_type: str, id: str | None = None) -> str:
    part = {"type": part_type}
    if id is not None:
        part["id"] = id
    return sse_message(
        to_json(
            {
                "type": "response.content_part.added",
                "output_index": 0,
                "content_index": 0,
                "part": part,
            }
        ),
        event="response.content_part.added",
    )


def _content_part_done(id: str | None = None) -> str:
    data = {
        "type": "response.content_part.done",
        "output_index": 0,
        "content_index": 0,
    }
    if id is not None:
        data["part"] = {"id": id}
    return sse_message(to_json(data), event="response.content_part.done")


def _tool_call_event_item(event: Event) -> dict[str, Any] | None:
    payload = cast(Any, event.payload)
    tool_result = (
        payload["result"]
        if event.type == EventType.TOOL_RESULT
        and isinstance(payload, dict)
        and "result" in payload
        else None
    )
    if isinstance(tool_result, ToolCallDiagnostic):
        call = payload.get("call") if isinstance(payload, dict) else None
        return _tool_call_diagnostic_item(
            tool_result,
            call if isinstance(call, ToolCall) else None,
        )
    if event.type is EventType.TOOL_DIAGNOSTIC:
        diagnostic = (
            payload.get("diagnostic") if isinstance(payload, dict) else None
        )
        if not isinstance(diagnostic, ToolCallDiagnostic):
            diagnostics = (
                payload.get("diagnostics")
                if isinstance(payload, dict)
                else None
            )
            if isinstance(diagnostics, list):
                diagnostic = next(
                    (
                        item
                        for item in diagnostics
                        if isinstance(item, ToolCallDiagnostic)
                    ),
                    None,
                )
        if not isinstance(diagnostic, ToolCallDiagnostic):
            return None
        call = payload.get("call") if isinstance(payload, dict) else None
        return _tool_call_diagnostic_item(
            diagnostic,
            call if isinstance(call, ToolCall) else None,
        )
    if tool_result is not None:
        tool_call = tool_result.call
    elif isinstance(payload, list):
        tool_call = payload[0]
    elif isinstance(payload, dict):
        tool_call = payload.get("call") or payload.get(0)
    else:
        tool_call = None
    if tool_call is None:
        return None
    item = {
        "type": "function_call",
        "id": str(tool_call.id),
        "name": tool_call.name,
        "arguments": tool_call.arguments,
    }
    if tool_result is not None:
        if isinstance(tool_result, ToolCallError):
            item["error"] = tool_call_error_payload(tool_result)
        elif isinstance(tool_result, ToolCallResult):
            item["result"] = tool_result.result
        else:
            item["result"] = tool_result
    return item


def _tool_call_diagnostic_item(
    diagnostic: ToolCallDiagnostic, call: ToolCall | None
) -> dict[str, Any]:
    diagnostic_payload = {
        "id": str(diagnostic.id),
        **tool_call_diagnostic_payload(diagnostic),
    }
    if diagnostic.call_id is not None:
        diagnostic_payload["call_id"] = str(diagnostic.call_id)
    if diagnostic.started_at is not None:
        diagnostic_payload["started_at"] = diagnostic.started_at.isoformat()
    if diagnostic.finished_at is not None:
        diagnostic_payload["finished_at"] = diagnostic.finished_at.isoformat()
    if diagnostic.duration_ms is not None:
        diagnostic_payload["duration_ms"] = diagnostic.duration_ms
    return {
        "type": "function_call",
        "id": str(call.id if call else diagnostic.call_id or diagnostic.id),
        "name": (
            call.name
            if call
            else diagnostic.canonical_name
            or diagnostic.requested_name
            or "tool"
        ),
        "arguments": call.arguments if call else None,
        "diagnostic": diagnostic_payload,
    }
