from __future__ import annotations

import json
import os
from types import SimpleNamespace
from typing import Any

import weave
from openai import OpenAI
from openai.types.responses import Response
from pydantic import Field, PrivateAttr

from wbal.object import WBALObject


class LM(WBALObject):
    """
    Base class for language models.

    Subclasses must implement the invoke() method to call the underlying LLM API.
    """

    def invoke(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        mcp_servers: list[dict[str, Any]] | None = None,
    ) -> Any:
        """
        Invoke the language model.

        Args:
            messages: List of message dicts with 'role' and 'content' keys.
            tools: Optional list of tool definitions in provider format.
            mcp_servers: Optional list of MCP server tool definitions.

        Returns:
            The LLM response object (format depends on provider).

        Raises:
            NotImplementedError: If not overridden by subclass.
        """
        raise NotImplementedError("Subclasses must implement invoke method")


class OpenAIResponsesLM(LM):
    """
    Generic OpenAI Responses API LM wrapper.

    Use this when you want to select arbitrary model names via config/YAML.

    Note: temperature is optional. Some models (especially reasoning models)
    reject the temperature parameter. Only set it if you need it.
    """

    model: str = "gpt-5-mini"
    temperature: float | None = None  # Optional - only sent if explicitly set
    reasoning: dict[str, Any] | None = None
    include: list[str] | None = None
    client: OpenAI = Field(
        default_factory=lambda: OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    )

    def observe(self) -> str:
        return f"OpenAIResponsesLM(model={self.model}, temperature={self.temperature})"

    @weave.op()
    def invoke(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        mcp_servers: list[dict[str, Any]] | None = None,
    ) -> Response:
        kwargs: dict[str, Any] = {
            "model": self.model,
            "input": messages,
        }
        # Only include temperature if explicitly set
        if self.temperature is not None:
            kwargs["temperature"] = self.temperature
        if self.reasoning is not None:
            kwargs["reasoning"] = self.reasoning
        if self.include is not None:
            kwargs["include"] = self.include
        if tools or mcp_servers:
            combined_tools = list(tools) if tools else []
            if mcp_servers:
                combined_tools.extend(mcp_servers)
            kwargs["tools"] = combined_tools
        return self.client.responses.create(**kwargs)


class ScriptedLM(LM):
    """
    Deterministic LM for tests/examples.

    The script is a list of steps. Supported step shapes:
      - {"exit": "<message>"}  -> emits a function_call to exit(exit_message=...)
    """

    script: list[dict[str, Any]] = Field(default_factory=list)
    _index: int = PrivateAttr(default=0)

    def observe(self) -> str:
        return f"ScriptedLM(steps={len(self.script)}, index={self._index})"

    def invoke(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        mcp_servers: list[dict[str, Any]] | None = None,
    ) -> Any:
        if self._index >= len(self.script):
            return SimpleNamespace(output=[], output_text="")

        step = self.script[self._index]
        self._index += 1

        if "exit" in step:
            msg = str(step.get("exit", ""))
            return SimpleNamespace(
                output=[
                    SimpleNamespace(
                        type="function_call",
                        name="exit",
                        arguments=json.dumps({"exit_message": msg}),
                        call_id=f"call_{self._index}",
                    )
                ],
                output_text="",
            )

        return SimpleNamespace(output=[], output_text="")


class GPT5Large(LM):
    model: str = "gpt-5"
    include: list[str] = ["reasoning.encrypted_content"]
    temperature: float = 1.0
    client: OpenAI = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    def observe(self) -> str:
        """
        Return a concise description of the model configuration.
        """
        return f"GPT5Large(model={self.model}, temperature={self.temperature})"

    @weave.op()
    def invoke(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        mcp_servers: list[dict[str, Any]] | None = None,
    ) -> Response:
        """
        Invoke the language model.
        """
        kwargs: dict[str, Any] = {
            "model": self.model,
            "input": messages,
            "temperature": self.temperature,
            "include": self.include,
        }
        # Combine tools and mcp_servers without mutating input lists
        if tools or mcp_servers:
            combined_tools = list(tools) if tools else []
            if mcp_servers:
                combined_tools.extend(mcp_servers)
            kwargs["tools"] = combined_tools
        response: Response = self.client.responses.create(**kwargs)
        return response


class GPT5MiniTester(LM):
    model: str = "gpt-5-mini"
    client: OpenAI = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    reasoning: dict[str, Any] = {"effort": "minimal"}
    temperature: float = 1.0

    def observe(self) -> str:
        """
        Return a concise description of the model configuration.
        """
        return f"GPT5MiniTester(model={self.model}, temperature={self.temperature})"

    @weave.op()
    def invoke(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        mcp_servers: list[dict[str, Any]] | None = None,
    ) -> Response:
        """
        Invoke the language model.
        """
        kwargs: dict[str, Any] = {
            "model": self.model,
            "input": messages,
            "temperature": self.temperature,
            "reasoning": self.reasoning,
        }
        # Combine tools and mcp_servers without mutating input lists
        if tools or mcp_servers:
            combined_tools = list(tools) if tools else []
            if mcp_servers:
                combined_tools.extend(mcp_servers)
            kwargs["tools"] = combined_tools
        return self.client.responses.create(**kwargs)
