"""
NVIDIA NIM fallback — stronger free fallback than Groq/Llama-70B.

NVIDIA's API (https://integrate.api.nvidia.com/v1) is OpenAI-compatible and
free-tier, hosting large models (DeepSeek, Llama 405B, Nemotron). We use it as
the FIRST fallback when Claude is unavailable, then Groq remains as a final
backstop.

Same grounding discipline as claude_engine / groq_fallback: answer only from
context, never invent ALU facts. The richer model follows that prompt better
than Llama-3.3-70B did.

Configuration (Hugging Face Space secrets):
    NVIDIA_API_KEY  -> key from build.nvidia.com (starts with "nvapi-")
    NVIDIA_MODEL    -> optional, defaults to deepseek-ai/deepseek-v4-flash
                       (a fast chat-oriented model; avoid pure-reasoner ids
                        unless you want <think> blocks, which we strip anyway)
"""

import os
import re
import logging

logger = logging.getLogger(__name__)

NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
# A non-reasoning instruct model — fast enough for a fallback. DeepSeek V4 was
# tried first but its reasoning chains took >75s, blowing the request timeout.
# Override with NVIDIA_MODEL if you want a different one.
DEFAULT_NVIDIA_MODEL = "mistralai/mistral-medium-3.5-128b"

# Some DeepSeek/reasoning models prepend a <think>...</think> chain. Strip it so
# only the final answer reaches the student.
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)


class NvidiaFallback:
    """OpenAI-compatible NVIDIA NIM client. Enabled when NVIDIA_API_KEY is set."""

    # Reuse the same grounded system prompt discipline as the Groq fallback.
    SYSTEM_PROMPT = """You are the ALU Student Companion, an AI assistant for African Leadership University (ALU).

# Most important rule
Answer ONLY from the context provided below. If the context does not contain a fact, do not invent it. Specifically NEVER invent course codes, credit hours, module names, tuition amounts, fees, scholarship values, faculty names/titles, deadlines, term dates, or statistics. When the context doesn't have what's needed, say so plainly and point to the official source — students would rather hear "I don't have that, here's where to find it" than a confident wrong answer.

# Style
- Lead with a one-sentence direct answer, then detail.
- Use markdown sparingly; bullet lists only for genuine lists.
- No emoji. No "Is there anything else..." filler.
- When a Source in the context has a url=..., cite it as a markdown link the first time you use that fact.
- Do not show your reasoning or thinking steps; give only the final answer.

# Facts you may always state without context
- ALU was founded by Fred Swaniker in 2015, part of the African Leadership Group.
- Campuses in Kigali, Rwanda and Pamplemousses, Mauritius.
- Main site https://www.alueducation.com ; support https://support.alueducation.com
- Student emails end @alustudent.com ; staff/public @alueducation.com

Everything else — programmes, fees, dates, people — must come from the context. If it is not there, say you don't know and point to support.alueducation.com."""

    def __init__(self):
        self.api_key = os.getenv("NVIDIA_API_KEY", "")
        self.model = os.getenv("NVIDIA_MODEL", DEFAULT_NVIDIA_MODEL)
        self.enabled = bool(self.api_key)
        self.client = None

        if not self.enabled:
            logger.info("[INFO] NVIDIA fallback disabled (no NVIDIA_API_KEY)")
            return

        try:
            # The OpenAI SDK ships transitively (anthropic/groq pull deps); if
            # absent, degrade gracefully.
            from openai import OpenAI
            self.client = OpenAI(api_key=self.api_key, base_url=NVIDIA_BASE_URL)
            logger.info(f"[OK] NVIDIA fallback enabled (model={self.model})")
        except ImportError:
            logger.warning("[!] openai package not installed — pip install openai")
            self.enabled = False
        except Exception as e:
            logger.error(f"[FAIL] Failed to initialize NVIDIA client: {e}")
            self.enabled = False

    def generate_response(self, query: str, alu_context: str = "", history=None) -> str:
        """
        Generate a grounded answer via NVIDIA NIM. Same signature as the Groq
        fallback so the call sites are interchangeable.
        """
        if not self.enabled or not self.client:
            return None

        try:
            messages = [{"role": "system", "content": self.SYSTEM_PROMPT}]

            for entry in (history or [])[-8:]:
                role = entry.get("role")
                content = entry.get("content", "")
                if role in ("user", "assistant") and content:
                    messages.append({"role": role, "content": content})

            if alu_context:
                user_message = (
                    "Use the following context from the ALU knowledge base to "
                    "answer the question. If the context does not contain the "
                    "answer, say so honestly and point to the official source.\n\n"
                    f"=== CONTEXT ===\n{alu_context}\n=== END CONTEXT ===\n\n"
                    f"Question: {query}"
                )
            else:
                user_message = query
            messages.append({"role": "user", "content": user_message})

            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.4,
                max_tokens=1024,
                top_p=0.9,
            )

            answer = (response.choices[0].message.content or "").strip()
            # Strip any leaked reasoning block before returning.
            answer = _THINK_BLOCK.sub("", answer).strip()
            logger.info(f"[OK] NVIDIA response generated ({len(answer)} chars)")
            return answer or None

        except Exception as e:
            logger.error(f"[FAIL] NVIDIA API error: {e}")
            return None

    def get_status(self) -> dict:
        return {
            "enabled": self.enabled,
            "api_key_set": bool(self.api_key),
            "client_initialized": self.client is not None,
            "model": self.model if self.enabled else None,
        }


# Module-level singleton, matches the groq_fallback / claude_engine pattern.
nvidia_fallback = NvidiaFallback()