"""
HTTP clients for Groq LLM and NVIDIA NIM.
- Async-first using httpx.AsyncClient
- Tenacity retry with exponential backoff
- Automatic JSON extraction from LLM responses
- Streaming support for real-time terminal output
"""
from __future__ import annotations

import json
import re
import asyncio
from typing import AsyncIterator, Optional
from dataclasses import dataclass

import httpx
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
    before_sleep_log,
)
import logging

logger = logging.getLogger("forge.api")


# ── Exceptions ────────────────────────────────────────────────────────────────

class APIError(Exception):
    def __init__(self, provider: str, status: int, message: str):
        self.provider = provider
        self.status = status
        super().__init__(f"{provider} API {status}: {message}")


class RateLimitError(APIError):
    pass


class AuthError(APIError):
    pass


# ── Response parsing ──────────────────────────────────────────────────────────

def strip_fences(text: str) -> str:
    """Remove markdown code fences from LLM output."""
    text = re.sub(r"^```[\w]*\s*\n?", "", text, flags=re.MULTILINE)
    text = re.sub(r"^```\s*$", "", text, flags=re.MULTILINE)
    return text.strip()


def extract_json(text: str) -> dict:
    """
    Robustly extract and parse a JSON object from LLM text.
    Handles: raw JSON, fenced JSON, JSON embedded in prose.
    """
    clean = strip_fences(text)

    # Try direct parse
    try:
        return json.loads(clean)
    except json.JSONDecodeError:
        pass

    # Find first balanced { } block
    depth = 0
    start = None
    for i, ch in enumerate(clean):
        if ch == "{":
            if start is None:
                start = i
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0 and start is not None:
                try:
                    return json.loads(clean[start : i + 1])
                except json.JSONDecodeError:
                    start = None
                    depth = 0

    raise ValueError(f"No valid JSON found in response (first 300 chars):\n{text[:300]}")


# ── Groq Client ───────────────────────────────────────────────────────────────

class GroqClient:
    BASE_URL = "https://api.groq.com/openai/v1/chat/completions"

    def __init__(self, api_key: str, timeout: int = 120):
        self.api_key = api_key
        self._headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        self._timeout = httpx.Timeout(timeout, connect=15.0)

    def _handle_error(self, resp: httpx.Response) -> None:
        try:
            body = resp.json()
            msg = body.get("error", {}).get("message", resp.text[:200])
        except Exception:
            msg = resp.text[:200]

        if resp.status_code == 401:
            raise AuthError("Groq", 401, "Invalid API key")
        if resp.status_code == 429:
            raise RateLimitError("Groq", 429, msg)
        raise APIError("Groq", resp.status_code, msg)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=15),
        retry=retry_if_exception_type(RateLimitError),
        reraise=True,
    )
    async def complete(
        self,
        messages: list[dict],
        model: str = "llama-3.3-70b-versatile",
        temperature: float = 0.3,
        max_tokens: int = 8000,
    ) -> str:
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(self.BASE_URL, json=payload, headers=self._headers)
            if not resp.is_success:
                self._handle_error(resp)
            data = resp.json()
            return data["choices"][0]["message"]["content"]

    async def complete_json(
        self,
        messages: list[dict],
        model: str,
        temperature: float = 0.2,
    ) -> dict:
        raw = await self.complete(messages, model=model, temperature=temperature)
        return extract_json(raw)

    async def stream(
        self,
        messages: list[dict],
        model: str,
        temperature: float = 0.2,
        max_tokens: int = 6000,
    ) -> AsyncIterator[str]:
        """Stream tokens for real-time display."""
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": True,
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            async with client.stream(
                "POST", self.BASE_URL, json=payload, headers=self._headers
            ) as resp:
                if not resp.is_success:
                    content = await resp.aread()
                    resp._content = content
                    self._handle_error(resp)
                async for line in resp.aiter_lines():
                    if line.startswith("data: "):
                        chunk = line[6:]
                        if chunk == "[DONE]":
                            break
                        try:
                            token = json.loads(chunk)["choices"][0]["delta"].get("content", "")
                            if token:
                                yield token
                        except Exception:
                            pass


# ── NVIDIA NIM Client ─────────────────────────────────────────────────────────

class NvidiaClient:
    BASE_URL = "https://integrate.api.nvidia.com/v1/chat/completions"

    def __init__(self, api_key: str, timeout: int = 180):
        self.api_key = api_key
        self._headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        self._timeout = httpx.Timeout(timeout, connect=20.0)

    def _handle_error(self, resp: httpx.Response) -> None:
        try:
            body = resp.json()
            msg = body.get("detail", body.get("message", resp.text[:200]))
        except Exception:
            msg = resp.text[:200]

        if resp.status_code == 401:
            raise AuthError("NVIDIA", 401, "Invalid API key")
        if resp.status_code == 429:
            raise RateLimitError("NVIDIA", 429, msg)
        raise APIError("NVIDIA", resp.status_code, msg)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=2, min=3, max=20),
        retry=retry_if_exception_type(RateLimitError),
        reraise=True,
    )
    async def generate_code(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str = "meta/codellama-70b-instruct",
        max_tokens: int = 4096,
    ) -> str:
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.12,
            "max_tokens": max_tokens,
            "stream": False,
        }
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(self.BASE_URL, json=payload, headers=self._headers)
            if not resp.is_success:
                self._handle_error(resp)
            data = resp.json()
            raw = data["choices"][0]["message"]["content"]
            return strip_fences(raw)

    async def generate_code_with_fallback(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str,
        groq_client: GroqClient,
        groq_model: str,
    ) -> tuple[str, str]:
        """
        Try NVIDIA, fall back to Groq.
        Returns (code, source) where source is 'nvidia' or 'groq'.
        """
        try:
            code = await self.generate_code(system_prompt, user_prompt, model)
            return code, "nvidia"
        except (APIError, Exception) as e:
            logger.warning(f"NVIDIA failed ({e}), falling back to Groq")
            raw = await groq_client.complete(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                model=groq_model,
                temperature=0.15,
                max_tokens=6000,
            )
            return strip_fences(raw), "groq"
