"""Google Gemini LLM implementation for Fluxibly.

Uses the Google Generative AI SDK with support for thinking,
function calling, code execution, and multimodal I/O.
"""

from __future__ import annotations

import os
from typing import Any

from fluxibly.exceptions import UnsupportedContentTypeError
from fluxibly.llm.base import BaseLLM, LLMConfig
from fluxibly.schema.response import (
    ContentItem,
    LLMMetadata,
    LLMOutput,
    LLMResponse,
    TokenUsage,
)
from fluxibly.schema.tools import Citation, ToolCall

try:
    import google.generativeai as genai
except ImportError as e:
    raise ImportError(
        "Google Generative AI package is required for GeminiLLM. "
        "Install it with: pip install fluxibly[gemini]"
    ) from e


class GeminiLLM(BaseLLM):
    """Google Gemini LLM implementation."""

    def __init__(self, config: LLMConfig) -> None:
        super().__init__(config)
        api_key = (
            config.api_key
            or os.environ.get("GOOGLE_API_KEY")
            or os.environ.get("GEMINI_API_KEY")
        )
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel(config.model)

    def prepare(self) -> dict[str, Any]:
        """Format for Gemini GenerateContent API.

        Gemini uses a unified Part structure for all content types.
        System instruction is set on the model, not in messages.
        """
        prepared: dict[str, Any] = {
            "contents": [],
        }

        # Extract system message
        for msg in self.messages:
            if msg["role"] == "system":
                prepared["system_instruction"] = msg["content"]
            else:
                role = "model" if msg["role"] == "assistant" else "user"
                parts = self._convert_content_to_parts(msg.get("content", ""))
                prepared["contents"].append({"role": role, "parts": parts})

        # Tools
        if self.tools:
            prepared["tools"] = self._format_tools()

        # Generation config
        gen_config: dict[str, Any] = {}
        if self.config.max_output_tokens:
            gen_config["max_output_tokens"] = self.config.max_output_tokens
        if self.config.temperature is not None:
            gen_config["temperature"] = self.config.temperature
        if self.config.top_p is not None:
            gen_config["top_p"] = self.config.top_p
        if self.config.stop:
            gen_config["stop_sequences"] = self.config.stop

        # Thinking (Gemini 2.5+)
        if self.config.reasoning and self.config.reasoning.enabled:
            gen_config["thinking_config"] = {
                "thinking_budget": self.config.reasoning.budget_tokens or 8192,
            }

        if gen_config:
            prepared["generation_config"] = gen_config

        return prepared

    def forward(self, **kwargs: Any) -> LLMResponse:
        """Execute Gemini inference."""
        prepared = self.prepare()

        self.logger.log_input(prepared)

        try:
            # Extract system instruction if present
            system = prepared.pop("system_instruction", None)
            if system:
                self.model = genai.GenerativeModel(
                    self.config.model,
                    system_instruction=system,
                )

            raw = self.model.generate_content(**prepared, **kwargs)
            result = self._parse_output(raw)
            self._apply_pricing(result)
            self.logger.log_output(result)
            return result
        except Exception as e:
            self.logger.log_error(e, prepared)
            raise

    # ── Response Parsing ─────────────────────────────────────────────

    def _parse_output(self, raw: Any) -> LLMResponse:
        """Parse Gemini GenerateContent output.

        Handles all Part types:
        - text (with optional thought=true for reasoning)
        - function_call       → tool call
        - inline_data         → image, audio, or other media output
        - executable_code     → code execution
        - code_execution_result → code execution result
        - Unknown Part fields → preserved with extra dict
        """
        if not raw.candidates:
            from fluxibly.exceptions import LLMError

            raise LLMError("Gemini returned empty response — no candidates")
        candidate = raw.candidates[0]
        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []
        reasoning_parts: list[str] = []
        citations: list[Citation] = []

        for part in candidate.content.parts:
            # Text (regular or thinking/reasoning)
            if hasattr(part, "text") and part.text:
                if getattr(part, "thought", False):
                    reasoning_parts.append(part.text)
                    content_items.append(
                        ContentItem(type="reasoning", reasoning_text=part.text)
                    )
                else:
                    text_parts.append(part.text)
                    content_items.append(
                        ContentItem(type="text", text=part.text)
                    )

            # Function calls
            elif hasattr(part, "function_call") and part.function_call:
                fc = part.function_call
                tc = ToolCall(
                    id=f"gemini_{fc.name}_{id(fc)}",
                    name=fc.name,
                    arguments=(
                        dict(fc.args) if hasattr(fc.args, "items") else fc.args
                    ),
                    type="function",
                )
                tool_calls.append(tc)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tc)
                )

            # Inline data — image, audio, video
            elif hasattr(part, "inline_data") and part.inline_data:
                mime = part.inline_data.mime_type
                data = part.inline_data.data

                if mime.startswith("image/"):
                    content_items.append(
                        ContentItem(
                            type="image",
                            image_data=data,
                            image_media_type=mime,
                        )
                    )
                elif mime.startswith("audio/"):
                    content_items.append(
                        ContentItem(
                            type="audio",
                            audio_data=data,
                            audio_media_type=mime,
                        )
                    )
                elif mime.startswith("video/"):
                    content_items.append(
                        ContentItem(
                            type="video",
                            video_data=data,
                            video_media_type=mime,
                        )
                    )
                else:
                    content_items.append(
                        ContentItem(
                            type="media",
                            extra={
                                "mime_type": mime,
                                "data_length": len(data) if data else 0,
                            },
                        )
                    )

            # Code execution
            elif hasattr(part, "executable_code") and part.executable_code:
                content_items.append(
                    ContentItem(
                        type="code",
                        code=part.executable_code.code,
                        code_language=str(
                            getattr(part.executable_code, "language", "python")
                        ).lower(),
                    )
                )

            # Code execution result
            elif (
                hasattr(part, "code_execution_result")
                and part.code_execution_result
            ):
                content_items.append(
                    ContentItem(
                        type="code_result",
                        code_output=part.code_execution_result.output,
                        extra={
                            "outcome": str(
                                getattr(
                                    part.code_execution_result, "outcome", ""
                                )
                            )
                        },
                    )
                )

            # Unknown Part types
            else:
                content_items.append(
                    ContentItem(
                        type="unknown",
                        extra=(
                            part.to_dict()
                            if hasattr(part, "to_dict")
                            else {"raw": str(part)}
                        ),
                    )
                )

        # Grounding citations
        if (
            hasattr(candidate, "grounding_metadata")
            and candidate.grounding_metadata
        ):
            gm = candidate.grounding_metadata
            if hasattr(gm, "grounding_chunks") and gm.grounding_chunks:
                for chunk in gm.grounding_chunks:
                    citations.append(
                        Citation(
                            type="url",
                            url=(
                                getattr(chunk.web, "uri", None)
                                if hasattr(chunk, "web")
                                else None
                            ),
                            title=(
                                getattr(chunk.web, "title", None)
                                if hasattr(chunk, "web")
                                else None
                            ),
                        )
                    )

        # Normalize finish reason
        finish_map = {
            "STOP": "stop",
            "MAX_TOKENS": "max_tokens",
            "SAFETY": "content_filter",
            "RECITATION": "content_filter",
        }

        usage_meta = (
            raw.usage_metadata if hasattr(raw, "usage_metadata") else None
        )

        images = [ci for ci in content_items if ci.type == "image"]
        audio = [ci for ci in content_items if ci.type == "audio"]

        return LLMResponse(
            output=LLMOutput(
                output_text="\n".join(text_parts),
                content=content_items,
                tool_calls=tool_calls,
                reasoning=(
                    "\n".join(reasoning_parts) if reasoning_parts else None
                ),
                citations=citations,
                images=images,
                audio=audio,
            ),
            metadata=LLMMetadata(
                model=self.config.model,
                id=f"gemini_{id(raw)}",
                created_at="",
                status=(
                    "completed"
                    if str(candidate.finish_reason) == "STOP"
                    else "incomplete"
                ),
                usage=TokenUsage(
                    input_tokens=(
                        getattr(usage_meta, "prompt_token_count", 0)
                        if usage_meta
                        else 0
                    ),
                    output_tokens=(
                        getattr(usage_meta, "candidates_token_count", 0)
                        if usage_meta
                        else 0
                    ),
                    total_tokens=(
                        getattr(usage_meta, "total_token_count", 0)
                        if usage_meta
                        else 0
                    ),
                ),
                stop_reason=finish_map.get(
                    str(candidate.finish_reason), str(candidate.finish_reason)
                ),
                provider="gemini",
            ),
        )

    # ── Content Conversion ───────────────────────────────────────────

    def _convert_content_to_parts(
        self, content: str | list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Convert our unified content format to Gemini Parts."""
        if isinstance(content, str):
            return [{"text": content}]

        parts: list[dict[str, Any]] = []
        for item in content:
            ptype = item.get("type", "text")
            if ptype == "text":
                parts.append({"text": item["text"]})
            elif ptype == "image":
                source = item.get("source", {})
                if source.get("type") == "base64":
                    parts.append(
                        {
                            "inline_data": {
                                "mime_type": source.get(
                                    "media_type", "image/png"
                                ),
                                "data": source["data"],
                            }
                        }
                    )
                elif source.get("type") == "url":
                    parts.append(
                        {
                            "file_data": {
                                "mime_type": source.get(
                                    "media_type", "image/png"
                                ),
                                "file_uri": source["url"],
                            }
                        }
                    )
            elif ptype == "audio":
                source = item.get("source", {})
                parts.append(
                    {
                        "inline_data": {
                            "mime_type": source.get("media_type", "audio/wav"),
                            "data": source.get("data", ""),
                        }
                    }
                )
            elif ptype == "video":
                source = item.get("source", {})
                if source.get("type") == "url":
                    parts.append(
                        {
                            "file_data": {
                                "mime_type": source.get(
                                    "media_type", "video/mp4"
                                ),
                                "file_uri": source["url"],
                            }
                        }
                    )
                elif source.get("type") == "base64":
                    parts.append(
                        {
                            "inline_data": {
                                "mime_type": source.get(
                                    "media_type", "video/mp4"
                                ),
                                "data": source["data"],
                            }
                        }
                    )
            elif ptype == "document":
                source = item.get("source", {})
                parts.append(
                    {
                        "inline_data": {
                            "mime_type": source.get(
                                "media_type", "application/pdf"
                            ),
                            "data": source.get("data", ""),
                        }
                    }
                )
            else:
                raise UnsupportedContentTypeError(
                    f"Gemini does not support content type: {ptype}"
                )
        return parts

    # ── Tool Formatting ──────────────────────────────────────────────

    def _format_tools(self) -> list[dict[str, Any]]:
        """Format tools for Gemini API."""
        function_declarations: list[dict[str, Any]] = []
        for tool in self.tools:
            ttype = tool.get("type", "function")
            if ttype == "function":
                func = tool.get("function", {})
                decl: dict[str, Any] = {
                    "name": func.get("name", ""),
                    "description": func.get("description", ""),
                }
                params = func.get("parameters")
                if params:
                    decl["parameters"] = self._convert_json_schema_to_gemini(
                        params
                    )
                function_declarations.append(decl)

        if function_declarations:
            return [{"function_declarations": function_declarations}]
        return []

    def _convert_json_schema_to_gemini(
        self, schema: dict[str, Any]
    ) -> dict[str, Any]:
        """Convert JSON Schema to Gemini's parameter format.

        Gemini doesn't support all JSON Schema features, so we simplify.
        """
        converted: dict[str, Any] = {
            "type": schema.get("type", "object").upper()
        }

        if "properties" in schema:
            converted["properties"] = {}
            for name, prop in schema["properties"].items():
                converted["properties"][name] = {
                    "type": prop.get("type", "string").upper(),
                    "description": prop.get("description", ""),
                }
                if "enum" in prop:
                    converted["properties"][name]["enum"] = prop["enum"]

        if "required" in schema:
            converted["required"] = schema["required"]

        return converted
