"""
LLM Client Module for TestGen AI.

This module provides a unified interface to interact with various LLM providers
(OpenAI, Anthropic, Gemini, Ollama) using LiteLLM.
"""

import os
from typing import Optional, Dict, Any, List
from enum import Enum

import litellm
from litellm import completion
from pydantic import BaseModel, Field

from testgen.config import config, LLMProvider


class LLMResponse(BaseModel):
    """Encapsulates the structured response from an LLM provider.

    This model provides a unified format for LLM outputs, including the
    generated content and detailed metadata such as token usage and cost.

    Attributes:
        content (str): The raw text generated by the LLM.
        model (str): Name of the specific model that produced the response.
        provider (str): The AI provider used (e.g., 'openai', 'anthropic').
        tokens_used (int): Total combined input and output tokens.
        input_tokens (int): Number of tokens in the input prompt.
        output_tokens (int): Number of tokens in the generated response.
        cost (float): Estimated monetary cost of the generation in USD.
    """
    content: str
    model: str = ""
    provider: str = ""
    tokens_used: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    cost: float = 0.0

    class Config:
        arbitrary_types_allowed = True


class LLMClient:
    """A unified, provider-agnostic client for interacting with Large Language Models.

    The LLMClient abstracts away the differences between various AI providers
    (OpenAI, Anthropic, Google Gemini, and local Ollama) by utilizing LiteLLM.
    It provides a consistent interface for text and test code generation,
    handling API keys, model formatting, and error scenarios.

    Attributes:
        provider (LLMProvider): The active LLM provider.
        model (str): The name of the model being used.
        temperature (float): The generation randomness parameter.
        max_tokens (int): The maximum tokens allowed for the output.

    Example:
        client = LLMClient()
        response = client.generate("Analyze this code for security vulnerabilities.")
        print(f"Generated Content: {response.content}")
    """
    
    def __init__(
        self,
        provider: Optional[LLMProvider] = None,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ):
        """Initialize a new LLM client with the specified provider and configuration.

        Sets up the environment for the selected provider (e.g., API keys) and
        configures LiteLLM's internal execution parameters.

        Args:
            provider (Optional[LLMProvider]): The LLM provider to use. If None,
                defaults to the provider specified in the configuration.
            model (Optional[str]): The specific model name to use. If None,
                defaults to the configured model for the provider.
            api_key (Optional[str]): Explicit API key. If None, attempts to
                retrieve the key from the configuration or environment variables.
            temperature (Optional[float]): Randomness parameter for generation.
                Defaults to the configured system value if None.
            max_tokens (Optional[int]): Maximum number of tokens to generate in
                the output. Defaults to the configured context limit if None.
        """
        # Use provided values or fall back to config
        self.provider = provider or config.llm_provider
        self.model = model or config.llm_model
        self.temperature = temperature or config.llm_temperature
        self.max_tokens = max_tokens or config.max_context_tokens
        
        # Set up API key
        self._setup_api_key(api_key)
        
        # Configure LiteLLM
        self._configure_litellm()
        
        # Validate setup
        config.validate_api_keys()
    
    def _setup_api_key(self, api_key: Optional[str] = None) -> None:
        """Configure the API key for the selected provider.

        Handles key precedence (explicit > config > env) and exports the key
        to the appropriate environment variable required by LiteLLM or the SDK.

        Args:
            api_key (Optional[str]): Explicit API key to use.
        """
        if api_key:
            # Use provided key
            self.api_key = api_key
        else:
            # Get from config
            self.api_key = config.get_api_key()
        
        # Set environment variables for LiteLLM
        if self.provider == LLMProvider.OPENAI and self.api_key:
            os.environ["OPENAI_API_KEY"] = self.api_key
        elif self.provider == LLMProvider.ANTHROPIC and self.api_key:
            os.environ["ANTHROPIC_API_KEY"] = self.api_key
        elif self.provider == LLMProvider.GEMINI and self.api_key:
            # For Google AI Studio (not Vertex AI)
            os.environ["GEMINI_API_KEY"] = self.api_key
            os.environ["GOOGLE_API_KEY"] = self.api_key  # LiteLLM also checks this
        # Ollama doesn't need API key
    
    def _configure_litellm(self) -> None:
        """Global configuration for the LiteLLM library.
        
        Disables telemetry and verbose logging for production use, and sets
        a standard request timeout.
        """
        # Disable verbose logging in production
        litellm.set_verbose = False
        
        # Set telemetry (can be disabled for privacy)
        litellm.telemetry = False
        
        # Set timeout
        litellm.request_timeout = 60
    
    def _format_model_name(self) -> str:
        """Format the model name with appropriate provider prefixes for LiteLLM.

        LiteLLM requires specific prefixes (e.g., 'gemini/', 'ollama/') for
        some providers to correctly route requests.

        Returns:
            str: The formatted model identifier ready for use with `litellm.completion`.
        """
        if self.provider == LLMProvider.GEMINI:
            # Gemini models need "gemini/" prefix for Google AI Studio
            if not self.model.startswith("gemini/"):
                return f"gemini/{self.model}"
        elif self.provider == LLMProvider.OLLAMA:
            # Ollama models need "ollama/" prefix
            if not self.model.startswith("ollama/"):
                return f"ollama/{self.model}"
        
        # OpenAI and Anthropic don't need prefixes
        return self.model
    
    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        stop: Optional[List[str]] = None,
    ) -> LLMResponse:
        """Generate a text completion from the LLM based on the provided prompt.

        This is the primary synchronous entry point for LLM interaction. It
        automatically routes the request to the correct provider implementation
        (direct SDK vs. LiteLLM).

        Args:
            prompt (str): The primary user instruction or question for the LLM.
            system_prompt (Optional[str]): High-level context or role-playing
                instructions for the model.
            temperature (Optional[float]): Override for the generation randomness.
            max_tokens (Optional[int]): Override for the maximum output length.
            stop (Optional[List[str]]): A list of sequences that will cause the
                model to stop generating.

        Returns:
            LLMResponse: A structured response object containing the generated content,
                usage statistics, and cost estimation.

        Raises:
            ValueError: If the generation fails due to API errors or invalid configuration.
        """
        # Use direct SDK for Gemini, LiteLLM for others
        if self.provider == LLMProvider.GEMINI:
            return self._generate_gemini(prompt, system_prompt, temperature, max_tokens)
        else:
            return self._generate_litellm(prompt, system_prompt, temperature, max_tokens, stop)
    
    def generate_test(
        self,
        source_code: str,
        language: str = "python",
        file_path: Optional[str] = None,
        **kwargs
    ) -> str:
        """Generate test code for the provided source code using a specialized prompt."""
        import ast as _ast

        system_prompt = (
            f"You are an expert test engineer specializing in {language} testing.\n"
            "Return ONLY valid, runnable Python source code — no markdown, no backticks, "
            "no explanations, no prose.\n"
            "CRITICAL FORMATTING RULES:\n"
            "1. Every string literal (including assertion strings) MUST be on a single line.\n"
            "2. Never wrap a long line by breaking a string across two lines.\n"
            "3. If an assertion string is long, shorten it — never split it with a newline.\n"
            "4. Docstrings may use triple quotes but MUST fit on one or three lines — never "
            "   wrap a single-quoted or double-quoted string literal mid-sentence.\n"
            "5. The output must be parseable by Python's ast.parse() without errors."
        )

        user_prompt = (
            f"Generate pytest unit tests for the following {language} code.\n\n"
            f"File: {file_path or 'unknown'}\n\n"
            f"```{language}\n{source_code}\n```\n\n"
            "Requirements:\n"
            "- Use pytest (functions or classes starting with test_/Test)\n"
            "- Cover normal cases, edge cases, and error handling\n"
            "- Each test has a one-line docstring\n"
            "- All string literals on a SINGLE line — never break them with a newline\n"
            "- Output ONLY the Python source code, no markdown fences\n\n"
            "Python test code:"
        )

        response = self.generate(user_prompt, system_prompt=system_prompt, **kwargs)
        test_code = response.content

        # Strip markdown code fences (```python ... ``` or ``` ... ```)
        test_code = self._strip_code_fences(test_code)

        # Repair common LLM formatting issues (line-wrapped string literals)
        test_code = self._repair_linebroken_strings(test_code)

        # Validate syntax; if still broken, emit a warning comment at the top
        try:
            _ast.parse(test_code)
        except SyntaxError as exc:
            # Prepend a warning so the user knows — file will still be saved
            test_code = (
                f"# WARNING: TestGen AI detected a syntax issue in this file: {exc}\n"
                f"# Please review and fix manually, or re-run 'testgen generate'.\n\n"
                + test_code
            )

        return test_code

    @staticmethod
    def _strip_code_fences(code: str) -> str:
        """Remove markdown code fences from LLM output."""
        lines = code.strip().splitlines()
        # Drop leading fence line (```python, ```py, ```, etc.)
        if lines and lines[0].lstrip().startswith("```"):
            lines = lines[1:]
        # Drop trailing fence line
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        # Handle rare double-wrapping
        result = "\n".join(lines).strip()
        if result.startswith("```"):
            result = result[result.index("\n") + 1:]
        if result.endswith("```"):
            result = result[:result.rindex("```")]
        return result.strip()

    @staticmethod
    def _repair_linebroken_strings(code: str) -> str:
        """
        Detect and repair string literals that the LLM has broken across two lines.

        Pattern: a line ends without closing its string quote, and the next
        line is a continuation ending properly (same quote type).  Join them.
        """
        import re
        lines = code.splitlines()
        out: List[str] = []
        i = 0
        while i < len(lines):
            line = lines[i]
            # Count unescaped single and double quotes to detect open strings
            # Quick heuristic: if odd number of unescaped " → might be open
            stripped = line.rstrip()
            # Check if there is an unclosed string on this line
            # Use a simple tokeniser: walk the line tracking quote state
            in_single = False
            in_double = False
            j = 0
            while j < len(stripped):
                c = stripped[j]
                if c == '\\':
                    j += 2
                    continue
                if c == '"' and not in_single:
                    in_double = not in_double
                elif c == "'" and not in_double:
                    in_single = not in_single
                j += 1
            # If a string is left open AND there is a next line, merge them
            if (in_single or in_double) and i + 1 < len(lines):
                combined = stripped + " " + lines[i + 1].lstrip()
                out.append(combined)
                i += 2
            else:
                out.append(line)
                i += 1
        return "\n".join(out)
    
    def _generate_gemini(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> LLMResponse:
        """Internal implementation for direct interaction with the Google Gemini SDK.

        Args:
            prompt (str): User instruction.
            system_prompt (Optional[str]): Role context.
            temperature (Optional[float]): Randomness.
            max_tokens (Optional[int]): Token limit.

        Returns:
            LLMResponse: Structured response with metadata.
        """
        try:
            from google import genai
            from google.genai import types as genai_types

            client = genai.Client(api_key=self.api_key)

            # Build prompt (system instruction prepended if provided)
            full_prompt = prompt
            if system_prompt:
                full_prompt = f"{system_prompt}\n\n{prompt}"

            # Strip any "models/" prefix — google.genai uses bare names
            model_name = self.model
            if model_name.startswith("models/"):
                model_name = model_name[len("models/"):]

            response = client.models.generate_content(
                model=model_name,
                contents=full_prompt,
                config=genai_types.GenerateContentConfig(
                    temperature=temperature or self.temperature,
                    max_output_tokens=max_tokens or self.max_tokens,
                ),
            )

            content = response.text

            # Use token metadata when available
            usage = getattr(response, "usage_metadata", None)
            input_tokens = getattr(usage, "prompt_token_count", 0) or int(len(full_prompt.split()) * 1.3)
            output_tokens = getattr(usage, "candidates_token_count", 0) or int(len(content.split()) * 1.3)

            return LLMResponse(
                content=content,
                model=self.model,
                provider=self.provider.value,
                tokens_used=input_tokens + output_tokens,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=0.0,  # Gemini free tier
            )

        except Exception as e:
            raise ValueError(f"Gemini generation failed: {str(e)}") from e
    
    def _generate_litellm(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        stop: Optional[List[str]] = None,
    ) -> LLMResponse:
        """Internal implementation for generating text via the LiteLLM proxy.

        Used for OpenAI, Anthropic, and local Ollama deployments.

        Args:
            prompt (str): User instruction.
            system_prompt (Optional[str]): Role context.
            temperature (Optional[float]): Randomness.
            max_tokens (Optional[int]): Token limit.
            stop (Optional[List[str]]): Stop sequences.

        Returns:
            LLMResponse: Structured response with metadata.
        """
        # Build messages
        messages = []
        
        if system_prompt:
            messages.append({
                "role": "system",
                "content": system_prompt
            })
        
        messages.append({
            "role": "user",
            "content": prompt
        })
        
        # Prepare parameters
        params = {
            "model": self._format_model_name(),
            "messages": messages,
            "temperature": temperature or self.temperature,
            "max_tokens": max_tokens or self.max_tokens,
        }
        
        if stop:
            params["stop"] = stop
        
        try:
            # Call LiteLLM
            response = completion(**params)
            
            # Extract data
            content = response.choices[0].message.content
            usage = response.usage if hasattr(response, 'usage') else None
            
            # Build response
            return LLMResponse(
                content=content,
                model=self.model,
                provider=self.provider.value,
                tokens_used=usage.total_tokens if usage else 0,
                input_tokens=usage.prompt_tokens if usage else 0,
                output_tokens=usage.completion_tokens if usage else 0,
                cost=self._estimate_cost(usage) if usage else 0.0
            )
            
        except Exception as e:
            raise ValueError(f"LLM generation failed: {str(e)}") from e
    
    async def generate_async(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        stop: Optional[List[str]] = None,
    ) -> LLMResponse:
        """Asynchronously generate a text completion.

        Useful for parallel generation tasks or integration with async frameworks.

        Args:
            prompt (str): User instruction.
            system_prompt (Optional[str]): Role context.
            temperature (Optional[float]): Randomness.
            max_tokens (Optional[int]): Token limit.
            stop (Optional[List[str]]): Stop sequences.

        Returns:
            LLMResponse: Structured response with metadata.
        """
        # Build messages
        messages = []
        
        if system_prompt:
            messages.append({
                "role": "system",
                "content": system_prompt
            })
        
        messages.append({
            "role": "user",
            "content": prompt
        })
        
        # Prepare parameters
        params = {
            "model": self._format_model_name(),
            "messages": messages,
            "temperature": temperature or self.temperature,
            "max_tokens": max_tokens or self.max_tokens,
        }
        
        if stop:
            params["stop"] = stop
        
        try:
            from litellm import acompletion
            
            # Call LiteLLM async
            response = await acompletion(**params)
            
            # Extract data
            content = response.choices[0].message.content
            usage = response.usage if hasattr(response, 'usage') else None
            
            # Build response
            return LLMResponse(
                content=content,
                model=self.model,
                provider=self.provider.value,
                tokens_used=usage.total_tokens if usage else 0,
                input_tokens=usage.prompt_tokens if usage else 0,
                output_tokens=usage.completion_tokens if usage else 0,
                cost=self._estimate_cost(usage) if usage else 0.0
            )
            
        except Exception as e:
            raise ValueError(f"LLM generation failed: {str(e)}") from e
    
    def _estimate_cost(self, usage: Any) -> float:
        """Calculate the estimated monetary cost of a generation based on token usage.

        Uses built-in price tables for common models. Provides a rough estimate
        for developer budgeting.

        Args:
            usage (Any): Token usage object returned by the model or LiteLLM.

        Returns:
            float: Estimated cost in USD, rounded to 6 decimal places.
        """
        if not usage or not hasattr(usage, 'prompt_tokens'):
            return 0.0
        
        input_tokens = usage.prompt_tokens
        output_tokens = usage.completion_tokens
        
        # Cost per 1M tokens
        costs = {
            "gpt-3.5-turbo": {"input": 0.5, "output": 1.5},
            "gpt-4": {"input": 10, "output": 30},
            "gpt-4-turbo": {"input": 10, "output": 30},
            "claude-3-5-haiku": {"input": 0.25, "output": 1.25},
            "claude-3-5-sonnet": {"input": 3, "output": 15},
            "claude-3": {"input": 3, "output": 15},
            "gemini": {"input": 0, "output": 0},  # Free tier
            "ollama": {"input": 0, "output": 0},  # Local/free
        }
        
        # Find matching model
        model_lower = self.model.lower()
        cost_data = None
        
        for model_prefix, cost in costs.items():
            if model_prefix in model_lower:
                cost_data = cost
                break
        
        if not cost_data:
            return 0.0
        
        # Calculate cost (per million tokens)
        input_cost = (input_tokens / 1_000_000) * cost_data["input"]
        output_cost = (output_tokens / 1_000_000) * cost_data["output"]
        
        return round(input_cost + output_cost, 6)
    
    def test_connection(self) -> bool:
        """Verify that the configured LLM provider and API key are functional.

        Sends a minimal 'ping' request to the model to ensure connectivity.

        Returns:
            bool: True if the model responds successfully, False otherwise.
        """
        try:
            response = self.generate(
                prompt="Say 'OK' if you can read this.",
                max_tokens=10
            )
            return len(response.content) > 0
        except Exception:
            return False
    
    def get_model_info(self) -> Dict[str, Any]:
        """Summarize the current LLM configuration and capability state.

        Returns:
            Dict[str, Any]: A dictionary containing provider, model name,
                generation settings, and API key availability status.
        """
        return {
            "provider": self.provider.value,
            "model": self.model,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "has_api_key": self.api_key is not None,
        }


# Convenience instance with default configuration
llm_client = LLMClient()
