"""Remote API client (Claude, OpenAI, Gemini via litellm) + local OpenAI-compatible."""

from __future__ import annotations

import json
import logging
import os
import re
import time

from bilgeai.core.exceptions import LLMError
from bilgeai.llm.base import LLMClient, LLMResponse

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "claude-sonnet-4-20250514"
MAX_RETRIES = 2
RETRY_BASE_DELAY = 1.0


class RemoteClient(LLMClient):
    """LLM client via litellm (supports Claude, OpenAI, Gemini, and local endpoints)."""

    def __init__(
        self,
        model: str = DEFAULT_MODEL,
        api_key_env: str = "ANTHROPIC_API_KEY",
        base_url: str | None = None,
        is_local: bool = False,
    ) -> None:
        self.model = model
        self.api_key_env = api_key_env
        self.base_url = base_url
        self.is_local = is_local

    @property
    def source_name(self) -> str:
        return "local_llm" if self.is_local else "remote_api"

    def _get_api_key(self) -> str | None:
        if not self.api_key_env:
            return "no-key-needed"
        return os.environ.get(self.api_key_env)

    def is_available(self) -> bool:
        """Check if litellm is installed and API key is set (or local)."""
        try:
            import litellm  # noqa: F401

            if self.is_local:
                return True
            return self._get_api_key() is not None
        except ImportError:
            return False

    def generate(self, prompt: str, system: str | None = None) -> LLMResponse:
        """Generate via litellm.completion."""
        return self._call_llm(prompt, system=system)

    def generate_structured(
        self,
        prompt: str,
        response_schema: dict | None = None,
        system: str | None = None,
    ) -> LLMResponse:
        """Generate with JSON output enforcement.

        For cloud providers: uses response_format={"type": "json_object"}.
        For local providers: falls back to _extract_json_from_response if needed.
        """
        # First attempt: try with response_format
        try:
            response = self._call_llm(
                prompt,
                system=system,
                response_format={"type": "json_object"},
            )
            # Validate it's actually JSON
            json.loads(response.content)
            return response
        except (json.JSONDecodeError, Exception) as e:
            if not self.is_local:
                logger.warning("Structured output failed: %s. Retrying without response_format.", e)

        # Fallback: call without response_format and extract JSON
        response = self._call_llm(prompt, system=system)
        extracted = _extract_json_from_response(response.content)
        return LLMResponse(
            content=json.dumps(extracted),
            model=response.model,
            source=response.source,
            tokens_used=response.tokens_used,
        )

    def validate_api_key(self) -> bool:
        """Make a minimal test call to verify the API key / endpoint works.

        Raises ImportError if litellm is not installed (caller should handle).
        Returns True if the key works, False if the key is invalid.
        """
        import litellm

        # Suppress litellm's verbose info messages during validation
        original_verbosity = getattr(litellm, "suppress_debug_info", False)
        litellm.suppress_debug_info = True
        logging.getLogger("LiteLLM").setLevel(logging.WARNING)

        try:
            api_key = self._get_api_key()
            if not api_key and not self.is_local:
                return False

            messages = [{"role": "user", "content": "Say hi"}]
            kwargs: dict = {
                "model": self.model,
                "messages": messages,
                "max_tokens": 5,
                "timeout": 10,
            }
            if api_key and api_key != "no-key-needed":
                kwargs["api_key"] = api_key
            if self.base_url:
                kwargs["api_base"] = self.base_url

            response = litellm.completion(**kwargs)
            content = response.choices[0].message.content or ""
            return bool(content)
        except Exception:
            return False
        finally:
            litellm.suppress_debug_info = original_verbosity

    def _call_llm(
        self,
        prompt: str,
        system: str | None = None,
        max_tokens: int = 4096,
        response_format: dict | None = None,
    ) -> LLMResponse:
        """Internal LLM call with retry logic."""
        import litellm

        api_key = self._get_api_key()
        if not api_key and not self.is_local:
            raise LLMError(
                f"API key not found. Set {self.api_key_env} environment variable.",
                suggestion=f"export {self.api_key_env}=your-api-key-here",
            )

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

        kwargs: dict = {
            "model": self.model,
            "messages": messages,
            "max_tokens": max_tokens,
        }
        if api_key and api_key != "no-key-needed":
            kwargs["api_key"] = api_key
        if self.base_url:
            kwargs["api_base"] = self.base_url
        if response_format:
            kwargs["response_format"] = response_format

        last_error: Exception | None = None
        for attempt in range(MAX_RETRIES + 1):
            try:
                response = litellm.completion(**kwargs)
                content = response.choices[0].message.content or ""
                tokens = None
                if response.usage:
                    tokens = response.usage.total_tokens

                return LLMResponse(
                    content=content,
                    model=self.model,
                    source="local" if self.is_local else "remote",
                    tokens_used=tokens,
                )
            except Exception as e:
                last_error = e
                error_str = str(e).lower()

                # Don't retry auth errors
                if "auth" in error_str or "api key" in error_str or "invalid" in error_str:
                    raise LLMError(
                        f"Invalid API key. Check {self.api_key_env}.",
                        suggestion=f"export {self.api_key_env}=your-valid-key",
                    ) from e

                if attempt < MAX_RETRIES:
                    delay = RETRY_BASE_DELAY * (2**attempt)
                    logger.warning(
                        "LLM call failed (attempt %d/%d): %s. Retrying in %.1fs...",
                        attempt + 1,
                        MAX_RETRIES + 1,
                        e,
                        delay,
                    )
                    time.sleep(delay)

        # All retries exhausted
        if "rate" in str(last_error).lower():
            raise LLMError(
                "Rate limit exceeded. Wait or switch provider.",
                suggestion="Wait a moment and try again.",
            ) from last_error
        if "connect" in str(last_error).lower() or "network" in str(last_error).lower():
            raise LLMError(
                "Cannot reach API. Check your network connection.",
                suggestion="Check your internet connection.",
            ) from last_error

        raise LLMError(
            f"LLM API call failed: {last_error}",
            suggestion="Check your API key and network connection.",
        ) from last_error


def _extract_json_from_response(text: str) -> dict:
    """Extract JSON from LLM response with 3-layer fallback.

    1. Direct json.loads()
    2. Extract from ```json ... ``` markdown code block
    3. Find first { ... last } and parse
    """
    text = text.strip()

    # Layer 1: Direct parse
    try:
        return json.loads(text)
    except (json.JSONDecodeError, ValueError):
        pass

    # Layer 2: Markdown code block
    pattern = r"```(?:json)?\s*\n?(.*?)\n?```"
    matches = re.findall(pattern, text, re.DOTALL)
    for match in matches:
        try:
            return json.loads(match.strip())
        except (json.JSONDecodeError, ValueError):
            continue

    # Layer 3: First { to last }
    first_brace = text.find("{")
    last_brace = text.rfind("}")
    if first_brace != -1 and last_brace > first_brace:
        try:
            return json.loads(text[first_brace : last_brace + 1])
        except (json.JSONDecodeError, ValueError):
            pass

    raise LLMError(
        "Failed to parse JSON from LLM response.",
        suggestion="Try a different model or provider.",
    )
