"""Abstract LLM client for unified AI provider integration."""

import asyncio
import json
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import Any

from ..logger import get_logger
from .pricing_manager import PricingManager

logger = get_logger("llm_client")


class LLMProvider(str, Enum):
    """Supported LLM providers."""

    OPENAI = "openai"
    ANTHROPIC = "anthropic"


@dataclass
class LLMResponse:
    """Standardized response from LLM providers."""

    content: str
    model: str
    usage: dict[str, int]  # tokens used, etc.
    raw_response: Any | None = None
    cost_breakdown: dict[str, Any] | None = None  # cost information from PricingManager


class LLMClientError(Exception):
    """Base exception for LLM client errors."""

    pass


class LLMAPIError(LLMClientError):
    """Exception for API-related errors."""

    pass


class LLMRateLimitError(LLMClientError):
    """Exception for rate limit errors."""

    pass


class LLMClient(ABC):
    """Abstract base class for LLM clients."""

    def __init__(self, api_key: str, model: str | None = None):
        """Initialize LLM client.

        Args:
            api_key: API key for the provider
            model: Model to use (provider-specific)
        """
        self.api_key = api_key
        self.model = model or self.get_default_model()
        self.pricing_manager = PricingManager()
        logger.info(f"Initializing {self.__class__.__name__} with model: {self.model}")
        logger.debug(
            f"Pricing manager initialized with {len(self.pricing_manager.get_supported_models())} supported models"
        )

    @abstractmethod
    def get_default_model(self) -> str:
        """Get the default model for this provider."""
        pass

    @abstractmethod
    async def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.1,
        max_tokens: int = 4000,
        response_format: str = "json",
    ) -> LLMResponse:
        """Make a completion request to the LLM.

        Args:
            system_prompt: System prompt for the LLM
            user_prompt: User prompt for the LLM
            temperature: Temperature for sampling (0.0 = deterministic)
            max_tokens: Maximum tokens in response
            response_format: Expected response format ("json" or "text")

        Returns:
            LLMResponse with the completion

        Raises:
            LLMAPIError: If the API request fails
            LLMRateLimitError: If rate limited
        """
        pass

    def validate_json_response(self, response_text: str) -> dict[str, Any]:
        """Validate and parse JSON response.

        Args:
            response_text: Raw response text

        Returns:
            Parsed JSON dictionary

        Raises:
            LLMClientError: If response is not valid JSON
        """
        try:
            # Handle potential markdown code blocks
            if response_text.strip().startswith("```"):
                lines = response_text.strip().split("\n")
                # Remove first and last lines if they're code block markers
                if lines[0].startswith("```") and lines[-1] == "```":
                    response_text = "\n".join(lines[1:-1])

            return json.loads(response_text)
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON response: {e}")
            logger.debug(f"Response text: {response_text[:500]}...")
            raise LLMClientError(f"Invalid JSON response: {e}")

    async def complete_with_retry(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.1,
        max_tokens: int = 4000,
        response_format: str = "json",
        max_retries: int = 3,
        retry_delay: float = 1.0,
    ) -> LLMResponse:
        """Complete with automatic retry on failure.

        Args:
            system_prompt: System prompt for the LLM
            user_prompt: User prompt for the LLM
            temperature: Temperature for sampling
            max_tokens: Maximum tokens in response
            response_format: Expected response format
            max_retries: Maximum number of retry attempts
            retry_delay: Base delay between retries (exponential backoff)

        Returns:
            LLMResponse with the completion

        Raises:
            LLMAPIError: If all retries fail
        """
        logger.debug(
            f"Starting completion with retry - max_retries: {max_retries}, "
            f"retry_delay: {retry_delay}s"
        )

        last_error = None
        start_time = asyncio.get_event_loop().time()

        for attempt in range(max_retries):
            try:
                logger.debug(f"Attempt {attempt + 1}/{max_retries} for completion")
                result = await self.complete(
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    response_format=response_format,
                )

                elapsed_time = asyncio.get_event_loop().time() - start_time
                if attempt > 0:
                    logger.info(
                        f"Completion successful on attempt {attempt + 1}, "
                        f"total time: {elapsed_time:.2f}s"
                    )
                return result

            except LLMRateLimitError as e:
                last_error = e
                if attempt < max_retries - 1:
                    delay = retry_delay * (2**attempt)  # Exponential backoff
                    logger.warning(
                        f"Rate limited on attempt {attempt + 1}/{max_retries}, "
                        f"retrying in {delay}s: {e}"
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(
                        f"Rate limit exceeded after {max_retries} attempts: {e}"
                    )
                    raise

            except LLMAPIError as e:
                last_error = e
                if attempt < max_retries - 1:
                    delay = retry_delay * (2**attempt)
                    logger.warning(
                        f"API error on attempt {attempt + 1}/{max_retries}, "
                        f"retrying in {delay}s: {e}"
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"API error after {max_retries} attempts: {e}")
                    raise

            except LLMClientError as e:
                # Client errors are typically not retryable
                logger.error(
                    f"Client error on attempt {attempt + 1}, not retrying: {e}"
                )
                raise

        # Should not reach here, but just in case
        total_time = asyncio.get_event_loop().time() - start_time
        logger.error(f"All {max_retries} attempts failed in {total_time:.2f}s")
        raise LLMAPIError(f"Failed after {max_retries} attempts: {last_error}")


class OpenAIClient(LLMClient):
    """OpenAI API client implementation."""

    def get_default_model(self) -> str:
        """Get the default OpenAI model."""
        return "gpt-4-turbo-preview"

    async def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.1,
        max_tokens: int = 4000,
        response_format: str = "json",
    ) -> LLMResponse:
        """Make a completion request to OpenAI."""
        try:
            import openai
        except ImportError:
            raise LLMClientError(
                "OpenAI client library not installed. Run: uv pip install openai"
            )

        logger.info(f"Making OpenAI completion request with model: {self.model}")
        logger.debug(
            f"Request parameters - Temperature: {temperature}, Max tokens: {max_tokens}, Format: {response_format}"
        )
        logger.debug(
            f"System prompt length: {len(system_prompt)} chars, User prompt length: {len(user_prompt)} chars"
        )

        # Estimate token usage
        estimated_input_tokens = (len(system_prompt) + len(user_prompt)) // 4
        logger.debug(f"Estimated input tokens: ~{estimated_input_tokens}")

        client = openai.AsyncOpenAI(api_key=self.api_key)

        try:
            # Prepare messages
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ]

            # Make the request
            kwargs = {
                "model": self.model,
                "messages": messages,
                "temperature": temperature,
                "max_tokens": max_tokens,
            }

            # Add response format if JSON expected
            if response_format == "json":
                kwargs["response_format"] = {"type": "json_object"}

            response = await client.chat.completions.create(**kwargs)  # type: ignore[call-overload]

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

            # Extract usage information
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

            logger.info(
                f"OpenAI completion successful. Tokens used: {usage['total_tokens']} "
                f"(prompt: {usage['prompt_tokens']}, completion: {usage['completion_tokens']})"
            )
            logger.debug(
                f"Response model: {response.model}, Content length: {len(content)} chars"
            )

            # Calculate cost using pricing manager
            cost_breakdown = self.pricing_manager.calculate_cost(
                model_name=response.model,
                prompt_tokens=usage["prompt_tokens"],
                completion_tokens=usage["completion_tokens"],
            )
            logger.debug(
                f"Cost calculation: ${cost_breakdown['total_cost']:.6f} "
                f"{cost_breakdown['currency']} "
                f"(pricing: {cost_breakdown['pricing_source']})"
            )

            return LLMResponse(
                content=content,
                model=response.model,
                usage=usage,
                raw_response=response,
                cost_breakdown=cost_breakdown,
            )

        except openai.RateLimitError as e:
            logger.error(f"OpenAI rate limit error: {e}")
            raise LLMRateLimitError(f"OpenAI rate limit exceeded: {e}")
        except openai.APIError as e:
            logger.error(f"OpenAI API error: {e}")
            raise LLMAPIError(f"OpenAI API error: {e}")
        except Exception as e:
            logger.error(f"Unexpected OpenAI error: {e}")
            raise LLMClientError(f"Unexpected error in OpenAI client: {e}")


class AnthropicClient(LLMClient):
    """Anthropic API client implementation."""

    def get_default_model(self) -> str:
        """Get the default Anthropic model."""
        return "claude-3-5-sonnet-20241022"

    async def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.1,
        max_tokens: int = 4000,
        response_format: str = "json",
    ) -> LLMResponse:
        """Make a completion request to Anthropic."""
        try:
            import anthropic
        except ImportError:
            raise LLMClientError(
                "Anthropic client library not installed. Run: pip install anthropic"
            )

        logger.info(f"Making Anthropic completion request with model: {self.model}")
        logger.debug(
            f"Request parameters - Temperature: {temperature}, Max tokens: {max_tokens}, Format: {response_format}"
        )
        logger.debug(
            f"System prompt length: {len(system_prompt)} chars, User prompt length: {len(user_prompt)} chars"
        )

        # Estimate token usage
        estimated_input_tokens = (len(system_prompt) + len(user_prompt)) // 4
        logger.debug(f"Estimated input tokens: ~{estimated_input_tokens}")

        client = anthropic.AsyncAnthropic(api_key=self.api_key)

        try:
            # Combine prompts for Anthropic format
            # Add JSON instruction if needed
            full_user_prompt = user_prompt
            if response_format == "json":
                full_user_prompt += "\n\nPlease respond with valid JSON only."

            # Make the request
            response = await client.messages.create(
                model=self.model,
                system=system_prompt,
                messages=[{"role": "user", "content": full_user_prompt}],
                temperature=temperature,
                max_tokens=max_tokens,
            )

            # Extract content
            content = response.content[0].text

            # Calculate usage
            usage = {
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens
                + response.usage.output_tokens,
            }

            logger.info(
                f"Anthropic completion successful. Tokens used: {usage['total_tokens']} "
                f"(input: {usage['prompt_tokens']}, output: {usage['completion_tokens']})"
            )
            logger.debug(
                f"Response model: {response.model}, Content length: {len(content)} chars"
            )

            # Calculate cost using pricing manager
            cost_breakdown = self.pricing_manager.calculate_cost(
                model_name=response.model,
                prompt_tokens=usage["prompt_tokens"],
                completion_tokens=usage["completion_tokens"],
            )
            logger.debug(
                f"Cost calculation: ${cost_breakdown['total_cost']:.6f} "
                f"{cost_breakdown['currency']} "
                f"(pricing: {cost_breakdown['pricing_source']})"
            )

            return LLMResponse(
                content=content,
                model=response.model,
                usage=usage,
                raw_response=response,
                cost_breakdown=cost_breakdown,
            )

        except anthropic.RateLimitError as e:
            logger.error(f"Anthropic rate limit error: {e}")
            raise LLMRateLimitError(f"Anthropic rate limit exceeded: {e}")
        except anthropic.APIError as e:
            logger.error(f"Anthropic API error: {e}")
            raise LLMAPIError(f"Anthropic API error: {e}")
        except Exception as e:
            logger.error(f"Unexpected Anthropic error: {e}")
            raise LLMClientError(f"Unexpected error in Anthropic client: {e}")


def create_llm_client(
    provider: LLMProvider, api_key: str, model: str | None = None
) -> LLMClient:
    """Factory function to create LLM client based on provider.

    Args:
        provider: LLM provider to use
        api_key: API key for the provider
        model: Optional model override

    Returns:
        Configured LLM client instance

    Raises:
        ValueError: If provider is not supported
    """
    logger.info(f"Creating LLM client for provider: {provider}")
    logger.debug(
        f"Model: {model or 'default'}, API key length: {len(api_key) if api_key else 0}"
    )

    try:
        if provider == LLMProvider.OPENAI:
            client = OpenAIClient(api_key=api_key, model=model)
            logger.info(
                f"Successfully created OpenAI client with model: {client.model}"
            )
            return client
        elif provider == LLMProvider.ANTHROPIC:
            client = AnthropicClient(api_key=api_key, model=model)
            logger.info(
                f"Successfully created Anthropic client with model: {client.model}"
            )
            return client
        else:
            logger.error(f"Unsupported LLM provider: {provider}")
            raise ValueError(f"Unsupported LLM provider: {provider}")
    except Exception as e:
        logger.error(f"Failed to create LLM client for {provider}: {e}")
        raise
