"""LLM provider adapters — OpenAI, Anthropic, LiteLLM."""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple

from .types import Message, Role, ToolDef


class LLMProvider(ABC):
    """Abstract LLM provider."""

    @abstractmethod
    def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        """
        Send a chat completion request.

        Returns:
            (content, tool_calls, usage)
            - content: text response (may be empty if tool calls)
            - tool_calls: list of {"id": ..., "name": ..., "arguments": {...}}
            - usage: {"prompt_tokens": N, "completion_tokens": N}
        """
        ...


class OpenAIProvider(LLMProvider):
    """OpenAI / OpenAI-compatible provider."""

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None) -> None:
        try:
            import openai
        except ImportError:
            raise ImportError(
                "openai package required. Install with: pip install novyx-agent[openai]"
            )
        kwargs: Dict[str, Any] = {}
        if api_key:
            kwargs["api_key"] = api_key
        if base_url:
            kwargs["base_url"] = base_url
        self._client = openai.OpenAI(**kwargs)

    def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        oai_messages = []
        for m in messages:
            msg: Dict[str, Any] = {"role": m.role.value, "content": m.content}
            if m.name:
                msg["name"] = m.name
            if m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            if m.tool_calls:
                msg["tool_calls"] = m.tool_calls
            oai_messages.append(msg)

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": oai_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if tools:
            kwargs["tools"] = [t.to_openai_schema() for t in tools]

        response = self._client.chat.completions.create(**kwargs)
        choice = response.choices[0]

        content = choice.message.content or ""
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "arguments": json.loads(tc.function.arguments),
                })

        usage = {}
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
            }

        return content, tool_calls, usage


class AnthropicProvider(LLMProvider):
    """Anthropic Claude provider."""

    def __init__(self, api_key: Optional[str] = None) -> None:
        try:
            import anthropic
        except ImportError:
            raise ImportError(
                "anthropic package required. Install with: pip install novyx-agent[anthropic]"
            )
        kwargs: Dict[str, Any] = {}
        if api_key:
            kwargs["api_key"] = api_key
        self._client = anthropic.Anthropic(**kwargs)

    def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        system_text = ""
        claude_messages = []
        for m in messages:
            if m.role == Role.SYSTEM:
                system_text += m.content + "\n"
            elif m.role == Role.TOOL:
                claude_messages.append({
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": m.tool_call_id or "",
                            "content": m.content,
                        }
                    ],
                })
            elif m.role == Role.ASSISTANT and m.tool_calls:
                content_blocks: List[Dict[str, Any]] = []
                if m.content:
                    content_blocks.append({"type": "text", "text": m.content})
                for tc in m.tool_calls:
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc["id"],
                        "name": tc["name"],
                        "input": tc["arguments"],
                    })
                claude_messages.append({"role": "assistant", "content": content_blocks})
            else:
                claude_messages.append({
                    "role": m.role.value if m.role != Role.USER else "user",
                    "content": m.content,
                })

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": claude_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if system_text.strip():
            kwargs["system"] = system_text.strip()
        if tools:
            kwargs["tools"] = [t.to_anthropic_schema() for t in tools]

        response = self._client.messages.create(**kwargs)

        content = ""
        tool_calls = []
        for block in response.content:
            if block.type == "text":
                content += block.text
            elif block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "arguments": block.input,
                })

        usage = {
            "prompt_tokens": response.usage.input_tokens,
            "completion_tokens": response.usage.output_tokens,
        }

        return content, tool_calls, usage


class LiteLLMProvider(LLMProvider):
    """LiteLLM universal provider — route to any model."""

    def __init__(self) -> None:
        try:
            import litellm  # noqa: F401
        except ImportError:
            raise ImportError(
                "litellm package required. Install with: pip install novyx-agent[litellm]"
            )

    def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        import litellm

        oai_messages = []
        for m in messages:
            msg: Dict[str, Any] = {"role": m.role.value, "content": m.content}
            if m.name:
                msg["name"] = m.name
            if m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            if m.tool_calls:
                msg["tool_calls"] = m.tool_calls
            oai_messages.append(msg)

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": oai_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if tools:
            kwargs["tools"] = [t.to_openai_schema() for t in tools]

        response = litellm.completion(**kwargs)
        choice = response.choices[0]

        content = choice.message.content or ""
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "arguments": json.loads(tc.function.arguments)
                    if isinstance(tc.function.arguments, str)
                    else tc.function.arguments,
                })

        usage = {}
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
            }

        return content, tool_calls, usage


def get_provider(provider: str, **kwargs: Any) -> LLMProvider:
    """Factory — resolve provider by name."""
    providers = {
        "openai": OpenAIProvider,
        "anthropic": AnthropicProvider,
        "litellm": LiteLLMProvider,
    }
    cls = providers.get(provider)
    if cls is None:
        raise ValueError(
            f"Unknown provider '{provider}'. Available: {', '.join(providers)}"
        )
    return cls(**kwargs)


# ---------------------------------------------------------------------------
# Async providers
# ---------------------------------------------------------------------------


class AsyncLLMProvider(ABC):
    """Abstract async LLM provider."""

    @abstractmethod
    async def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        """
        Send an async chat completion request.

        Returns:
            (content, tool_calls, usage)
            - content: text response (may be empty if tool calls)
            - tool_calls: list of {"id": ..., "name": ..., "arguments": {...}}
            - usage: {"prompt_tokens": N, "completion_tokens": N}
        """
        ...


class AsyncOpenAIProvider(AsyncLLMProvider):
    """Async OpenAI / OpenAI-compatible provider."""

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None) -> None:
        try:
            import openai
        except ImportError:
            raise ImportError(
                "openai package required. Install with: pip install novyx-agent[openai]"
            )
        kwargs: Dict[str, Any] = {}
        if api_key:
            kwargs["api_key"] = api_key
        if base_url:
            kwargs["base_url"] = base_url
        self._client = openai.AsyncOpenAI(**kwargs)

    async def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        oai_messages = []
        for m in messages:
            msg: Dict[str, Any] = {"role": m.role.value, "content": m.content}
            if m.name:
                msg["name"] = m.name
            if m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            if m.tool_calls:
                msg["tool_calls"] = m.tool_calls
            oai_messages.append(msg)

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": oai_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if tools:
            kwargs["tools"] = [t.to_openai_schema() for t in tools]

        response = await self._client.chat.completions.create(**kwargs)
        choice = response.choices[0]

        content = choice.message.content or ""
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "arguments": json.loads(tc.function.arguments),
                })

        usage = {}
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
            }

        return content, tool_calls, usage


class AsyncAnthropicProvider(AsyncLLMProvider):
    """Async Anthropic Claude provider."""

    def __init__(self, api_key: Optional[str] = None) -> None:
        try:
            import anthropic
        except ImportError:
            raise ImportError(
                "anthropic package required. Install with: pip install novyx-agent[anthropic]"
            )
        kwargs: Dict[str, Any] = {}
        if api_key:
            kwargs["api_key"] = api_key
        self._client = anthropic.AsyncAnthropic(**kwargs)

    async def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        system_text = ""
        claude_messages = []
        for m in messages:
            if m.role == Role.SYSTEM:
                system_text += m.content + "\n"
            elif m.role == Role.TOOL:
                claude_messages.append({
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": m.tool_call_id or "",
                            "content": m.content,
                        }
                    ],
                })
            elif m.role == Role.ASSISTANT and m.tool_calls:
                content_blocks: List[Dict[str, Any]] = []
                if m.content:
                    content_blocks.append({"type": "text", "text": m.content})
                for tc in m.tool_calls:
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc["id"],
                        "name": tc["name"],
                        "input": tc["arguments"],
                    })
                claude_messages.append({"role": "assistant", "content": content_blocks})
            else:
                claude_messages.append({
                    "role": m.role.value if m.role != Role.USER else "user",
                    "content": m.content,
                })

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": claude_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if system_text.strip():
            kwargs["system"] = system_text.strip()
        if tools:
            kwargs["tools"] = [t.to_anthropic_schema() for t in tools]

        response = await self._client.messages.create(**kwargs)

        content = ""
        tool_calls = []
        for block in response.content:
            if block.type == "text":
                content += block.text
            elif block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "arguments": block.input,
                })

        usage = {
            "prompt_tokens": response.usage.input_tokens,
            "completion_tokens": response.usage.output_tokens,
        }

        return content, tool_calls, usage


class AsyncLiteLLMProvider(AsyncLLMProvider):
    """Async LiteLLM universal provider — route to any model."""

    def __init__(self) -> None:
        try:
            import litellm  # noqa: F401
        except ImportError:
            raise ImportError(
                "litellm package required. Install with: pip install novyx-agent[litellm]"
            )

    async def chat(
        self,
        messages: List[Message],
        tools: List[ToolDef],
        *,
        model: str,
        temperature: float,
        max_tokens: int,
    ) -> Tuple[str, List[Dict[str, Any]], Dict[str, int]]:
        import litellm

        oai_messages = []
        for m in messages:
            msg: Dict[str, Any] = {"role": m.role.value, "content": m.content}
            if m.name:
                msg["name"] = m.name
            if m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            if m.tool_calls:
                msg["tool_calls"] = m.tool_calls
            oai_messages.append(msg)

        kwargs: Dict[str, Any] = {
            "model": model,
            "messages": oai_messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if tools:
            kwargs["tools"] = [t.to_openai_schema() for t in tools]

        response = await litellm.acompletion(**kwargs)
        choice = response.choices[0]

        content = choice.message.content or ""
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append({
                    "id": tc.id,
                    "name": tc.function.name,
                    "arguments": json.loads(tc.function.arguments)
                    if isinstance(tc.function.arguments, str)
                    else tc.function.arguments,
                })

        usage = {}
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
            }

        return content, tool_calls, usage


def get_async_provider(provider: str, **kwargs: Any) -> AsyncLLMProvider:
    """Factory — resolve async provider by name."""
    providers = {
        "openai": AsyncOpenAIProvider,
        "anthropic": AsyncAnthropicProvider,
        "litellm": AsyncLiteLLMProvider,
    }
    cls = providers.get(provider)
    if cls is None:
        raise ValueError(
            f"Unknown provider '{provider}'. Available: {', '.join(providers)}"
        )
    return cls(**kwargs)
