from __future__ import annotations

from urllib.parse import quote

from ..models import ProviderResponse, ToolProposal
from ..scenarios import Scenario
from .base import (
    SYSTEM_INSTRUCTION,
    TOOL_PARAMETERS,
    ProviderAdapter,
    post_json,
    scenario_prompt,
)


def _gemini_schema(value):
    """Convert the shared JSON schema to Gemini's supported schema subset."""
    if isinstance(value, dict):
        return {
            key: _gemini_schema(item)
            for key, item in value.items()
            if key not in {"additionalProperties", "$schema"}
        }
    if isinstance(value, list):
        return [_gemini_schema(item) for item in value]
    return value


class GeminiAdapter(ProviderAdapter):
    provider = "gemini"

    def propose(self, scenario: Scenario) -> ProviderResponse:
        payload = {
            "systemInstruction": {"parts": [{"text": SYSTEM_INSTRUCTION}]},
            "contents": [
                {"role": "user", "parts": [{"text": scenario_prompt(scenario)}]}
            ],
            "tools": [
                {
                    "functionDeclarations": [
                        {
                            "name": "propose_action",
                            "description": "Propose one synthetic action for policy evaluation.",
                            "parameters": _gemini_schema(TOOL_PARAMETERS),
                        }
                    ]
                }
            ],
            "toolConfig": {
                "functionCallingConfig": {
                    "mode": "ANY",
                    "allowedFunctionNames": ["propose_action"]
                }
            },
            "generationConfig": {"maxOutputTokens": 500},
        }
        model = quote(self.model, safe="")
        key = quote(self._api_key, safe="")
        raw = post_json(
            f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}",
            {"Content-Type": "application/json"},
            payload,
            self.timeout_seconds,
        )
        proposals: list[ToolProposal] = []
        text_parts: list[str] = []
        for candidate in raw.get("candidates", []):
            for part in candidate.get("content", {}).get("parts", []):
                if "functionCall" in part:
                    function_call = part["functionCall"]
                    proposals.append(
                        ToolProposal(
                            name=str(function_call.get("name", "")),
                            arguments=dict(function_call.get("args", {})),
                            call_id=str(function_call.get("id", "")),
                        )
                    )
                elif "text" in part:
                    text_parts.append(str(part["text"]))
        usage = raw.get("usageMetadata", {})
        return ProviderResponse(
            provider=self.provider,
            model=self.model,
            text="\n".join(text_parts),
            proposals=tuple(proposals),
            input_tokens=int(usage.get("promptTokenCount", 0)),
            output_tokens=int(
                usage.get("candidatesTokenCount", 0)
                + usage.get("thoughtsTokenCount", 0)
            ),
            raw_stop_reason=str(
                raw.get("candidates", [{}])[0].get("finishReason", "")
                if raw.get("candidates")
                else ""
            ),
        )
