"""AI integration for gar - LLM code review.

This module provides a robust interface for calling LLM APIs
with proper error handling, retry logic, and response validation.
"""

import json
import logging
import os
import time
from typing import Callable, Optional

import httpx
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)

# Configuration
DEFAULT_TIMEOUT = 120.0  # seconds
DEFAULT_MAX_TOKENS = 8000
DEFAULT_TEMPERATURE = 0.3
DEFAULT_MAX_RETRIES = 3
RETRY_DELAY = 2.0  # seconds between retries


class LLMError(Exception):
    """Base exception for LLM-related errors."""
    pass


class LLMConnectionError(LLMError):
    """Failed to connect to LLM API."""
    pass


class LLMResponseError(LLMError):
    """Invalid response from LLM API."""
    pass


class LLMParseError(LLMError):
    """Failed to parse LLM response."""
    pass


def generate_review(
    model: str,
    prompt: str | tuple[str, str] | list[dict[str, str]],
    temperature: float = DEFAULT_TEMPERATURE,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    max_retries: int = DEFAULT_MAX_RETRIES,
    stream_callback: Callable[[str], None] | None = None,
) -> str:
    """Generate code review using LLM.
    
    Args:
        model: Model name (e.g., "qwen3-coder:30b")
        prompt: Either string, tuple of (system, user), or list of messages
        temperature: Sampling temperature (0.0-1.0)
        max_tokens: Max output tokens
        max_retries: Max retry attempts
        stream_callback: Optional callback function for streaming chunks
        
    Returns:
        Response content (preferably JSON string)
        
    Raises:
        LLMError: For any LLM-related errors
    """
    messages = _build_messages(prompt)
    return _call_llm_with_retry(messages, model, temperature, max_tokens, max_retries, stream_callback=stream_callback)


def _build_messages(prompt: str | tuple[str, str] | list[dict[str, str]]) -> list[dict[str, str]]:
    """Convert various prompt formats to message list."""
    if isinstance(prompt, list):
        return [{**msg} for msg in prompt]
    elif isinstance(prompt, tuple):
        system_prompt, user_prompt = prompt
        return [
            {"role": "system", "content": system_prompt or ""},
            {"role": "user", "content": user_prompt},
        ]
    else:
        return [
            {"role": "system", "content": "You are a code review assistant."},
            {"role": "user", "content": str(prompt)},
        ]


def _call_llm_with_retry(
    messages: list[dict[str, str]],
    model: str,
    temperature: float,
    max_tokens: int,
    max_retries: int,
    stream_callback: Callable[[str], None] | None = None,
) -> str:
    """Call LLM API with retry logic."""
    url = _get_api_url()
    body = _build_request_body(model, messages, temperature, max_tokens)
    headers = {"Content-Type": "application/json"}
    
    logger.info(f"model={model}, max_tokens={max_tokens}")
    
    last_error: Optional[Exception] = None
    
    for attempt in range(max_retries):
        try:
            # stream_callback은 더 이상 사용하지 않고, 전체 응답을 한 번에 받는다.
            content = _make_request(url, body, headers)
            validated_content = _validate_and_extract_json(content)
            return validated_content
            
        except httpx.ConnectError as e:
            last_error = LLMConnectionError(f"Failed to connect: {e}")
            logger.warning(f"Connection error (attempt {attempt + 1}/{max_retries}): {e}")
            
        except httpx.TimeoutException as e:
            last_error = LLMConnectionError(f"Request timeout: {e}")
            logger.warning(f"Timeout (attempt {attempt + 1}/{max_retries}): {e}")
            
        except httpx.HTTPStatusError as e:
            last_error = LLMResponseError(f"HTTP error {e.response.status_code}: {e}")
            logger.warning(f"HTTP error (attempt {attempt + 1}/{max_retries}): {e}")
            
        except LLMParseError as e:
            # Don't retry parse errors - the response is what it is
            logger.warning(f"Parse error: {e}")
            raise
            
        except Exception as e:
            last_error = LLMError(f"Unexpected error: {e}")
            logger.warning(f"Unexpected error (attempt {attempt + 1}/{max_retries}): {e}")
        
        # Wait before retry (except on last attempt)
        if attempt < max_retries - 1:
            time.sleep(RETRY_DELAY * (attempt + 1))
    
    raise last_error or LLMError("All retry attempts failed")


def _get_api_url() -> str:
    """Get API URL from environment."""
    base_url = os.getenv("GAR_LLM_BASE_URL")
    if not base_url:
        raise LLMConnectionError("MODEL_URL not set in .env file")
    
    # Ensure /chat/completions endpoint
    if "/chat/completions" not in base_url:
        base_url = base_url.rstrip("/")
        return f"{base_url}/chat/completions"
    return base_url


def _build_request_body(
    model: str,
    messages: list[dict[str, str]],
    temperature: float,
    max_tokens: int,
) -> dict:
    """Build request body for OpenAI-compatible API."""
    return {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_completion_tokens": max_tokens,
    }


def _make_request(url: str, body: dict, headers: dict) -> str:
    """Make HTTP POST request and return content (non-streaming)."""
    response = httpx.post(
        url,
        json=body,
        headers=headers,
        timeout=DEFAULT_TIMEOUT,
    )
    response.raise_for_status()

    response_data = response.json()
    content = _extract_content(response_data)

    if not content:
        raise LLMResponseError("Empty response from LLM")

    logger.info(f"LLM response received (length: {len(content)})")
    return content


def _extract_content(response_data: dict) -> str:
    """Extract content from API response."""
    # OpenAI-style response
    choices = response_data.get("choices")
    if choices and isinstance(choices, list) and len(choices) > 0:
        content = choices[0].get("message", {}).get("content")
        if content:
            return content
    
    # Ollama-style response
    message = response_data.get("message", {})
    if "content" in message:
        return message["content"]
    
    raise LLMResponseError(f"Could not extract content from response: {list(response_data.keys())}")


def _validate_and_extract_json(content: str) -> str:
    """Validate response and extract JSON if needed.
    
    Handles various response formats:
    - Pure JSON
    - JSON wrapped in ```json ... ```
    - JSON wrapped in ``` ... ```
    """
    content = content.strip()
    
    # Check if response starts with JSON
    if content.startswith("{"):
        # Validate it's complete JSON
        if _is_valid_json(content):
            return content
        else:
            logger.warning("Response starts with '{' but is not valid JSON - may be truncated")
            # Return anyway, let caller handle parse error
            return content
    
    # Try to extract from markdown code block
    extracted = _extract_from_markdown(content)
    if extracted:
        return extracted
    
    # If content doesn't look like JSON at all, return as-is
    logger.warning(f"Response doesn't appear to be JSON (starts with: {content[:50]}...)")
    return content


def _is_valid_json(content: str) -> bool:
    """Check if content is valid JSON."""
    try:
        json.loads(content)
        return True
    except json.JSONDecodeError:
        return False


def _extract_from_markdown(content: str) -> Optional[str]:
    """Extract JSON from markdown code block."""
    # Try ```json ... ```
    if "```json" in content:
        try:
            json_str = content.split("```json")[1].split("```")[0].strip()
            if json_str and _is_valid_json(json_str):
                logger.info("Extracted valid JSON from ```json block")
                return json_str
        except (IndexError, ValueError):
            pass
    
    # Try ``` ... ```
    if "```" in content:
        try:
            parts = content.split("```")
            if len(parts) >= 3:
                json_str = parts[1].strip()
                # Remove language hint if present (e.g., "json\n{...")
                if json_str.startswith(("json", "JSON")):
                    json_str = json_str[4:].strip()
                if json_str and _is_valid_json(json_str):
                    logger.info("Extracted valid JSON from code block")
                    return json_str
        except (IndexError, ValueError):
            pass
    
    return None
