app/core/llm.py

import asyncio
import json
from typing import Any, Optional, cast

import httpx
from tenacity import (
    RetryCallState,
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from app.core.config import get_settings
from app.core.logger import get_logger
from app.core.metrics import metrics
from app.core.circuit_breaker import llm_breaker

logger = get_logger(__name__)

OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"


def _log_retry(rs: RetryCallState) -> None:
    error = str(rs.outcome.exception()) if rs.outcome is not None else "unknown"
    logger.warning(
        "LLM call failed, retrying",
        extra={"attempt": rs.attempt_number, "error": error},
    )


async def _wait_for_rate_limit(exc: httpx.HTTPStatusError) -> None:
    retry_after = exc.response.headers.get("Retry-After")
    if retry_after:
        await asyncio.sleep(float(retry_after))


@retry(
    retry=retry_if_exception_type((httpx.HTTPStatusError, httpx.TimeoutException)),
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
    before_sleep=_log_retry,
)
async def call_llm(
    prompt: str, system_prompt: Optional[str] = None, max_tokens: int = 2000
) -> str:
    """Async call to OpenRouter API with retry logic and circuit breaker."""
    if llm_breaker.is_open():
        raise RuntimeError("LLM circuit breaker is OPEN — upstream is unavailable")

    settings = get_settings()
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    headers = {
        "Authorization": f"Bearer {settings.openrouter_api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co/spaces",
        "X-Title": "AI Code Review Agent",
    }

    payload = {
        "model": settings.llm_model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.3,
    }

    logger.info(
        "Calling LLM", extra={"model": settings.llm_model, "max_tokens": max_tokens}
    )
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(OPENROUTER_API_URL, headers=headers, json=payload)
            if response.status_code == 429:
                await _wait_for_rate_limit(
                    httpx.HTTPStatusError(
                        "rate limited", request=response.request, response=response
                    )
                )
                response.raise_for_status()

        data = response.json()
        usage = data.get("usage", {})
        total_tokens = usage.get("total_tokens", 0)
        if total_tokens:
            metrics.record_tokens("llm_total", total_tokens)
        llm_breaker.record_success()
        return str(data["choices"][0]["message"]["content"])
    except (httpx.HTTPStatusError, httpx.TimeoutException) as exc:
        llm_breaker.record_failure()
        raise exc


async def call_llm_for_json(prompt: str, max_tokens: int = 2000, task: str = "default") -> dict:
    """
    Call the LLM expecting a JSON response.
    If the first response fails to parse, sends a corrective follow-up once.
    Returns a dict (empty on total failure).
    """

    response = await call_llm_routed(prompt, task=task, max_tokens=max_tokens)

    def _try_parse(raw: str) -> dict[str, Any] | None:
        try:
            cleaned = (
                raw.strip()
                .removeprefix("```json")
                .removeprefix("```")
                .removesuffix("```")
                .strip()
            )
            return cast(dict[str, Any], json.loads(cleaned))
        except json.JSONDecodeError:
            return None

    result = _try_parse(response)
    if result is not None:
        return result

    logger.warning(
        "LLM returned invalid JSON, sending corrective prompt",
        extra={"raw": response[:200]},
    )
    corrective_prompt = (
        f"Your previous response was not valid JSON. "
        f"Here is what you returned:\n\n{response}\n\n"
        f"Return ONLY the corrected valid JSON object, no extra text, no markdown."
    )
    retry_response = await call_llm(corrective_prompt, max_tokens=max_tokens)
    result = _try_parse(retry_response)
    if result is not None:
        return result

    logger.error(
        "LLM failed to return valid JSON after corrective prompt",
        extra={"raw": retry_response[:200]},
    )
    return {}

async def call_llm_routed(
    prompt: str,
    *,
    task: str = "default",
    system_prompt: Optional[str] = None,
    max_tokens: int = 2000,
) -> str:
    """
    Route to cheap or expensive model based on task type.
    
    cheap tasks:  repo_analysis, test_generation, report_summary
    expensive tasks: bug_detection, code_review, security_review
    """
    settings = get_settings()
    CHEAP_TASKS = {"repo_analysis", "test_generation", "report_summary"}
    model = (
        settings.llm_model_cheap if task in CHEAP_TASKS else settings.llm_model_expensive
    )

    return await _call_llm_with_model(
        prompt=prompt,
        model=model,
        system_prompt=system_prompt,
        max_tokens=max_tokens,
        task=task,
    )


async def _call_llm_with_model(
    prompt: str,
    model: str,
    system_prompt: Optional[str] = None,
    max_tokens: int = 2000,
    task: str = "default",
) -> str:
    """Internal: call LLM with an explicit model string (bypasses settings.llm_model)."""
    if llm_breaker.is_open():
        raise RuntimeError("LLM circuit breaker is OPEN — upstream is unavailable")

    settings = get_settings()
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    headers = {
        "Authorization": f"Bearer {settings.openrouter_api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co/spaces",
        "X-Title": "AI Code Review Agent",
    }
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.3,
    }

    logger.info("Calling LLM (routed)", extra={"model": model, "task": task, "max_tokens": max_tokens})
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(OPENROUTER_API_URL, headers=headers, json=payload)
            response.raise_for_status()

        data = response.json()
        usage = data.get("usage", {})
        total_tokens = usage.get("total_tokens", 0)
        if total_tokens:
            metrics.record_tokens(f"llm_{task}", total_tokens)
        llm_breaker.record_success()
        return str(data["choices"][0]["message"]["content"])
    except (httpx.HTTPStatusError, httpx.TimeoutException) as exc:
        llm_breaker.record_failure()
        raise exc