"""OpenAI LLM implementation for Fluxibly.

Supports both Responses API (default) and Chat Completions API.
"""

from __future__ import annotations

import json
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 openai
except ImportError as e:
    raise ImportError(
        "OpenAI package is required for OpenAILLM. "
        "Install it with: pip install fluxibly[openai]"
    ) from e


class OpenAILLM(BaseLLM):
    """OpenAI LLM implementation.

    Supports:
    - Responses API (default, recommended) — api_type="responses"
    - Chat Completions API (legacy) — api_type="chat"
    """

    def __init__(self, config: LLMConfig) -> None:
        super().__init__(config)
        self.client = openai.OpenAI(
            api_key=config.api_key or os.environ.get("OPENAI_API_KEY"),
            base_url=config.api_base,
            timeout=config.timeout,
            max_retries=config.max_retries,
        )

    def prepare(self) -> dict[str, Any]:
        """Format messages and tools for OpenAI API."""
        if self.config.api_type == "responses":
            return self._prepare_responses_api()
        return self._prepare_chat_api()

    def _prepare_responses_api(self) -> dict[str, Any]:
        """Prepare inputs for OpenAI Responses API format."""
        prepared: dict[str, Any] = {
            "model": self.config.model,
            "input": [],
        }

        # Extract system message as "instructions"
        for msg in self.messages:
            if msg["role"] == "system":
                prepared["instructions"] = msg["content"]
            elif msg["role"] == "assistant" and msg.get("tool_calls"):
                # Assistant message with tool calls → emit function_call items
                for tc in msg["tool_calls"]:
                    args = tc.get("arguments", tc.get("input", ""))
                    if isinstance(args, dict):
                        args = json.dumps(args)
                    prepared["input"].append(
                        {
                            "type": "function_call",
                            "call_id": tc["id"],
                            "name": tc["name"],
                            "arguments": args,
                        }
                    )
            elif msg["role"] == "tool":
                # Tool result → emit function_call_output item
                prepared["input"].append(
                    {
                        "type": "function_call_output",
                        "call_id": msg["tool_call_id"],
                        "output": str(msg.get("content", "")),
                    }
                )
            else:
                # Regular user/assistant message — convert multimodal content
                prepared["input"].append(
                    self._convert_message_for_responses(msg)
                )

        # Format tools
        if self.tools:
            prepared["tools"] = self._format_tools_responses()

        # Skills — inject into shell tool's environment (Agent Skills standard)
        if self.skills:
            self._inject_skills_into_shell(prepared)

        # Generation params
        reasoning_enabled = bool(
            self.config.reasoning and self.config.reasoning.enabled
        )
        if self.config.max_output_tokens:
            prepared["max_output_tokens"] = self.config.max_output_tokens
        if not reasoning_enabled:
            # Reasoning models reject temperature/top_p
            if self.config.temperature is not None:
                prepared["temperature"] = self.config.temperature
            if self.config.top_p is not None:
                prepared["top_p"] = self.config.top_p
        if reasoning_enabled:
            prepared["reasoning"] = {"effort": self.config.reasoning.effort}

        return prepared

    def _prepare_chat_api(self) -> dict[str, Any]:
        """Prepare inputs for OpenAI Chat Completions API format."""
        prepared: dict[str, Any] = {
            "model": self.config.model,
            "messages": [
                self._convert_message_for_chat(m) for m in self.messages
            ],
        }

        reasoning_enabled = bool(
            self.config.reasoning and self.config.reasoning.enabled
        )

        if self.tools:
            prepared["tools"] = self._format_tools_chat()

        if self.config.max_output_tokens:
            prepared["max_tokens"] = self.config.max_output_tokens
        if not reasoning_enabled:
            # Reasoning models reject temperature/top_p/penalties
            if self.config.temperature is not None:
                prepared["temperature"] = self.config.temperature
            if self.config.top_p is not None:
                prepared["top_p"] = self.config.top_p
            if self.config.frequency_penalty is not None:
                prepared["frequency_penalty"] = self.config.frequency_penalty
            if self.config.presence_penalty is not None:
                prepared["presence_penalty"] = self.config.presence_penalty
        if self.config.stop:
            prepared["stop"] = self.config.stop

        return prepared

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

        self.logger.log_input(prepared)

        try:
            if self.config.api_type == "responses":
                raw = self.client.responses.create(**prepared)
                result = self._parse_responses_output(raw)
            else:
                raw = self.client.chat.completions.create(**prepared)
                result = self._parse_chat_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_responses_output(self, raw: Any) -> LLMResponse:
        """Parse OpenAI Responses API output to LLMResponse."""
        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []
        reasoning_parts: list[str] = []
        citations: list[Citation] = []

        for item in raw.output:
            if item.type == "message":
                for block in item.content:
                    if block.type == "output_text":
                        text_parts.append(block.text)
                        block_citations: list[Citation] | None = None
                        if hasattr(block, "annotations") and block.annotations:
                            block_citations = []
                            for ann in block.annotations:
                                cit = Citation(
                                    type=getattr(ann, "type", "url"),
                                    url=getattr(ann, "url", None),
                                    title=getattr(ann, "title", None),
                                    text=getattr(ann, "text", None),
                                    start_index=getattr(
                                        ann, "start_index", None
                                    ),
                                    end_index=getattr(ann, "end_index", None),
                                )
                                block_citations.append(cit)
                                citations.append(cit)
                        content_items.append(
                            ContentItem(
                                type="text",
                                text=block.text,
                                annotations=block_citations,
                            )
                        )
                    elif block.type == "refusal":
                        content_items.append(
                            ContentItem(
                                type="refusal",
                                refusal=block.refusal,
                            )
                        )

            elif item.type == "reasoning":
                for summary_item in item.summary or []:
                    reasoning_parts.append(summary_item.text)
                    content_items.append(
                        ContentItem(
                            type="reasoning",
                            reasoning_text=summary_item.text,
                        )
                    )

            elif item.type == "function_call":
                tc = ToolCall(
                    id=item.call_id,
                    name=item.name,
                    arguments=item.arguments,
                    type="function",
                )
                tool_calls.append(tc)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tc)
                )

            elif item.type == "image_generation_call":
                content_items.append(
                    ContentItem(
                        type="image",
                        image_data=getattr(item, "result", None),
                        image_media_type="image/png",
                    )
                )

            else:
                # Unknown output types — preserve for forward-compatibility
                content_items.append(
                    ContentItem(
                        type=item.type,
                        extra=(
                            item.model_dump()
                            if hasattr(item, "model_dump")
                            else {"raw": str(item)}
                        ),
                    )
                )

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

        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,
            ),
            metadata=LLMMetadata(
                model=raw.model,
                id=raw.id,
                created_at=str(getattr(raw, "created_at", "")),
                status=raw.status,
                usage=TokenUsage(
                    input_tokens=raw.usage.input_tokens,
                    output_tokens=raw.usage.output_tokens,
                    total_tokens=raw.usage.total_tokens,
                    reasoning_tokens=getattr(
                        getattr(raw.usage, "output_tokens_details", None),
                        "reasoning_tokens",
                        None,
                    ),
                ),
                stop_reason=self._normalize_stop_reason_responses(raw),
                provider="openai",
                raw_response=(
                    raw.model_dump() if hasattr(raw, "model_dump") else None
                ),
            ),
        )

    def _parse_chat_output(self, raw: Any) -> LLMResponse:
        """Parse OpenAI Chat Completions API output to LLMResponse."""
        if not raw.choices:
            from fluxibly.exceptions import LLMError

            raise LLMError("OpenAI returned empty response — no choices")
        choice = raw.choices[0]
        message = choice.message

        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []

        # Text content
        if message.content:
            text_parts.append(message.content)
            content_items.append(
                ContentItem(type="text", text=message.content)
            )

        # Refusal
        if hasattr(message, "refusal") and message.refusal:
            content_items.append(
                ContentItem(type="refusal", refusal=message.refusal)
            )

        # Audio output (gpt-4o-audio models)
        if hasattr(message, "audio") and message.audio:
            content_items.append(
                ContentItem(
                    type="audio",
                    audio_data=message.audio.data,
                    audio_media_type="audio/wav",
                    audio_transcript=getattr(
                        message.audio, "transcript", None
                    ),
                )
            )

        # Tool calls
        if message.tool_calls:
            for tc in message.tool_calls:
                tool_call = ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=tc.function.arguments,
                    type="function",
                )
                tool_calls.append(tool_call)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tool_call)
                )

        audio_items = [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,
                audio=audio_items,
            ),
            metadata=LLMMetadata(
                model=raw.model,
                id=raw.id,
                created_at=str(getattr(raw, "created", "")),
                status=(
                    "completed"
                    if choice.finish_reason == "stop"
                    else "incomplete"
                ),
                usage=TokenUsage(
                    input_tokens=raw.usage.prompt_tokens,
                    output_tokens=raw.usage.completion_tokens,
                    total_tokens=raw.usage.total_tokens,
                    reasoning_tokens=getattr(
                        getattr(raw.usage, "completion_tokens_details", None),
                        "reasoning_tokens",
                        None,
                    ),
                ),
                stop_reason=self._normalize_stop_reason_chat(
                    choice.finish_reason
                ),
                provider="openai",
                raw_response=(
                    raw.model_dump() if hasattr(raw, "model_dump") else None
                ),
            ),
        )

    # ── Stop Reason Normalization ────────────────────────────────────

    def _normalize_stop_reason_responses(self, raw: Any) -> str:
        status = raw.status
        if status == "completed":
            has_tool_calls = any(
                item.type == "function_call" for item in raw.output
            )
            return "tool_calls" if has_tool_calls else "stop"
        if status == "incomplete":
            reason = getattr(raw, "incomplete_details", None)
            if reason and getattr(reason, "reason", "") == "max_output_tokens":
                return "max_tokens"
            return "incomplete"
        if status == "failed":
            return "error"
        return status

    def _normalize_stop_reason_chat(self, finish_reason: str) -> str:
        mapping = {
            "stop": "stop",
            "tool_calls": "tool_calls",
            "length": "max_tokens",
            "content_filter": "content_filter",
            "function_call": "tool_calls",
        }
        return mapping.get(finish_reason, finish_reason)

    # ── Message Conversion ───────────────────────────────────────────

    def _convert_message_for_responses(
        self, msg: dict[str, Any]
    ) -> dict[str, Any]:
        """Convert our unified message format to Responses API format."""
        content = msg.get("content", "")

        if isinstance(content, str):
            return msg

        # Multimodal content — convert each part
        converted_parts: list[dict[str, Any]] = []
        for part in content:
            ptype = part.get("type", "text")
            if ptype == "text":
                converted_parts.append(
                    {"type": "input_text", "text": part["text"]}
                )
            elif ptype == "image":
                source = part.get("source", {})
                if source.get("type") == "base64":
                    data_uri = (
                        f"data:{source['media_type']};base64,{source['data']}"
                    )
                    converted_parts.append(
                        {"type": "input_image", "image_url": data_uri}
                    )
                elif source.get("type") == "url":
                    converted_parts.append(
                        {"type": "input_image", "image_url": source["url"]}
                    )
            elif ptype == "audio":
                source = part.get("source", {})
                converted_parts.append(
                    {
                        "type": "input_audio",
                        "data": source.get("data", ""),
                        "format": source.get("media_type", "audio/wav").split(
                            "/"
                        )[-1],
                    }
                )
            elif ptype == "document":
                source = part.get("source", {})
                if source.get("type") == "base64":
                    data_uri = (
                        f"data:{source['media_type']};base64,{source['data']}"
                    )
                    converted_parts.append(
                        {
                            "type": "input_file",
                            "file_data": data_uri,
                            "filename": part.get("filename", "document"),
                        }
                    )
            else:
                raise UnsupportedContentTypeError(
                    f"OpenAI Responses API does not support content type: {ptype}"
                )

        return {**msg, "content": converted_parts}

    def _convert_message_for_chat(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Convert our unified message format to Chat Completions API format."""
        content = msg.get("content", "")

        if isinstance(content, str):
            return msg

        # Multimodal content
        converted_parts: list[dict[str, Any]] = []
        for part in content:
            ptype = part.get("type", "text")
            if ptype == "text":
                converted_parts.append({"type": "text", "text": part["text"]})
            elif ptype == "image":
                source = part.get("source", {})
                if source.get("type") == "base64":
                    url = (
                        f"data:{source['media_type']};base64,{source['data']}"
                    )
                elif source.get("type") == "url":
                    url = source["url"]
                else:
                    continue
                converted_parts.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": url},
                    }
                )
            elif ptype == "audio":
                source = part.get("source", {})
                converted_parts.append(
                    {
                        "type": "input_audio",
                        "input_audio": {
                            "data": source.get("data", ""),
                            "format": source.get(
                                "media_type", "audio/wav"
                            ).split("/")[-1],
                        },
                    }
                )
            else:
                raise UnsupportedContentTypeError(
                    f"OpenAI Chat API does not support content type: {ptype}"
                )

        return {**msg, "content": converted_parts}

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

    def _format_tools_responses(self) -> list[dict[str, Any]]:
        """Format tools for Responses API."""
        formatted: list[dict[str, Any]] = []
        for tool in self.tools:
            ttype = tool.get("type", "function")
            if ttype == "function":
                func = tool.get("function", {})
                formatted.append(
                    {
                        "type": "function",
                        "name": func.get("name", ""),
                        "description": func.get("description", ""),
                        "parameters": func.get("parameters", {}),
                    }
                )
            elif ttype == "web_search":
                formatted.append({"type": "web_search_preview"})
            elif ttype == "file_search":
                fs = tool.get("file_search", {})
                formatted.append(
                    {
                        "type": "file_search",
                        "vector_store_ids": fs.get("vector_store_ids", []),
                    }
                )
            elif ttype == "mcp":
                mcp_config = tool.get("mcp", {})
                formatted.append(
                    {
                        "type": "mcp",
                        "server_label": mcp_config.get("server_label", ""),
                        "server_url": mcp_config.get("server_url", ""),
                        "require_approval": mcp_config.get(
                            "require_approval", "never"
                        ),
                    }
                )
        return formatted

    def _format_tools_chat(self) -> list[dict[str, Any]]:
        """Format tools for Chat Completions API."""
        formatted: list[dict[str, Any]] = []
        for tool in self.tools:
            ttype = tool.get("type", "function")
            if ttype == "function":
                formatted.append(
                    {
                        "type": "function",
                        "function": tool.get("function", {}),
                    }
                )
        return formatted

    # ── Skills (Agent Skills standard) ────────────────────────────

    def _inject_skills_into_shell(self, prepared: dict[str, Any]) -> None:
        """Inject Agent Skills into the shell tool's environment.

        OpenAI skills go into: tools[shell].environment.skills
        If no shell tool exists, one is created automatically.
        """
        formatted_skills = self._format_skills_openai()
        if not formatted_skills:
            return

        tools = prepared.setdefault("tools", [])

        # Find existing shell tool
        shell_tool = None
        for tool in tools:
            if tool.get("type") in ("shell", "shell_exec"):
                shell_tool = tool
                break

        if shell_tool is None:
            # Create a shell tool to host the skills
            shell_tool = {
                "type": "shell",
                "environment": {"type": "container_auto"},
            }
            tools.append(shell_tool)

        # Inject skills into the shell environment
        env = shell_tool.setdefault("environment", {"type": "container_auto"})
        env["skills"] = formatted_skills

    def _format_skills_openai(self) -> list[dict[str, Any]]:
        """Format skills for OpenAI Responses API."""
        formatted: list[dict[str, Any]] = []
        for skill in self.skills:
            source = skill.get("source", "custom")

            if source == "inline":
                formatted.append(
                    {
                        "type": "inline",
                        "name": skill.get("name", skill["skill_id"]),
                        "description": skill.get("description", ""),
                        "source": {
                            "type": "base64",
                            "media_type": skill.get(
                                "media_type", "application/zip"
                            ),
                            "data": skill.get("data", ""),
                        },
                    }
                )
            else:
                entry: dict[str, Any] = {
                    "type": "skill_reference",
                    "skill_id": skill["skill_id"],
                }
                version = skill.get("version", "latest")
                entry["version"] = (
                    int(version)
                    if isinstance(version, str) and version.isdigit()
                    else version
                )
                formatted.append(entry)

        return formatted
