"""
OpenAI-compatible LLM client module
Provides support for OpenAI API and any OpenAI-compatible endpoints.
"""

import json
import logging
import os
from typing import Any

from .llm_interface import (
    ChatMessage,
    LLMFactory,
    LLMInterface,
    LLMResponse,
    MessageRole,
)


logger = logging.getLogger(__name__)


class OpenAICompatibleClient(LLMInterface):
    """OpenAI-compatible LLM client implementation"""

    def __init__(
        self,
        model: str = "gpt-4o",
        api_key: str | None = None,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> None:
        """
        Initialize OpenAI-compatible client
        Args:
            model: Model name (default: gpt-4o)
            api_key: API key (can be None for local endpoints)
            base_url: API base URL (default: https://api.openai.com/v1)
            **kwargs: Other configuration parameters
        Raises:
            ImportError: If openai package is not installed
            ValueError: If api_key is required but not provided
        """
        super().__init__(model=model, **kwargs)

        # Import openai package
        try:
            from openai import AsyncOpenAI

            self._AsyncOpenAI = AsyncOpenAI
        except ImportError:
            raise ImportError(
                "OpenAI package is not installed. Install it with: pip install openai"
            )

        # Get api_key from parameter or environment variable
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")

        # Set base_url
        if base_url is None:
            base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
        self.base_url = base_url

        # Check if api_key is required (not for local endpoints)
        self._is_local_endpoint = self._is_local_url(self.base_url)

        if not self.api_key and not self._is_local_endpoint:
            raise ValueError(
                "API key is required for OpenAI API. "
                "Set OPENAI_API_KEY environment variable or use --api-key option. "
                "For local endpoints (vLLM, LM Studio), api_key is optional."
            )

        # Initialize client
        client_kwargs: dict[str, Any] = {
            "api_key": self.api_key,
            "base_url": self.base_url,
        }

        # Add timeout if provided
        if "timeout" in kwargs:
            client_kwargs["timeout"] = kwargs["timeout"]

        self.client = self._AsyncOpenAI(**client_kwargs)

        logger.info(
            f"OpenAI-compatible client initialized: model={model}, "
            f"base_url={self.base_url}, local={self._is_local_endpoint}"
        )

    def _is_local_url(self, url: str) -> bool:
        """Check if the URL is a local endpoint"""
        local_indicators = [
            "localhost",
            "127.0.0.1",
            "0.0.0.0",
            "::1",
        ]
        return any(indicator in url.lower() for indicator in local_indicators)

    def _convert_message_role(self, role: MessageRole) -> str:
        """Convert message role to OpenAI format"""
        role_mapping = {
            MessageRole.SYSTEM: "system",
            MessageRole.USER: "user",
            MessageRole.ASSISTANT: "assistant",
        }
        return role_mapping.get(role, "user")

    async def chat(self, messages: list[ChatMessage]) -> LLMResponse:
        """
        Chat interface
        Args:
            messages: List of chat messages
        Returns:
            LLM response
        Raises:
            Exception: If API call fails
        """
        try:
            # Import OpenAI types for proper type hints
            from openai.types.chat import ChatCompletionMessageParam

            # Convert messages to OpenAI format
            # Type: ignore because we're constructing valid message dicts
            openai_messages: list[ChatCompletionMessageParam] = [  # type: ignore
                {
                    "role": self._convert_message_role(msg.role),
                    "content": msg.content or "",  # Handle None content
                }
                for msg in messages
            ]

            # Call OpenAI API
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=openai_messages,
            )

            # Extract response content
            content = response.choices[0].message.content or ""

            # Build usage info
            usage = None
            if hasattr(response, "usage") and response.usage:
                usage = {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                }

            return LLMResponse(
                content=content,
                metadata={
                    "model": response.model,
                    "provider": "openai-compatible",
                    "base_url": self.base_url,
                    "finish_reason": response.choices[0].finish_reason,
                },
                usage=usage,
                model=self.model,
            )
        except Exception as e:
            logger.error(f"OpenAI-compatible chat failed: {e}")
            raise

    async def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
        """
        Text completion interface
        Args:
            prompt: Input prompt
            **kwargs: Additional parameters (not used, kept for interface compatibility)
        Returns:
            LLM response
        Raises:
            Exception: If API call fails
        """
        _ = kwargs  # Reserved for future use
        try:
            # Use chat interface for completion
            messages = [ChatMessage(role=MessageRole.USER, content=prompt)]
            return await self.chat(messages)
        except Exception as e:
            logger.error(f"OpenAI-compatible completion failed: {e}")
            raise

    async def extract_structured(
        self, content: str, schema: dict[str, Any], **kwargs: Any
    ) -> dict[str, Any]:
        """
        Structured data extraction interface
        Args:
            content: Input content
            schema: Output schema definition
            **kwargs: Additional parameters
        Returns:
            Extracted structured data
        Raises:
            Exception: If extraction fails
        """
        try:
            # Build extraction prompt
            prompt = self._build_extraction_prompt(content, schema)

            # Call completion interface
            response = await self.complete(prompt, **kwargs)

            # Try to parse as JSON
            response_text = response.content.strip()

            # Try to extract JSON from markdown code blocks
            if "```json" in response_text:
                start = response_text.find("```json") + 7
                end = response_text.rfind("```")
                if end > start:
                    response_text = response_text[start:end].strip()
            elif "```" in response_text:
                start = response_text.find("```") + 3
                end = response_text.rfind("```")
                if end > start:
                    response_text = response_text[start:end].strip()

            try:
                extracted_data = json.loads(response_text)
                return {
                    "extracted_data": extracted_data,
                    "schema": schema,
                    "model": self.model,
                    "provider": "openai-compatible",
                }
            except json.JSONDecodeError:
                # If not valid JSON, return raw content with error
                return {
                    "extracted_data": response_text,
                    "schema": schema,
                    "model": self.model,
                    "provider": "openai-compatible",
                    "parse_error": "Could not parse as JSON",
                    "raw_response": response_text[:500],  # First 500 chars
                }
        except Exception as e:
            logger.error(f"OpenAI-compatible structured extraction failed: {e}")
            raise

    def _build_extraction_prompt(self, content: str, schema: dict[str, Any]) -> str:
        """Build extraction prompt for structured data"""
        prompt = f"""Extract structured information from the following content according to the schema.

Content:
{content}

Schema:
{schema}

Please return the extracted data in valid JSON format. If a field is not found, use null.
Return ONLY the JSON object without any additional text or formatting."""
        return prompt

    def get_model_info(self) -> dict[str, Any]:
        """Get model information"""
        info = super().get_model_info()
        info.update(
            {
                "provider": "OpenAI-compatible",
                "base_url": self.base_url,
                "local_endpoint": self._is_local_endpoint,
                "api_key_configured": bool(self.api_key),
            }
        )
        return info

    async def test_connection(self) -> bool:
        """
        Test API connection
        Returns:
            True if connection successful, False otherwise
        """
        try:
            test_message = ChatMessage(
                role=MessageRole.USER, content="Hello, this is a connection test."
            )
            response = await self.chat([test_message])
            return bool(response.content)
        except Exception as e:
            logger.error(f"OpenAI-compatible connection test failed: {e}")
            return False

    async def __aenter__(self):
        """Async context manager entry"""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        _ = exc_type, exc_val, exc_tb
        await self.client.close()


# Register OpenAI-compatible provider
LLMFactory.register("openai", OpenAICompatibleClient)


# Convenience function
def create_openai_client(
    model: str = "gpt-4o",
    api_key: str | None = None,
    base_url: str | None = None,
    **kwargs: Any,
) -> OpenAICompatibleClient:
    """
    Create OpenAI-compatible client
    Args:
        model: Model name
        api_key: API key (optional for local endpoints)
        base_url: API base URL
        **kwargs: Configuration parameters
    Returns:
        OpenAI-compatible client instance
    """
    return OpenAICompatibleClient(
        model=model, api_key=api_key, base_url=base_url, **kwargs
    )
