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,
    StopReason,
    Tool,
    ToolCall,
    Usage,
    make_non_stream_performance,
    resolve_tokens,
)


class GoogleEngine(BaseEngine):
    BASE_URL = "https://generativelanguage.googleapis.com/v1beta"

    def __init__(
        self,
        model: str,
        api_key: str,
        max_tokens: int = 8192,
        thinking_budget: int | None = None,
    ) -> None:
        self.provider = "google"
        self.model = model
        self.api_key = api_key
        self.max_tokens = max_tokens
        # NOTE -1 = dynamic budget (Gemini decides); pass 0 to disable thinking.
        self.thinking_budget = thinking_budget if thinking_budget is not None else -1

    ### PAYLOAD BUILDERS

    def _build_contents(self, messages: list[Message]) -> tuple[str | None, list[dict]]:
        """Returns (system_instruction, contents).
        Gemini takes system prompt separately via systemInstruction.

        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 entry in contents.
        """

        system: str | None = None

        non_system = [m for m in messages if m.role.value != "system"]
        if not non_system:
            sys_content = next(
                (m.content for m in messages if m.role.value == "system"), None
            )
            return None, [{"role": "user", "parts": [{"text": sys_content}]}]

        # NOTE reasoning signature is only echoed on the last assistant turn so that
        # cross-engine conversation history doesn't send provider-specific data to the wrong API.
        last_assistant_idx = max(
            (i for i, m in enumerate(non_system) if m.role.value == "assistant"),
            default=-1,
        )

        contents = []
        for idx, m in enumerate(non_system):
            if m.role.value == "tool":
                # NOTE tool_call_id holds the function name for Gemini.
                contents.append(
                    {
                        "role": "user",
                        "parts": [
                            {
                                "functionResponse": {
                                    "name": m.tool_call_id,
                                    "response": {"result": m.content},
                                }
                            }
                        ],
                    }
                )
            elif m.role.value == "assistant":
                parts: list[dict] = []
                if m.content:
                    parts.append({"text": m.content})
                for tc in m.tool_calls or []:
                    parts.append(
                        {
                            "functionCall": {
                                "name": tc.name,
                                "args": tc.arguments,
                            }
                        }
                    )
                # NOTE thoughtSignature goes directly in parts[0] (text or functionCall part),
                # only on the last assistant turn.
                if parts and m.reasoning_signature and idx == last_assistant_idx:
                    parts[0]["thoughtSignature"] = m.reasoning_signature
                contents.append({"role": "model", "parts": parts})
            else:
                contents.append({"role": "user", "parts": [{"text": m.content}]})

        # extract system from the original messages list (it was skipped above).
        for m in messages:
            if m.role.value == "system":
                system = m.content
                break

        return system, contents

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

    def _build_payload(self, messages: list[Message], tools: list[Tool] | None) -> dict:
        system, contents = self._build_contents(messages)
        thinking_config: dict = {"includeThoughts": True}
        if self.thinking_budget > 0:
            thinking_config["thinkingBudget"] = self.thinking_budget
        payload: dict = {
            "contents": contents,
            "generationConfig": {
                "maxOutputTokens": self.max_tokens,
                "thinkingConfig": thinking_config,
            },
        }
        if system:
            payload["systemInstruction"] = {"parts": [{"text": system}]}
        if tools:
            payload["tools"] = self._build_tools(tools)
        return payload

    ### RESPONSE PARSERS

    def _parse_response(self, raw: dict) -> Response:
        """Parse a non-streaming generateContent response."""
        candidate = raw["candidates"][0]
        parts = candidate.get("content", {}).get("parts", [])

        content: str | None = None
        tool_calls: list[ToolCall] | None = None
        reasoning_content: str | None = None
        reasoning_signature: str | None = None

        for part in parts:
            if part.get("thought"):
                # NOTE thought parts carry the chain-of-thought text and signature.
                reasoning_content = part.get("text") or None
                reasoning_signature = part.get("thoughtSignature") or None
            elif "text" in part:
                content = part["text"] or None
            elif "functionCall" in part:
                fc = part["functionCall"]
                tool_calls = tool_calls or []
                # NOTE gemini has no call id — use the function name as id so the
                # agent can set tool_call_id = tc.id on the tool result message.
                tool_calls.append(
                    ToolCall(
                        id=fc["name"],
                        name=fc["name"],
                        arguments=fc.get("args", {}),
                    )
                )

        finish_reason = candidate.get("finishReason")
        stop_map = {
            "STOP": StopReason.END_TURN,
            "MAX_TOKENS": StopReason.MAX_TOKENS,
            # NOTE we set this ourselves, doesn't exist on google's end
            "TOOL": StopReason.TOOL_USE,
        }

        if tool_calls:
            stop_reason = "TOOL"
        elif finish_reason:
            stop_reason = finish_reason
        else:  # pragma: no cover
            stop_reason = ""

        usage = raw.get("usageMetadata", {})
        in_tok = usage.get("promptTokenCount", 0)
        # NOTE candidatesTokenCount may exclude thoughtsTokenCount; trust totalTokenCount.
        out_tok = usage.get("candidatesTokenCount", 0)
        tot_tok = usage.get("totalTokenCount", 0)
        cache_read_tokens = usage.get("cachedContentTokenCount", 0)
        input_tokens, output_tokens, total_tokens = resolve_tokens(
            in_tok, out_tok, tot_tok
        )

        return Response(
            provider=self.provider,
            model=self.model,
            content=content,
            tool_calls=tool_calls,
            stop_reason=stop_map.get(stop_reason, 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,
            ),
        )

    ### INTERFACES

    def _url(self, method: str) -> str:
        return f"{self.BASE_URL}/models/{self.model}:{method}?key={self.api_key}"

    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(
                self._url("generateContent"),
                json=self._build_payload(messages, tools),
                timeout=120,
            )
            if not resp.is_success:
                raise RuntimeError(f"Google 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]:
        stop_map = {
            "STOP": StopReason.END_TURN,
            "MAX_TOKENS": StopReason.MAX_TOKENS,
            # NOTE we set this ourselves, doesn't exist on google's end
            "TOOL": StopReason.TOOL_USE,
        }
        input_tokens = output_tokens = total_tokens = 0
        cache_read_tokens = 0
        stop_reason: str = StopReason.END_TURN.value
        has_tool: bool = False

        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",
                self._url("streamGenerateContent") + "&alt=sse",
                json=self._build_payload(messages, tools),
                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

                    chunk = json.loads(line[6:])
                    try:
                        candidate = chunk["candidates"][0]
                    except (KeyError, IndexError):  # pragma: no cover
                        continue

                    parts = candidate.get("content", {}).get("parts", [])
                    finish_reason = candidate.get("finishReason")

                    for part in parts:
                        if part.get("thought"):
                            rc = part.get("text") or None
                            rs = part.get("thoughtSignature") or None
                            if rc:
                                if t_reasoning_start is None:
                                    t_reasoning_start = now
                                yield Response(
                                    provider=self.provider,
                                    model=self.model,
                                    streaming=True,
                                    reasoning_content=rc,
                                )
                            if rs:
                                yield Response(
                                    provider=self.provider,
                                    model=self.model,
                                    streaming=True,
                                    reasoning_signature=rs,
                                )
                        elif "text" in part:
                            text = part["text"] or None
                            if text:
                                if t_content_start is None:
                                    t_content_start = now
                                    if (
                                        t_reasoning_start is not None
                                        and t_reasoning_end is None
                                    ):  # pragma: no cover
                                        t_reasoning_end = now
                                yield Response(
                                    provider=self.provider,
                                    model=self.model,
                                    streaming=True,
                                    content=text,
                                )
                        elif "functionCall" in part:
                            fc = part["functionCall"]
                            has_tool = True
                            yield Response(
                                provider=self.provider,
                                model=self.model,
                                streaming=True,
                                tool_calls=[
                                    ToolCall(
                                        id=fc["name"],
                                        name=fc["name"],
                                        arguments=fc.get("args", {}),
                                    )
                                ],
                            )

                    if finish_reason:
                        usage = chunk.get("usageMetadata", {})
                        in_tok = usage.get("promptTokenCount", 0)
                        out_tok = usage.get("candidatesTokenCount", 0)
                        tot_tok = usage.get("totalTokenCount", 0)
                        cache_read_tokens = usage.get("cachedContentTokenCount", 0)
                        input_tokens, output_tokens, total_tokens = resolve_tokens(
                            in_tok, out_tok, tot_tok
                        )
                        stop_reason = "TOOL" if has_tool else finish_reason

        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
        throughput = int(output_tokens / 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=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cache_read_tokens=cache_read_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,
            ),
        )
