from __future__ import annotations

# standard
import json
import time
from collections.abc import AsyncIterator

# third party
import httpx

# custom
from ..base import BaseEngine
from ..types import (
    Message,
    Performance,
    Response,
    Role,
    StopReason,
    Tool,
    ToolCall,
    Usage,
    make_non_stream_performance,
    resolve_tokens,
)


class AnthropicEngine(BaseEngine):
    BASE_URL = "https://api.anthropic.com/v1"
    API_VERSION = "2023-06-01"

    def __init__(
        self,
        model: str,
        api_key: str,
        max_tokens: int = 8192,
        thinking_budget: int | None = None,
    ) -> None:
        self.provider = "anthropic"
        self.model = model
        self.api_key = api_key
        self.max_tokens = max_tokens
        # NOTE budget defaults to max_tokens - 1024 (leave room for the response).
        # must be >= 1024 and < max_tokens per Anthropic's requirements.
        self.thinking_budget = (
            thinking_budget
            if thinking_budget is not None
            else max(1024, max_tokens - 1024)
        )

    ### PAYLOAD BUILDERS

    def _split_system(
        self, messages: list[Message]
    ) -> tuple[str | None, list[Message]]:
        """Anthropic takes system prompt separately from the messages list.

        If there are no non-system messages (system-only conversation), the
        system content is promoted to a user message so the API receives at
        least one message in the messages array.
        """

        system = next((m.content for m in messages if m.role.value == "system"), None)
        rest = [m for m in messages if m.role.value != "system"]
        if system and not rest:
            return None, [Message(role=Role.USER, content=system)]
        return system, rest

    def _build_messages(self, messages: list[Message]) -> list[dict]:
        # NOTE reasoning is only echoed on the last assistant turn so that cross-engine
        # conversation history doesn't send provider-specific signatures to the wrong API.
        last_assistant_idx = max(
            (i for i, m in enumerate(messages) if m.role.value == "assistant"),
            default=-1,
        )

        result = []
        for idx, m in enumerate(messages):
            if m.role.value == "tool":
                result.append(
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_use_id": m.tool_call_id,
                                "content": m.content,
                            }
                        ],
                    }
                )
            elif m.role.value == "assistant":
                parts: list[dict] = []
                if m.reasoning_content and idx == last_assistant_idx:
                    parts.append(
                        {
                            "type": "thinking",
                            "thinking": m.reasoning_content,
                            "signature": m.reasoning_signature,
                        }
                    )
                if m.content:
                    parts.append({"type": "text", "text": m.content})
                for tc in m.tool_calls or []:
                    parts.append(
                        {
                            "type": "tool_use",
                            "id": tc.id,
                            "name": tc.name,
                            "input": tc.arguments,
                        }
                    )
                result.append({"role": "assistant", "content": parts})
            else:
                result.append({"role": m.role.value, "content": m.content})
        return result

    def _build_tools(self, tools: list[Tool]) -> list[dict]:
        return [
            {
                "name": t.name,
                "description": t.description,
                "input_schema": t.parameters,
            }
            for t in tools
        ]

    def _build_payload(
        self,
        messages: list[Message],
        tools: list[Tool] | None,
        stream: bool = False,
    ) -> dict:
        system, msgs = self._split_system(messages)
        payload: dict = {
            "model": self.model,
            "messages": self._build_messages(msgs),
            "max_tokens": self.max_tokens,
            "thinking": {"type": "enabled", "budget_tokens": self.thinking_budget},
        }
        if system:
            payload["system"] = system
        if tools:
            payload["tools"] = self._build_tools(tools)
        if stream:
            payload["stream"] = True
        return payload

    ### RESPONSE PARSERS

    def _parse_response(self, raw: dict) -> Response:
        content: str | None = None
        tool_calls: list[ToolCall] | None = None
        reasoning_content: str | None = None
        reasoning_signature: str | None = None

        for block in raw.get("content", []):
            if block["type"] == "text":
                content = block["text"]
            elif block["type"] == "thinking":
                reasoning_content = block.get("thinking")
                reasoning_signature = block.get("signature")
            elif block["type"] == "tool_use":
                tool_calls = tool_calls or []
                tool_calls.append(
                    ToolCall(
                        id=block["id"],
                        name=block["name"],
                        arguments=block["input"],
                    )
                )

        usage = raw.get("usage", {})
        in_tok = usage.get("input_tokens", 0)
        out_tok = usage.get("output_tokens", 0)
        cache_read_tokens = usage.get("cache_read_input_tokens", 0)
        cache_write_tokens = usage.get("cache_creation_input_tokens", 0)
        input_tokens, output_tokens, total_tokens = resolve_tokens(in_tok, out_tok, 0)
        stop_map = {
            "end_turn": StopReason.END_TURN,
            "tool_use": StopReason.TOOL_USE,
            "max_tokens": StopReason.MAX_TOKENS,
        }

        return Response(
            provider=self.provider,
            model=self.model,
            content=content,
            tool_calls=tool_calls,
            stop_reason=stop_map.get(
                raw.get("stop_reason", StopReason.END_TURN),
                StopReason.END_TURN,
            ),
            reasoning_content=reasoning_content,
            reasoning_signature=reasoning_signature,
            usage=Usage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cache_read_tokens=cache_read_tokens,
                cache_write_tokens=cache_write_tokens,
            ),
        )

    ### INTERFACES

    def _headers(self) -> dict:
        return {
            "x-api-key": self.api_key,
            "anthropic-version": self.API_VERSION,
            "content-type": "application/json",
        }

    async def chat(
        self,
        messages: list[Message],
        tools: list[Tool] | None = None,
    ) -> Response:
        t_start = time.monotonic()
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{self.BASE_URL}/messages",
                headers=self._headers(),
                json=self._build_payload(messages, tools),
                timeout=120,
            )
            if not resp.is_success:
                raise RuntimeError(
                    f"Anthropic API error {resp.status_code}: {resp.text}"
                )
            response = self._parse_response(resp.json())
        total_duration = time.monotonic() - t_start
        response.performance = make_non_stream_performance(response, total_duration)
        return response

    async def stream(
        self,
        messages: list[Message],
        tools: list[Tool] | None = None,
    ) -> AsyncIterator[Response]:
        input_tokens = 0
        output_tokens = 0
        cache_read_tokens = 0
        cache_write_tokens = 0
        stop_reason: str = StopReason.END_TURN.value
        stop_map = {
            "end_turn": StopReason.END_TURN,
            "tool_use": StopReason.TOOL_USE,
            "max_tokens": StopReason.MAX_TOKENS,
        }

        # tool call accumulators keyed by content_block index.
        tool_buf: dict[int, dict] = {}

        t_start = time.monotonic()
        t_first_chunk: float | None = None
        t_reasoning_start: float | None = None
        t_reasoning_end: float | None = None
        t_content_start: float | None = None

        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                f"{self.BASE_URL}/messages",
                headers=self._headers(),
                json=self._build_payload(messages, tools, stream=True),
                timeout=120,
            ) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):  # pragma: no cover
                        continue
                    now = time.monotonic()
                    if t_first_chunk is None:
                        t_first_chunk = now

                    event = json.loads(line[6:])
                    etype = event.get("type")
                    idx: int = event.get("index", 0)

                    if etype == "message_start":
                        msg_usage = event.get("message", {}).get("usage", {})
                        input_tokens = msg_usage.get("input_tokens", 0)
                        cache_read_tokens = msg_usage.get("cache_read_input_tokens", 0)
                        cache_write_tokens = msg_usage.get(
                            "cache_creation_input_tokens", 0
                        )

                    elif etype == "content_block_start":
                        block = event.get("content_block", {})
                        if block.get("type") == "tool_use":
                            tool_buf[idx] = {
                                "id": block.get("id", ""),
                                "name": block.get("name", ""),
                                "arguments": "",
                            }

                    elif etype == "content_block_delta":
                        delta = event.get("delta", {})
                        dtype = delta.get("type")

                        if dtype == "thinking_delta":
                            if t_reasoning_start is None:
                                t_reasoning_start = now
                            yield Response(
                                provider=self.provider,
                                model=self.model,
                                streaming=True,
                                reasoning_content=delta.get("thinking", ""),
                            )
                        elif dtype == "signature_delta":
                            yield Response(
                                provider=self.provider,
                                model=self.model,
                                streaming=True,
                                reasoning_signature=delta.get("signature"),
                            )
                        elif dtype == "text_delta":
                            if t_content_start is None:
                                t_content_start = now
                                if (
                                    t_reasoning_start is not None
                                    and t_reasoning_end is None
                                ):
                                    t_reasoning_end = now
                            yield Response(
                                provider=self.provider,
                                model=self.model,
                                streaming=True,
                                content=delta.get("text", ""),
                            )
                        elif dtype == "input_json_delta" and idx in tool_buf:
                            tool_buf[idx]["arguments"] += delta.get("partial_json", "")

                    elif etype == "content_block_stop" and idx in tool_buf:
                        buf = tool_buf.pop(idx)
                        try:
                            args = json.loads(buf["arguments"])
                        except json.JSONDecodeError:  # pragma: no cover
                            args = {}
                        yield Response(
                            provider=self.provider,
                            model=self.model,
                            streaming=True,
                            tool_calls=[
                                ToolCall(id=buf["id"], name=buf["name"], arguments=args)
                            ],
                        )

                    elif etype == "message_delta":
                        stop_reason = event.get("delta", {}).get(
                            "stop_reason", StopReason.END_TURN
                        )
                        output_tokens = event.get("usage", {}).get("output_tokens", 0)

        t_end = time.monotonic()
        total_duration = t_end - t_start
        latency = (t_first_chunk - t_start) if t_first_chunk else total_duration
        reasoning_duration = (
            ((t_reasoning_end or t_end) - t_reasoning_start)
            if t_reasoning_start
            else 0.0
        )
        content_duration = (t_end - t_content_start) if t_content_start else 0.0

        in_tok, out_tok, tot_tok = resolve_tokens(input_tokens, output_tokens, 0)
        throughput = int(out_tok / total_duration) if total_duration > 0 else 0

        yield Response(
            provider=self.provider,
            model=self.model,
            streaming=True,
            stop_reason=stop_map.get(stop_reason, StopReason.END_TURN),
            usage=Usage(
                input_tokens=in_tok,
                output_tokens=out_tok,
                total_tokens=tot_tok,
                cache_read_tokens=cache_read_tokens,
                cache_write_tokens=cache_write_tokens,
            ),
            performance=Performance(
                latency=round(latency, 2),
                reasoning_duration=round(reasoning_duration, 2),
                content_duration=round(content_duration, 2),
                total_duration=round(total_duration, 2),
                throughput=throughput,
            ),
        )
