import os
import json
import re
import time
import logging
from typing import Dict, Any, Optional
from openai import OpenAI, APIError, AuthenticationError, RateLimitError

from .config import settings

logger = logging.getLogger("document_intelligence.nvidia")


class NVIDIAClient:
    """Client for invoking NVIDIA LLM inference endpoints via OpenAI compatible API."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
    ):
        self.api_key = api_key or settings.NVIDIA_API_KEY
        self.base_url = base_url or settings.NVIDIA_BASE_URL
        self.model = model or settings.NVIDIA_MODEL
        self._client: Optional[OpenAI] = None

    def _get_client(self) -> OpenAI:
        """Lazily initialize and return the OpenAI client configured for NVIDIA."""
        if not self.api_key:
            raise ValueError(
                "NVIDIA_API_KEY is not set. Please provide a valid NVIDIA API Key via environment variables."
            )
        if self._client is None:
            self._client = OpenAI(
                base_url=self.base_url,
                api_key=self.api_key,
                timeout=60.0
            )
        return self._client

    def generate_document_intelligence(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.2,
        top_p: float = 0.95,
        max_tokens: int = 4096,
    ) -> Dict[str, Any]:
        """
        Send document extraction text and prompt to NVIDIA LLM.
        Returns parsed JSON dict, reasoning content (if any), and latency.
        """
        client = self._get_client()
        start_time = time.time()

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

        logger.info(
            "Invoking NVIDIA model '%s' at '%s'", self.model, self.base_url
        )

        try:
            extra_body = {}
            if "nemotron" in self.model.lower():
                extra_body["chat_template_kwargs"] = {"enable_thinking": False}

            completion = client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                max_tokens=min(max_tokens, 2048),
                extra_body=extra_body if extra_body else None,
                stream=False,
            )

            raw_choice = completion.choices[0]
            raw_content = raw_choice.message.content or ""
            reasoning = getattr(raw_choice.message, "reasoning_content", None)

            latency = time.time() - start_time
            logger.info("NVIDIA model invocation completed in %.3f seconds", latency)

            parsed_json = self._extract_json(raw_content)

            return {
                "parsed_data": parsed_json,
                "raw_text": raw_content,
                "reasoning": reasoning,
                "latency_seconds": round(latency, 4),
                "model": self.model,
            }

        except AuthenticationError as e:
            logger.error("NVIDIA Authentication failed. Verify NVIDIA_API_KEY.")
            raise ValueError("Authentication with NVIDIA API failed. Please check your NVIDIA_API_KEY.") from e
        except RateLimitError as e:
            logger.error("NVIDIA Rate Limit exceeded: %s", e)
            raise RuntimeError("NVIDIA API rate limit exceeded. Please retry shortly.") from e
        except APIError as e:
            logger.error("NVIDIA API error: %s", e)
            raise RuntimeError(f"NVIDIA API error: {e.message}") from e
        except Exception as e:
            logger.error("Unexpected error during NVIDIA model call: %s", e)
            raise RuntimeError(f"Failed to communicate with NVIDIA LLM: {str(e)}") from e

    def _extract_json(self, text: str) -> Dict[str, Any]:
        """Extract and clean valid JSON from model response text."""
        cleaned = text.strip()

        # Handle markdown code fences ```json ... ```
        json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned)
        if json_match:
            candidate = json_match.group(1).strip()
            try:
                return json.loads(candidate)
            except json.JSONDecodeError:
                pass

        # Handle raw json string
        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            pass

        # Try to find the first '{' and last '}'
        start_idx = cleaned.find("{")
        end_idx = cleaned.rfind("}")
        if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
            sub = cleaned[start_idx : end_idx + 1]
            try:
                return json.loads(sub)
            except json.JSONDecodeError as e:
                logger.error("Failed to parse JSON substring: %s", e)

        raise ValueError(
            f"Could not parse valid JSON from NVIDIA model output: {cleaned[:200]}..."
        )


# Global instance
nvidia_client = NVIDIAClient()