"""
LLM Agent for code generation using LangChain.

Provides a clean abstraction over different LLM providers (OpenAI, Anthropic)
with built-in retry logic, error correction, and production-grade features.
"""

from typing import Optional, Dict, Any, Literal
from functools import wraps
import time
import logging
import os

logger = logging.getLogger(__name__)


def retry_on_llm_error(max_retries: int = 3, delay: float = 1.0):
    """
    Decorator to retry LLM calls on failure with exponential backoff.

    Args:
        max_retries: Maximum number of retry attempts
        delay: Initial delay in seconds (doubles with each retry)
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        wait_time = delay * (2 ** attempt)
                        logger.warning(
                            f"LLM call failed (attempt {attempt + 1}/{max_retries}): {e}. "
                            f"Retrying in {wait_time}s..."
                        )
                        time.sleep(wait_time)
                    else:
                        logger.error(f"LLM call failed after {max_retries} attempts: {e}")
            raise last_exception
        return wrapper
    return decorator


class LLMAgent:
    """
    LLM agent for generating code from natural language using LangChain.

    Supports multiple LLM providers with consistent interface:
    - OpenAI (GPT-4, GPT-3.5-turbo)
    - Anthropic (Claude-3-Sonnet, Claude-3-Opus)

    Features:
    - Automatic retry with exponential backoff
    - Error correction capabilities
    - Token usage tracking
    - Verbose mode for debugging

    Example:
        >>> agent = LLMAgent(model="gpt-4", temperature=0.1)
        >>> code = agent.generate_code(
        ...     query="What is the conversion rate?",
        ...     schema="user_id, event, timestamp",
        ...     examples="...",
        ...     glossary="conversion: users who purchased"
        ... )
    """

    SUPPORTED_PROVIDERS = {
        "gpt-4": "openai",
        "gpt-3.5-turbo": "openai",
        "gpt-4-turbo": "openai",
        "claude-3-sonnet": "anthropic",
        "claude-3-opus": "anthropic",
        "claude-3-haiku": "anthropic",
    }

    def __init__(
        self,
        model: str = "gpt-4",
        temperature: float = 0.1,
        provider: Optional[Literal["openai", "anthropic"]] = None,
        max_tokens: Optional[int] = None,
        verbose: bool = False,
    ):
        """
        Initialize LLM agent.

        Args:
            model: Model name (gpt-4, gpt-3.5-turbo, claude-3-sonnet, etc.)
            temperature: Sampling temperature (0-1). Lower = more deterministic
            provider: LLM provider (auto-detected if None)
            max_tokens: Maximum tokens in response (provider default if None)
            verbose: Enable verbose logging

        Raises:
            ValueError: If model is unsupported or API key is missing
            ImportError: If required LangChain package is not installed
        """
        self.model = model
        self.temperature = temperature
        self.verbose = verbose
        self.max_tokens = max_tokens

        # Auto-detect provider if not specified
        if provider is None:
            provider = self._detect_provider(model)
        self.provider = provider

        # Validate temperature
        if not 0 <= temperature <= 1:
            raise ValueError(f"Temperature must be between 0 and 1, got {temperature}")

        # Initialize LLM
        self.llm = self._initialize_llm()

        # Track usage
        self.total_tokens_used = 0
        self.calls_made = 0

        if verbose:
            logger.setLevel(logging.DEBUG)
            logger.debug(f"Initialized LLMAgent with {model} (provider: {provider})")

    def _detect_provider(self, model: str) -> str:
        """Auto-detect provider from model name."""
        if model in self.SUPPORTED_PROVIDERS:
            return self.SUPPORTED_PROVIDERS[model]

        # Fallback: guess by model name prefix
        if model.startswith("gpt"):
            return "openai"
        elif model.startswith("claude"):
            return "anthropic"
        else:
            raise ValueError(
                f"Unknown model: {model}. Supported models: {list(self.SUPPORTED_PROVIDERS.keys())}"
            )

    def _initialize_llm(self):
        """Initialize LLM client based on provider."""
        if self.provider == "openai":
            return self._initialize_openai()
        elif self.provider == "anthropic":
            return self._initialize_anthropic()
        else:
            raise ValueError(f"Unknown provider: {self.provider}")

    def _initialize_openai(self):
        """Initialize OpenAI LLM client."""
        try:
            from langchain_openai import ChatOpenAI
        except ImportError:
            raise ImportError(
                "langchain-openai required for OpenAI models. "
                "Install with: pip install langchain-openai"
            )

        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError(
                "OPENAI_API_KEY environment variable not set. "
                "Add it to your .env file or set it in your environment."
            )

        kwargs = {
            "model": self.model,
            "temperature": self.temperature,
            "api_key": api_key,
        }

        if self.max_tokens:
            kwargs["max_tokens"] = self.max_tokens

        return ChatOpenAI(**kwargs)

    def _initialize_anthropic(self):
        """Initialize Anthropic LLM client."""
        try:
            from langchain_anthropic import ChatAnthropic
        except ImportError:
            raise ImportError(
                "langchain-anthropic required for Anthropic models. "
                "Install with: pip install langchain-anthropic"
            )

        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            raise ValueError(
                "ANTHROPIC_API_KEY environment variable not set. "
                "Add it to your .env file or set it in your environment."
            )

        # Map friendly names to actual model IDs
        model_map = {
            "claude-3-sonnet": "claude-3-sonnet-20240229",
            "claude-3-opus": "claude-3-opus-20240229",
            "claude-3-haiku": "claude-3-haiku-20240307",
        }

        model_id = model_map.get(self.model, self.model)

        kwargs = {
            "model": model_id,
            "temperature": self.temperature,
            "api_key": api_key,
        }

        if self.max_tokens:
            kwargs["max_tokens"] = self.max_tokens

        return ChatAnthropic(**kwargs)

    @retry_on_llm_error(max_retries=3, delay=1.0)
    def generate_code(
        self,
        query: str,
        schema: str,
        examples: str,
        glossary: str,
        system_prompt: Optional[str] = None,
        custom_instructions: Optional[str] = None,
    ) -> str:
        """
        Generate Python code from natural language query.

        Args:
            query: User's natural language question
            schema: Schema description (columns, data types, stats)
            examples: Similar example queries with code
            glossary: Business term definitions
            system_prompt: Custom system prompt (uses default if None)
            custom_instructions: Additional instructions to append

        Returns:
            Generated Python code as string

        Raises:
            ValueError: If inputs are invalid
            RuntimeError: If LLM call fails after retries
        """
        # Validate inputs
        if not query or not query.strip():
            raise ValueError("query cannot be empty")

        if self.verbose:
            logger.debug(f"Generating code for query: {query[:100]}...")

        # Use default system prompt if not provided
        if system_prompt is None:
            from .prompt_templates import PromptTemplates
            system_prompt = PromptTemplates.get_system_prompt()

        # Create user prompt
        from .prompt_templates import PromptTemplates
        user_prompt = PromptTemplates.get_user_prompt(
            query=query,
            schema=schema,
            examples=examples,
            glossary=glossary
        )

        # Add custom instructions if provided
        if custom_instructions:
            user_prompt += f"\n\n**Additional Instructions:**\n{custom_instructions}"

        # Create chat messages
        from langchain.prompts import ChatPromptTemplate

        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            ("user", user_prompt)
        ])

        # Generate code
        chain = prompt | self.llm
        response = chain.invoke({})

        # Track usage
        self.calls_made += 1
        if hasattr(response, 'response_metadata'):
            usage = response.response_metadata.get('token_usage', {})
            tokens = usage.get('total_tokens', 0)
            self.total_tokens_used += tokens
            if self.verbose:
                logger.debug(f"Tokens used: {tokens} (total: {self.total_tokens_used})")

        # Extract code from response
        code = self._extract_code(response.content)

        if self.verbose:
            logger.debug(f"Generated code ({len(code)} chars)")

        return code

    @retry_on_llm_error(max_retries=2, delay=0.5)
    def correct_error(
        self,
        query: str,
        code: str,
        error: str,
        attempt: int = 1
    ) -> str:
        """
        Correct code that failed execution.

        Args:
            query: Original natural language query
            code: Failed Python code
            error: Error message from execution
            attempt: Current correction attempt number

        Returns:
            Corrected Python code

        Raises:
            RuntimeError: If correction fails after retries
        """
        if self.verbose:
            logger.debug(f"Correcting error (attempt {attempt}): {error[:100]}...")

        # Create error correction prompt
        from .prompt_templates import PromptTemplates
        prompt_text = PromptTemplates.get_error_correction_prompt(
            query=query,
            code=code,
            error=error
        )

        from langchain.prompts import ChatPromptTemplate

        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a Python debugging expert. Fix the code to handle the error."),
            ("user", prompt_text)
        ])

        # Get correction
        chain = prompt | self.llm
        response = chain.invoke({})

        # Track usage
        self.calls_made += 1

        # Extract corrected code
        corrected_code = self._extract_code(response.content)

        if self.verbose:
            logger.debug(f"Generated corrected code ({len(corrected_code)} chars)")

        return corrected_code

    def _extract_code(self, response: str) -> str:
        """
        Extract Python code from LLM response.

        Handles markdown code blocks and other formatting.
        """
        code = response.strip()

        # Remove markdown code blocks
        if code.startswith("```python"):
            code = code[len("```python"):].strip()
        elif code.startswith("```"):
            code = code[3:].strip()

        if code.endswith("```"):
            code = code[:-3].strip()

        return code

    def get_usage_stats(self) -> Dict[str, Any]:
        """
        Get usage statistics.

        Returns:
            dict: Statistics including tokens used, calls made, etc.
        """
        return {
            "model": self.model,
            "provider": self.provider,
            "calls_made": self.calls_made,
            "total_tokens_used": self.total_tokens_used,
            "avg_tokens_per_call": (
                self.total_tokens_used / self.calls_made
                if self.calls_made > 0
                else 0
            ),
        }

    def reset_stats(self) -> None:
        """Reset usage statistics."""
        self.total_tokens_used = 0
        self.calls_made = 0
        if self.verbose:
            logger.debug("Usage statistics reset")

    def __repr__(self) -> str:
        return (
            f"LLMAgent(model='{self.model}', provider='{self.provider}', "
            f"temperature={self.temperature}, calls={self.calls_made})"
        )
