"""
Indian Legal AI Assistant — Hugging Face Spaces entry point.

Loads pre-built vector DB artifacts from ./data/ and launches the Gradio UI.
No PDF processing, no package installation, no Kaggle/Colab paths.
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from pathlib import Path
from threading import Lock
import inspect
import json
import os
import pickle
import re
import time
from typing import Dict, Iterable, List, Tuple

import gradio as gr
import numpy as np
import requests

# ── Artifact location ─────────────────────────────────────────────────────────
# All four files (chunks_data.pkl, embeddings.npy, legal_faiss_index.index,
# manifest.json) must be uploaded to the Space's data/ folder.
ARTIFACT_DIR = Path(os.getenv("ARTIFACT_DIR", "./data"))

# ── Model & inference settings ────────────────────────────────────────────────
EMBEDDING_MODEL_NAME    = "BAAI/bge-large-en-v1.5"
TOP_K                   = 3
MIN_RELEVANCE_SCORE     = 0.30
MAX_QUERY_CHARS         = 1800
MAX_HISTORY_TURNS       = 4
MAX_CONTEXT_WORDS       = 500
MAX_OUTPUT_TOKENS       = 1200
BATCH_SIZE              = 32
REQUESTS_PER_MINUTE     = 8
MIN_REQUEST_INTERVAL    = 2.5
QUEUE_MAX_SIZE          = 16
DEFAULT_CONCURRENCY     = 2

DEFAULT_PROVIDER  = "openai"
ALLOWED_PROVIDERS = ("openai", "groq", "gemini")
PROVIDER_SPECS = {
    "openai": {"label": "OpenAI",          "base_url": "https://api.openai.com/v1"},
    "groq":   {"label": "Groq",            "base_url": "https://api.groq.com/openai/v1"},
    "gemini": {"label": "Gemini (Google)", "base_url": "https://generativelanguage.googleapis.com/v1beta/openai"},
}


# ── Exceptions ────────────────────────────────────────────────────────────────
class UserFacingError(RuntimeError):
    """Safe exception whose message is shown directly in the UI."""


# ── Helpers ───────────────────────────────────────────────────────────────────
def sanitize_text(value: str | None, max_chars: int) -> str:
    text = (value or "").replace("\x00", " ").strip()
    text = re.sub(r"\r\n?", "\n", text)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text[:max_chars].rstrip() if len(text) > max_chars else text


def redact_sensitive(text: str) -> str:
    for pattern in [
        r"sk-[A-Za-z0-9_-]{12,}",
        r"hf_[A-Za-z0-9]{12,}",
        r"gsk_[A-Za-z0-9_-]{12,}",
        r"AIza[A-Za-z0-9_-]{35}",
    ]:
        text = re.sub(pattern, "[REDACTED]", text)
    return text


def safe_error_message(exc: Exception) -> str:
    if isinstance(exc, UserFacingError):
        return str(exc)
    detail = redact_sensitive(str(exc)).strip()
    return f"Something went wrong. Detail: {detail[:180]}" if detail else "Something went wrong."


def normalize_rows(values: np.ndarray) -> np.ndarray:
    arr = np.asarray(values, dtype=np.float32)
    if arr.ndim == 1:
        arr = arr.reshape(1, -1)
    norms = np.linalg.norm(arr, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return arr / norms


def detect_intent(query: str) -> str:
    q = query.lower().strip()
    legal_indicators = [
        # Acts and codes
        "section", "act", "law", "ipc", "bns", "bnss", "bnns", "pocso", "ndps",
        "crpc", "it act", "constitution", "article", "schedule", "clause",
        # People and status
        "accused", "victim", "complainant", "witness", "juvenile", "minor",
        "child", "abscond", "detain", "detention", "custody", "prisoner",
        # Proceedings
        "court", "tribunal", "magistrate", "judge", "appeal", "petition",
        "bail", "arrest", "warrant", "summons", "charge", "trial", "acquit",
        # Liability and culpability
        "criminal", "civil", "liable", "liability", "culpable", "negligence",
        "intent", "mens rea", "actus reus", "abetment", "conspiracy",
        # Outcomes
        "penalty", "offense", "offence", "crime", "rights", "punishment",
        "sentence", "fine", "imprisonment", "provisions", "legal", "illegal",
        "exempt", "immunity", "pardon", "remission",
    ]
    if any(ind in q for ind in legal_indicators):
        return "LEGAL"
    if re.search(r"\b(section|sec|article|art)\s*\d+", q):
        return "LEGAL"
    if len(q.split()) <= 3:
        return "CONVERSATIONAL"
    return "CONVERSATIONAL"


def history_to_messages(history: list | None, max_turns: int, max_chars: int) -> list[dict]:
    if not history:
        return []
    normalized: list[dict] = []
    for item in history:
        if isinstance(item, dict):
            role = str(item.get("role", "")).strip().lower()
            content = item.get("content")
            if role in {"user", "assistant"} and isinstance(content, str):
                c = sanitize_text(content, max_chars)
                if c:
                    normalized.append({"role": role, "content": c})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            u = sanitize_text(str(item[0] or ""), max_chars)
            a = sanitize_text(str(item[1] or ""), max_chars)
            if u:
                normalized.append({"role": "user",      "content": u})
            if a:
                normalized.append({"role": "assistant", "content": a})
    return normalized[-(max_turns * 2):]


# ── Rate limiter ──────────────────────────────────────────────────────────────
@dataclass(frozen=True)
class RateLimitPolicy:
    requests_per_minute: int
    min_interval_seconds: float


class SessionRateLimiter:
    def __init__(self, policy: RateLimitPolicy) -> None:
        self.policy = policy
        self._events: dict[str, deque[float]] = {}
        self._lock = Lock()

    def check(self, session_id: str) -> None:
        if not session_id:
            return
        now = time.monotonic()
        with self._lock:
            bucket = self._events.setdefault(session_id, deque())
            while bucket and now - bucket[0] > 60:
                bucket.popleft()
            if bucket and now - bucket[-1] < self.policy.min_interval_seconds:
                wait = self.policy.min_interval_seconds - (now - bucket[-1])
                raise UserFacingError(f"Please wait {wait:.1f}s before sending another request.")
            if len(bucket) >= self.policy.requests_per_minute:
                raise UserFacingError("Too many requests. Pause for a minute and try again.")
            bucket.append(now)


# ── Embedder ──────────────────────────────────────────────────────────────────
class SentenceTransformerEmbedder:
    def __init__(self, model_name: str) -> None:
        from sentence_transformers import SentenceTransformer  # type: ignore
        # HF Spaces CPU — no GPU available on free tier
        self.model = SentenceTransformer(model_name, device="cpu")

    def encode(self, texts: list[str]) -> np.ndarray:
        return np.asarray(
            self.model.encode(texts, show_progress_bar=False,
                              batch_size=BATCH_SIZE, convert_to_numpy=True),
            dtype=np.float32,
        )


# ── Vector index ──────────────────────────────────────────────────────────────
class NumpyVectorIndex:
    def __init__(self, embeddings: np.ndarray) -> None:
        self.embeddings = normalize_rows(embeddings)

    def search(self, query_embedding: np.ndarray, top_k: int):
        scores = np.matmul(normalize_rows(query_embedding), self.embeddings.T)
        indices = np.argsort(-scores, axis=1)[:, :top_k]
        return np.take_along_axis(scores, indices, axis=1).astype(np.float32), indices.astype(np.int64)


# ── Retriever ─────────────────────────────────────────────────────────────────
class OptimizedLegalRetriever:
    QUERY_PREFIX = "Represent this query for retrieving relevant legal passages: "
    EXPANSIONS = {
        r"\bBNS\b":    "Bharatiya Nyaya Sanhita",
        r"\bBNNS\b":   "Bharatiya Nagarik Suraksha Sanhita",
        r"\bCrPC\b":   "Criminal Procedure Code",
        r"\bPOCSO\b":  "Protection of Children from Sexual Offences",
        r"\bIT Act\b": "Information Technology Act",
        r"\bSec\.?\b": "Section",
        r"\bArt\.?\b": "Article",
    }

    def __init__(self, embedding_model, vector_index, all_chunks, chunk_metadata) -> None:
        self.embedding_model = embedding_model
        self.index           = vector_index
        self.all_chunks      = all_chunks
        self.chunk_metadata  = chunk_metadata

    def preprocess(self, query: str) -> str:
        for abbr, full in self.EXPANSIONS.items():
            query = re.sub(abbr, full, query, flags=re.IGNORECASE)
        return query.strip()

    def retrieve(self, query: str, top_k: int = 3) -> list[dict]:
        processed = self.preprocess(query)
        qe = normalize_rows(self.embedding_model.encode([self.QUERY_PREFIX + processed]))
        scores, indices = self.index.search(qe, top_k * 2)
        keywords = set(processed.lower().split())
        candidates = []
        for score, idx in zip(scores[0], indices[0]):
            idx = int(idx)
            if 0 <= idx < len(self.all_chunks):
                chunk = self.all_chunks[idx]
                overlap = len(keywords & set(chunk.lower().split())) / max(len(keywords), 1)
                candidates.append({
                    "chunk":    chunk,
                    "metadata": self.chunk_metadata[idx],
                    "score":    float(score) + overlap * 0.15,
                    "index":    idx,
                })
        candidates.sort(key=lambda x: x["score"], reverse=True)
        return candidates[:top_k]


def format_retrieved_context(chunks_data: list[dict], max_words: int) -> str:
    parts = []
    for i, data in enumerate(chunks_data, 1):
        meta  = data["metadata"]
        chunk = data["chunk"]
        words = chunk.split()
        if len(words) > max_words:
            chunk = " ".join(words[:max_words]) + " ..."
        source = meta.get("source_file", "Unknown").replace(".pdf", "")
        sec    = f" - {meta.get('section_type','Section')} {meta['section_number']}" \
                 if "section_number" in meta else ""
        parts.append(f"[Reference {i}: {source}{sec}]\n{chunk}")
    return "\n\n---\n\n".join(parts)


# ── Load artifacts ────────────────────────────────────────────────────────────
def load_runtime():
    manifest_path  = ARTIFACT_DIR / "manifest.json"
    chunks_path    = ARTIFACT_DIR / "chunks_data.pkl"
    embeddings_path= ARTIFACT_DIR / "embeddings.npy"

    for p in (manifest_path, chunks_path, embeddings_path):
        if not p.exists():
            raise UserFacingError(
                f"Artifact file not found: {p}\n"
                "Make sure all 4 files are uploaded to the Space's data/ folder."
            )

    with manifest_path.open("r", encoding="utf-8") as f:
        manifest = json.load(f)
    with chunks_path.open("rb") as f:
        data = pickle.load(f)
    embeddings = np.load(embeddings_path).astype(np.float32)

    # Try FAISS first, fall back to numpy
    faiss_path = ARTIFACT_DIR / "legal_faiss_index.index"
    index = None
    if faiss_path.exists():
        try:
            import faiss  # type: ignore
            index = faiss.read_index(str(faiss_path))
            print("FAISS index loaded.")
        except Exception as e:
            print(f"FAISS load failed ({e}), using numpy search.")

    if index is None:
        index = NumpyVectorIndex(embeddings)

    embedder  = SentenceTransformerEmbedder(EMBEDDING_MODEL_NAME)
    retriever = OptimizedLegalRetriever(
        embedding_model=embedder,
        vector_index=index,
        all_chunks=list(data["all_chunks"]),
        chunk_metadata=list(data["chunk_metadata"]),
    )
    print(f"Runtime ready — {manifest.get('vector_count','?')} chunks from "
          f"{manifest.get('source_pdf_count','?')} PDFs.")
    return retriever, manifest


# ── Provider streaming ────────────────────────────────────────────────────────
class ProviderClient:
    def __init__(self) -> None:
        self.session = requests.Session()

    def stream_chat(
        self,
        provider_name: str,
        model_name: str,
        api_key: str,
        messages: list[dict],
        *,
        temperature: float,
        max_tokens: int,
    ) -> Iterable[str]:
        key = provider_name.strip().lower()
        if key not in ALLOWED_PROVIDERS:
            raise UserFacingError(f"Unknown provider '{provider_name}'.")
        api_key = sanitize_text(api_key, 500)
        if not api_key:
            raise UserFacingError("API key is required. Open ⚙️ Settings and paste your key.")
        model = sanitize_text(model_name, 200)
        if not model:
            raise UserFacingError("Model name is required. Open ⚙️ Settings and enter a model.")

        spec = PROVIDER_SPECS[key]
        try:
            with self.session.post(
                f"{spec['base_url']}/chat/completions",
                headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
                json={"model": model, "messages": messages,
                      "temperature": temperature, "max_tokens": max_tokens, "stream": True},
                timeout=(10, 180),
                stream=True,
            ) as resp:
                if resp.status_code >= 400:
                    raw = ""
                    try:
                        ep = resp.json()
                        if isinstance(ep, list):
                            ep = ep[0] if ep else {}
                        if isinstance(ep, dict):
                            err = ep.get("error") or {}
                            raw = (err.get("message") if isinstance(err, dict) else str(err)) \
                                  or ep.get("message") or ""
                        else:
                            raw = str(ep)
                    except ValueError:
                        raw = resp.text[:300]
                    raw = redact_sensitive(str(raw)).strip()

                    s = resp.status_code
                    if s == 401:
                        msg = f"❌ Invalid API key for {spec['label']}. Check the key has no extra spaces."
                    elif s == 403:
                        msg = f"❌ Access denied ({spec['label']} 403). Your plan may not include this model."
                    elif s == 404:
                        msg = f"❌ Model '{model}' not found on {spec['label']} (404). Check the model name."
                    elif s == 429:
                        msg = f"⏳ Rate limit exceeded on {spec['label']} (429). Wait a moment and retry."
                    elif s == 500:
                        msg = f"🔥 {spec['label']} server error (500). Wait a minute and retry."
                    elif s in (502, 503, 504):
                        msg = f"🔌 {spec['label']} temporarily unavailable ({s}). Retry in a few seconds."
                    else:
                        msg = f"❌ {spec['label']} rejected the request (HTTP {s})."
                    if raw and raw.lower() not in msg.lower():
                        msg += f"\n\nProvider said: {raw[:300]}"
                    raise UserFacingError(msg)

                for raw_line in resp.iter_lines(decode_unicode=True):
                    if not raw_line:
                        continue
                    line = raw_line.strip()
                    if not line.startswith("data:"):
                        continue
                    data = line[5:].strip()
                    if data == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                    except json.JSONDecodeError:
                        continue
                    if isinstance(chunk, list):
                        chunk = chunk[0] if chunk else {}
                    if not isinstance(chunk, dict):
                        continue
                    choices = chunk.get("choices")
                    if not choices or not isinstance(choices, list):
                        continue
                    fc = choices[0]
                    if not isinstance(fc, dict):
                        continue
                    delta = fc.get("delta") or {}
                    if not isinstance(delta, dict):
                        continue
                    token = delta.get("content")
                    if token:
                        yield token

        except requests.RequestException as exc:
            raise UserFacingError(
                f"{spec['label']} request failed. Check your internet connection."
            ) from exc


# ── Bootstrap ─────────────────────────────────────────────────────────────────
RETRIEVER: OptimizedLegalRetriever | None = None
MANIFEST:  dict | None = None
LOAD_ERROR: Exception | None = None
CLIENT = ProviderClient()
RATE_LIMITER = SessionRateLimiter(
    RateLimitPolicy(requests_per_minute=REQUESTS_PER_MINUTE,
                    min_interval_seconds=MIN_REQUEST_INTERVAL)
)

try:
    RETRIEVER, MANIFEST = load_runtime()
except Exception as _e:
    LOAD_ERROR = _e
    print(f"ERROR loading runtime: {_e}")


# ── Message builders ──────────────────────────────────────────────────────────
def build_conversational_messages(message: str, history: list) -> list[dict]:
    return [
        {"role": "system",
         "content": "You are a helpful Indian legal AI assistant. "
                    "For casual conversation reply briefly in 2-3 sentences."},
        *history_to_messages(history, MAX_HISTORY_TURNS, MAX_QUERY_CHARS),
        {"role": "user", "content": message},
    ]


def build_legal_messages(query: str, chunks_data: list[dict]) -> list[dict]:
    context = format_retrieved_context(chunks_data, MAX_CONTEXT_WORDS)
    return [
        {"role": "system",
         "content": "You are an expert on Indian law. Answer only from the provided legal "
                    "references. Cite sections, acts, and provisions where possible. "
                    "If the context is incomplete, say so plainly."},
        {"role": "user",
         "content": f"Legal References:\n{context}\n\nQuestion: {query}\n\n"
                    "Provide a grounded answer based only on the references above."},
    ]


def stream_reply(provider: str, model: str, key: str,
                 messages: list[dict], *, temperature: float):
    full = ""
    for token in CLIENT.stream_chat(provider, model, key, messages,
                                    temperature=temperature, max_tokens=MAX_OUTPUT_TOKENS):
        full += token
        yield full.strip()


def test_connection(provider: str, model: str, key: str) -> str:
    try:
        msgs = [{"role": "system",
                 "content": "Reply in one short sentence confirming the API works."},
                {"role": "user", "content": "Say: API connection successful."}]
        sample = ""
        for p in stream_reply(provider, model, key, msgs, temperature=0.1):
            sample = p
            if len(sample) >= 120:
                break
        return f"✅ Connection OK. Reply: {sample[:160]}" if sample \
               else "Connection reached the provider but returned no text."
    except Exception as exc:
        return safe_error_message(exc)


# ── Main chat handler ─────────────────────────────────────────────────────────
def chat(message: str, history: list,
         provider_name: str, model_name: str, user_api_key: str,
         request: gr.Request):
    try:
        sid = getattr(request, "session_hash", None) or "anon"
        RATE_LIMITER.check(sid)

        msg = sanitize_text(message, MAX_QUERY_CHARS)
        if not msg:
            raise UserFacingError("Please enter a question.")

        intent = detect_intent(msg)

        if intent == "CONVERSATIONAL":
            yield "Thinking..."
            for p in stream_reply(provider_name, model_name, user_api_key,
                                   build_conversational_messages(msg, history),
                                   temperature=0.6):
                yield p
            return

        if RETRIEVER is None:
            err = safe_error_message(LOAD_ERROR) if LOAD_ERROR else "Vector DB not loaded."
            yield f"⚠️ The vector database is not ready.\n\n{err}"
            return

        chunks = RETRIEVER.retrieve(msg, top_k=TOP_K)
        if not chunks or all(c["score"] < MIN_RELEVANCE_SCORE for c in chunks):
            yield ("No relevant legal documents found for this question. "
                   "Try being more specific — mention the act, section, or offence.")
            return

        # Build footnote before streaming so it always appears
        sources_line = "\n\n---\n*Sources: " + " · ".join(
            dict.fromkeys(
                c["metadata"].get("source_file", "?").replace(".pdf", "")
                for c in chunks
            )
        ) + "*"

        yield "⚖️ Searching legal documents..."

        current = ""
        try:
            for p in stream_reply(provider_name, model_name, user_api_key,
                                   build_legal_messages(msg, chunks),
                                   temperature=0.1):
                current = p
                yield current
        except Exception as stream_exc:
            err_msg = safe_error_message(stream_exc)
            yield (current.strip() + "\n\n" + err_msg if current else err_msg) + sources_line
            return

        yield current.strip() + sources_line

    except Exception as exc:
        yield safe_error_message(exc)


# ── Gradio UI ─────────────────────────────────────────────────────────────────
def _qkw(blocks, **kw):
    sup = inspect.signature(blocks.queue).parameters
    return {k: v for k, v in kw.items() if k in sup}


def build_demo() -> gr.Blocks:

    PROVIDER_HELP = {
        "openai": "e.g. gpt-4o · gpt-4o-mini · gpt-3.5-turbo",
        "groq":   "e.g. llama3-70b-8192 · mixtral-8x7b-32768 · gemma2-9b-it",
        "gemini": "e.g. gemini-1.5-flash · gemini-1.5-pro",
    }

    with gr.Blocks(title="Indian Legal AI Assistant") as demo:

        # ── Title ─────────────────────────────────────────────────────────────
        gr.Markdown(
            "# 🏛️ Indian Legal AI Assistant\n"
            "Ask questions about Indian legal documents "
            "(POCSO, NDPS, IT Act, BNS, BNSS, and 40+ other acts)"
        )

        # ── Settings accordion ─────────────────────────────────────────────────
        with gr.Accordion("⚙️  Model & API Settings — click to expand", open=False):
            provider = gr.Dropdown(
                choices=list(ALLOWED_PROVIDERS),
                value=DEFAULT_PROVIDER,
                label="Provider",
            )
            model_hint = gr.Markdown(
                f"<small>{PROVIDER_HELP[DEFAULT_PROVIDER]}</small>"
            )
            model_name = gr.Textbox(
                label="Model name",
                placeholder="Type the exact model string for your provider",
            )
            user_api_key = gr.Textbox(
                type="password",
                label="API Key",
                placeholder="Paste your key here — never stored or logged",
            )
            test_btn = gr.Button("🔍 Test Connection", variant="primary")
            key_status = gr.Markdown(
                "<small>Enter key + model, then click Test Connection.</small>"
            )

        # ── Chat ───────────────────────────────────────────────────────────────
        ci = gr.ChatInterface(
            fn=chat,
            additional_inputs=[provider, model_name, user_api_key],
            additional_inputs_accordion=gr.Accordion(visible=False, open=False),
            title=None,
            description=None,
            submit_btn="Send ➤",
            stop_btn="Stop",
            show_progress="minimal",
            save_history=False,
            flagging_mode="never",
            api_name=False,
            examples=[
                ["Hello, how can you help me?"],
                ["What does POCSO stand for?"],
                ["What is the punishment for child abuse under POCSO Act?"],
                ["Explain Section 27 of NDPS Act"],
                ["What are the provisions for cybercrime in IT Act?"],
                ["Thank you for your help!"],
            ],
        )

        # ── Event wiring ───────────────────────────────────────────────────────
        provider.change(
            fn=lambda p: f"<small>{PROVIDER_HELP.get(p, '')}</small>",
            inputs=provider,
            outputs=model_hint,
        )
        test_btn.click(
            fn=test_connection,
            inputs=[provider, model_name, user_api_key],
            outputs=key_status,
        )

    demo.queue(
        **_qkw(demo, api_open=False, max_size=QUEUE_MAX_SIZE,
               default_concurrency_limit=DEFAULT_CONCURRENCY)
    )
    return demo
demo = build_demo()

if __name__ == "__main__":
    demo.launch()