#!/usr/bin/env python3
"""
Indian Legal AI Assistant with integrated India Code lookup.

CPU-only Hugging Face Spaces app.

Features:
- Uses India Code as the legal web source: https://www.indiacode.nic.in/
- Reads India Code HTML pages.
- Reads text-based PDFs using pypdf.
- Falls back to OCR for scanned/image PDFs using PyMuPDF + Tesseract.
- Streams output for faster perceived response time.
- Automatically trims prompt/context to avoid context-window overflow.
- Supports follow-up questions using recent chat history.
- Footer: Developed by Rohan R.
"""

import os
import re
import io
import html
import time
import traceback
from functools import lru_cache
from urllib.parse import urljoin, urlparse, quote_plus, urldefrag, parse_qs

import requests
from bs4 import BeautifulSoup
from PIL import Image
from pypdf import PdfReader
import fitz  # PyMuPDF
import pytesseract

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama


# =================================================
# Model configuration
# =================================================

MODEL_REPO = "invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF"
MODEL_FILE = "llama-3.2-1b-instruct.Q4_K_M.gguf"

# Balanced CPU-only configuration.
# If the Space is too slow or memory-limited, reduce N_CTX to 3072 or 2048.
N_CTX = int(os.getenv("N_CTX", "4096"))
N_THREADS = int(os.getenv("N_THREADS", "2"))
N_THREADS_BATCH = int(os.getenv("N_THREADS_BATCH", "2"))
N_BATCH = int(os.getenv("N_BATCH", "512"))
N_GPU_LAYERS = 0

# Detailed but bounded output.
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "1200"))
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.25"))
TOP_P = float(os.getenv("TOP_P", "0.9"))

MODEL_CACHE_DIR = os.getenv("MODEL_CACHE_DIR", "./models")

llm = None


# =================================================
# India Code / extraction configuration
# =================================================

INDIACODE_HOME = "https://www.indiacode.nic.in/"
ALLOWED_DOMAINS = {
    "indiacode.nic.in",
    "www.indiacode.nic.in",
}

REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "12"))

# Retrieval limits. Increasing these may improve detail but will slow the app.
MAX_DISCOVERY_RESULTS = int(os.getenv("MAX_DISCOVERY_RESULTS", "12"))
MAX_CRAWL_PAGES = int(os.getenv("MAX_CRAWL_PAGES", "24"))
MAX_CONTEXT_DOCS = int(os.getenv("MAX_CONTEXT_DOCS", "6"))

# PDF text extraction.
MAX_PDF_TEXT_PAGES = int(os.getenv("MAX_PDF_TEXT_PAGES", "8"))

# OCR fallback. Kept small for CPU Spaces.
MAX_OCR_PAGES = int(os.getenv("MAX_OCR_PAGES", "3"))
OCR_DPI_SCALE = float(os.getenv("OCR_DPI_SCALE", "1.4"))
MIN_PDF_TEXT_CHARS_BEFORE_OCR = int(os.getenv("MIN_PDF_TEXT_CHARS_BEFORE_OCR", "250"))

# Context trimming.
MAX_TEXT_PER_DOC = int(os.getenv("MAX_TEXT_PER_DOC", "3000"))
PROMPT_SAFETY_MARGIN = int(os.getenv("PROMPT_SAFETY_MARGIN", "160"))

HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/120.0 Safari/537.36"
    )
}


# =================================================
# Basic helpers
# =================================================

def clean_text(text):
    """Normalize whitespace and decode HTML entities."""
    if not text:
        return ""

    text = html.unescape(text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def normalize_url(url, base=INDIACODE_HOME):
    """Resolve relative URLs, remove fragments, and normalize."""
    if not url:
        return ""

    url = urljoin(base, url)
    url, _fragment = urldefrag(url)
    return url.strip()


def is_indiacode_url(url):
    """Allow only India Code URLs."""
    try:
        domain = urlparse(url).netloc.lower()
        return domain in ALLOWED_DOMAINS
    except Exception:
        return False


def looks_like_pdf_url(url):
    """Detect PDF-ish URLs."""
    return ".pdf" in url.lower()


def query_terms(query):
    """Extract useful terms for lightweight relevance scoring."""
    stopwords = {
        "the", "a", "an", "and", "or", "of", "in", "on", "to", "for", "with",
        "under", "section", "sections", "act", "acts", "law", "laws", "what",
        "is", "are", "explain", "about", "current", "latest", "india", "indian",
        "tell", "me", "please", "does", "do", "give", "details", "detailed",
    }

    terms = re.findall(r"[a-zA-Z0-9]+", query.lower())
    return [t for t in terms if len(t) >= 3 and t not in stopwords]


def score_text_against_query(text, query):
    """Simple lexical relevance score."""
    text_l = (text or "").lower()
    terms = query_terms(query)

    if not terms:
        return 0

    score = 0

    for term in terms:
        count = text_l.count(term)
        if count:
            score += min(count, 5)

    exact_query = clean_text(query).lower()
    if exact_query and exact_query in text_l:
        score += 10

    return score


def make_snippet(text, query, max_chars=MAX_TEXT_PER_DOC):
    """Create a focused snippet around the first query-term hit."""
    text = clean_text(text)

    if not text:
        return ""

    lower = text.lower()
    terms = query_terms(query)

    first_hit = None

    for term in terms:
        idx = lower.find(term)
        if idx != -1:
            first_hit = idx
            break

    if first_hit is None:
        return text[:max_chars]

    start = max(first_hit - 400, 0)
    end = min(start + max_chars, len(text))
    return text[start:end]


def safe_get(url):
    """HTTP GET with basic error handling."""
    try:
        response = requests.get(url, headers=HEADERS, timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        return response
    except Exception as error:
        print(f"GET failed: {url} :: {error}")
        return None


# =================================================
# HTML extraction
# =================================================

def extract_links_from_html(html_text, base_url):
    """Extract India Code links from HTML."""
    links = []

    try:
        soup = BeautifulSoup(html_text, "html.parser")

        for a in soup.find_all("a", href=True):
            href = a.get("href", "")
            text = clean_text(a.get_text(" "))
            url = normalize_url(href, base_url)

            if is_indiacode_url(url):
                links.append(
                    {
                        "url": url,
                        "anchor": text,
                    }
                )

    except Exception as error:
        print(f"Link extraction failed for {base_url}: {error}")

    return links


def extract_text_from_html(html_text):
    """Extract readable text from an HTML page."""
    try:
        soup = BeautifulSoup(html_text, "html.parser")

        for tag in soup(["script", "style", "nav", "footer", "header", "aside", "form"]):
            tag.decompose()

        parts = []

        if soup.title:
            parts.append(clean_text(soup.title.get_text(" ")))

        for tag in soup.find_all(
            ["h1", "h2", "h3", "h4", "p", "li", "td", "th"],
            limit=350,
        ):
            text = clean_text(tag.get_text(" "))
            if len(text) >= 20:
                parts.append(text)

        return clean_text(" ".join(parts))

    except Exception as error:
        print(f"HTML extraction failed: {error}")
        return ""


# =================================================
# PDF extraction with OCR fallback
# =================================================

def extract_pdf_text_with_pypdf(pdf_bytes, max_pages=MAX_PDF_TEXT_PAGES):
    """Extract embedded text from text-based PDFs."""
    try:
        reader = PdfReader(io.BytesIO(pdf_bytes))
        parts = []

        pages_to_read = min(len(reader.pages), max_pages)

        for page_index in range(pages_to_read):
            try:
                text = reader.pages[page_index].extract_text() or ""
                text = clean_text(text)

                if text:
                    parts.append(f"[PDF text page {page_index + 1}] {text}")

            except Exception as page_error:
                print(f"pypdf page extraction failed: {page_error}")

        return clean_text(" ".join(parts))

    except Exception as error:
        print(f"pypdf extraction failed: {error}")
        return ""


def ocr_pdf_with_tesseract(pdf_bytes, max_pages=MAX_OCR_PAGES):
    """
    Lightweight OCR fallback for scanned PDFs.

    Uses:
    - PyMuPDF to render PDF pages to images.
    - Tesseract OCR via pytesseract.

    OCR is intentionally limited because CPU Spaces are resource-constrained.
    """
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        parts = []

        pages_to_read = min(len(doc), max_pages)

        for page_index in range(pages_to_read):
            try:
                page = doc.load_page(page_index)
                matrix = fitz.Matrix(OCR_DPI_SCALE, OCR_DPI_SCALE)
                pix = page.get_pixmap(matrix=matrix, alpha=False)

                img = Image.frombytes(
                    "RGB",
                    [pix.width, pix.height],
                    pix.samples,
                )

                # Grayscale usually improves speed and reduces OCR workload.
                img = img.convert("L")

                text = pytesseract.image_to_string(img, lang="eng")
                text = clean_text(text)

                if text:
                    parts.append(f"[OCR PDF page {page_index + 1}] {text}")

            except Exception as page_error:
                print(f"OCR failed on page {page_index + 1}: {page_error}")

        doc.close()

        return clean_text(" ".join(parts))

    except Exception as error:
        print(f"OCR PDF extraction failed: {error}")
        return ""


def extract_pdf_text(pdf_bytes):
    """
    Try embedded PDF text first.
    If too little text is found, fall back to OCR.
    """
    text = extract_pdf_text_with_pypdf(pdf_bytes)

    if len(text) >= MIN_PDF_TEXT_CHARS_BEFORE_OCR:
        return text, "pdf-text"

    print("PDF appears scanned or has too little extractable text. Running OCR fallback...")

    ocr_text = ocr_pdf_with_tesseract(pdf_bytes)

    if ocr_text:
        if text:
            return text + "\n\n" + ocr_text, "pdf-text-plus-ocr"
        return ocr_text, "pdf-ocr"

    if text:
        return text, "pdf-text-low"

    return (
        "[PDF detected, but no readable text could be extracted. "
        "The PDF may be scanned, low quality, encrypted, or OCR failed.]",
        "pdf-unreadable",
    )


# =================================================
# India Code document fetching
# =================================================

@lru_cache(maxsize=256)
def fetch_document_text(url):
    """
    Fetch and extract text from an India Code HTML or PDF URL.
    Cached for faster follow-up questions.
    """
    if not is_indiacode_url(url):
        return {
            "url": url,
            "title": "Blocked non-India-Code URL",
            "text": "",
            "links": [],
            "type": "blocked",
        }

    response = safe_get(url)

    if response is None:
        return {
            "url": url,
            "title": "Fetch failed",
            "text": "",
            "links": [],
            "type": "failed",
        }

    ctype = response.headers.get("content-type", "").lower()

    if "application/pdf" in ctype or looks_like_pdf_url(url):
        text, pdf_type = extract_pdf_text(response.content)

        return {
            "url": url,
            "title": url.split("/")[-1] or "India Code PDF",
            "text": text,
            "links": [],
            "type": pdf_type,
        }

    html_text = response.text
    links = extract_links_from_html(html_text, url)
    text = extract_text_from_html(html_text)

    title = "India Code page"

    try:
        soup = BeautifulSoup(html_text, "html.parser")
        if soup.title:
            title = clean_text(soup.title.get_text(" "))
    except Exception:
        pass

    return {
        "url": url,
        "title": title,
        "text": text,
        "links": links,
        "type": "html",
    }


# =================================================
# India Code discovery and crawling
# =================================================

def india_code_seed_urls():
    """
    Core logical entry points on India Code.

    Includes:
    - homepage
    - Central Acts browse pages
    - repealed Acts
    - spent Acts
    """
    return [
        INDIACODE_HOME,

        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=shorttitle",
        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=actno",
        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=actyear",
        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=enactmentdate",
        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=ministry",
        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=department",

        "https://www.indiacode.nic.in/repealed-act/repealed-act.jsp",
        "https://www.indiacode.nic.in/spent-act/spent-act.jsp",
    ]


def discover_from_indiacode_home():
    """Discover additional India Code navigation links from the homepage."""
    discovered = []

    home_doc = fetch_document_text(INDIACODE_HOME)

    for link in home_doc.get("links", []):
        url = link.get("url", "")
        if is_indiacode_url(url):
            discovered.append(url)

    return discovered


def duckduckgo_site_discovery(query):
    """
    Site-restricted URL discovery only.

    This function uses a search engine only to discover India Code URLs.
    The app fetches and reads the India Code pages directly.
    """
    discovered = []

    search_query = f"site:indiacode.nic.in {query}"
    search_url = f"https://duckduckgo.com/html/?q={quote_plus(search_query)}"

    response = safe_get(search_url)

    if response is None:
        return discovered

    try:
        soup = BeautifulSoup(response.text, "html.parser")

        for a in soup.select(".result__title a"):
            url = a.get("href", "").strip()

            if "uddg=" in url:
                try:
                    parsed = urlparse(url)
                    qs = parse_qs(parsed.query)
                    if "uddg" in qs:
                        url = qs["uddg"][0]
                except Exception:
                    pass

            url = normalize_url(url)

            if is_indiacode_url(url):
                discovered.append(url)

            if len(discovered) >= MAX_DISCOVERY_RESULTS:
                break

    except Exception as error:
        print(f"Site discovery failed: {error}")

    return discovered


def relevant_link_filter(link, query):
    """
    Decide whether a link is worth crawling.
    """
    url = link.get("url", "")
    anchor = link.get("anchor", "")

    if not is_indiacode_url(url):
        return False

    url_l = url.lower()
    anchor_l = anchor.lower()

    if looks_like_pdf_url(url):
        return True

    important_patterns = [
        "/handle/",
        "/bitstream/",
        "/browse",
        "repealed-act",
        "spent-act",
        "download",
        "pdf",
    ]

    if any(pattern in url_l for pattern in important_patterns):
        return True

    terms = query_terms(query)

    if any(term in url_l or term in anchor_l for term in terms):
        return True

    return False


def crawl_indiacode_for_query(query):
    """
    Crawl India Code for relevant sources.

    Strategy:
    1. Start from logical India Code seed URLs.
    2. Add homepage-discovered links.
    3. Add site-restricted discovered India Code URLs.
    4. Fetch and score pages.
    5. Follow relevant India Code links.
    """
    seeds = []
    seeds.extend(india_code_seed_urls())
    seeds.extend(discover_from_indiacode_home())
    seeds.extend(duckduckgo_site_discovery(query))

    queue = []
    seen = set()

    for url in seeds:
        url = normalize_url(url)

        if is_indiacode_url(url) and url not in seen:
            queue.append(url)
            seen.add(url)

    visited = set()
    scored_docs = []

    while queue and len(visited) < MAX_CRAWL_PAGES:
        url = queue.pop(0)

        if url in visited:
            continue

        visited.add(url)

        doc = fetch_document_text(url)

        text = doc.get("text", "")
        title = doc.get("title", "India Code document")
        doc_type = doc.get("type", "html")

        combined = f"{title} {url} {text}"
        score = score_text_against_query(combined, query)

        if score > 0 or doc_type.startswith("pdf"):
            scored_docs.append(
                {
                    "url": url,
                    "title": title,
                    "type": doc_type,
                    "score": score,
                    "text": text,
                }
            )

        for link in doc.get("links", []):
            link_url = normalize_url(link.get("url", ""), url)

            if link_url in seen:
                continue

            if relevant_link_filter(link, query):
                queue.append(link_url)
                seen.add(link_url)

    scored_docs.sort(key=lambda item: item["score"], reverse=True)

    useful_docs = []

    for doc in scored_docs:
        if doc.get("text"):
            useful_docs.append(doc)

        if len(useful_docs) >= MAX_CONTEXT_DOCS:
            break

    return useful_docs


def build_indiacode_context(query):
    """
    Build raw India Code context.

    Final prompt sizing is handled later by build_safe_prompt().
    """
    docs = crawl_indiacode_for_query(query)

    if not docs:
        return (
            "No directly relevant readable content was retrieved from India Code for this query.",
            [],
        )

    blocks = []

    for index, doc in enumerate(docs, start=1):
        snippet = make_snippet(
            doc.get("text", ""),
            query,
            max_chars=MAX_TEXT_PER_DOC,
        )

        block = (
            f"[India Code Source {index}]\n"
            f"Title: {doc.get('title', 'India Code document')}\n"
            f"Type: {doc.get('type', 'html')}\n"
            f"URL: {doc.get('url')}\n"
            f"Relevant excerpt:\n{snippet}\n"
        )

        blocks.append(block)

    return "\n\n".join(blocks), docs


def format_sources(docs):
    """Append source links to the final response."""
    if not docs:
        return "\n\nIndia Code sources checked: No readable India Code source was retrieved."

    lines = ["\n\nIndia Code sources checked:"]

    for index, doc in enumerate(docs, start=1):
        title = doc.get("title", "India Code document")
        url = doc.get("url", "")
        dtype = doc.get("type", "html")
        lines.append(f"{index}. {title} [{dtype}]\n   {url}")

    return "\n".join(lines)


# =================================================
# Model loading
# =================================================

def load_model():
    """Download and load the GGUF model once."""
    global llm

    if llm is not None:
        return llm

    print(f"Downloading model from {MODEL_REPO}...")
    print(f"Model file: {MODEL_FILE}")

    model_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILE,
        cache_dir=MODEL_CACHE_DIR,
    )

    print(f"Model downloaded to: {model_path}")
    print("Loading model into memory...")

    llm = Llama(
        model_path=model_path,
        n_ctx=N_CTX,
        n_threads=N_THREADS,
        n_threads_batch=N_THREADS_BATCH,
        n_batch=N_BATCH,
        n_gpu_layers=N_GPU_LAYERS,
        verbose=False,
    )

    print("Model loaded successfully.")
    return llm


# =================================================
# Prompt management and context-window safety
# =================================================

def extract_recent_history(history, max_turns=5):
    """Extract recent chat history for follow-up questions."""
    if not history:
        return ""

    recent_history = history[-max_turns:]
    conversation = ""

    for item in recent_history:
        if isinstance(item, dict):
            role = item.get("role", "")
            content = item.get("content", "")

            if role == "user":
                conversation += f"User: {content}\n"
            elif role == "assistant":
                conversation += f"Assistant: {content}\n"

        elif isinstance(item, (list, tuple)) and len(item) == 2:
            user_msg, bot_msg = item
            conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"

    return conversation


def base_prompt_template(message, conversation, indiacode_context):
    """Main legal assistant prompt."""
    current_date = time.strftime("%Y-%m-%d")

    return f"""You are an Indian legal AI assistant.

Current date: {current_date}

You must answer using:
1. The user's question.
2. The recent conversation.
3. The integrated India Code context below.

Critical source rules:
- Use India Code context as the primary legal source.
- Do not invent legal provisions, case names, dates, citations, section numbers, or statutory text.
- If the India Code context is missing, weak, OCR-based, unreadable, or incomplete, clearly say so.
- If text came from OCR, mention that OCR can contain recognition errors.
- If a precise answer cannot be verified from India Code, say exactly what could not be verified.
- If the user asks for latest/current law, rely on the India Code context provided.

Answer style:
- Give a very detailed answer.
- Use clear headings and subheadings.
- Explain the legal position step by step.
- Mention relevant Act names, sections, definitions, procedures, rights, duties, exceptions, and practical implications where supported by the context.
- Include caveats where the source text is incomplete or OCR-based.
- For follow-up questions, use the recent conversation to maintain continuity.
- Do not claim to be a lawyer.
- Do not present the answer as formal legal advice.
- For specific legal matters, advise consulting a qualified lawyer.

Recent conversation:
{conversation}

Integrated India Code context:
{indiacode_context}

User question:
{message}

Detailed answer:"""


def count_tokens(model, prompt):
    """Count prompt tokens using llama-cpp tokenizer."""
    try:
        return len(model.tokenize(prompt.encode("utf-8"), add_bos=True))
    except Exception:
        # Fallback estimate: about 4 chars per token.
        return max(1, len(prompt) // 4)


def build_safe_prompt(model, message, history, source_docs):
    """
    Build the most detailed prompt that fits inside the model context window.

    This prevents:
    Requested tokens (...) exceed context window (...)
    """
    max_prompt_tokens = max(256, N_CTX - MAX_TOKENS - PROMPT_SAFETY_MARGIN)

    history_options = [5, 4, 3, 2, 1, 0]
    doc_options = [6, 5, 4, 3, 2, 1]
    chars_options = [3000, 2400, 1800, 1400, 1000, 700]

    for history_turns in history_options:
        conversation = extract_recent_history(history, max_turns=history_turns)

        for doc_count in doc_options:
            docs = source_docs[:doc_count]

            for chars_per_doc in chars_options:
                blocks = []

                for index, doc in enumerate(docs, start=1):
                    snippet = make_snippet(
                        doc.get("text", ""),
                        message,
                        max_chars=chars_per_doc,
                    )

                    blocks.append(
                        f"[India Code Source {index}]\n"
                        f"Title: {doc.get('title', 'India Code document')}\n"
                        f"Type: {doc.get('type', 'html')}\n"
                        f"URL: {doc.get('url')}\n"
                        f"Relevant excerpt:\n{snippet}\n"
                    )

                context = "\n\n".join(blocks)

                if not context:
                    context = (
                        "No directly relevant readable content was retrieved from India Code "
                        "for this query."
                    )

                prompt = base_prompt_template(
                    message=message,
                    conversation=conversation,
                    indiacode_context=context,
                )

                prompt_tokens = count_tokens(model, prompt)

                if prompt_tokens <= max_prompt_tokens:
                    return prompt, docs, prompt_tokens

    # Emergency fallback: never crash.
    fallback_context = (
        "India Code context was retrieved, but it was too long for the local model context window. "
        "Use only the available high-level information and clearly state that detailed source text "
        "could not fit into the prompt."
    )

    prompt = base_prompt_template(
        message=message,
        conversation=extract_recent_history(history, max_turns=1),
        indiacode_context=fallback_context,
    )

    return prompt, source_docs[:1], count_tokens(model, prompt)


def safe_generation_max_tokens(prompt_tokens):
    """
    Keep output generation within context window while allowing detailed answers.
    """
    available = N_CTX - prompt_tokens - PROMPT_SAFETY_MARGIN

    if available <= 0:
        return 128

    return max(128, min(MAX_TOKENS, available))


# =================================================
# Main chat function with streaming
# =================================================

def chat(message, history):
    """
    Gradio ChatInterface function.

    Integrated flow:
    user message -> India Code lookup -> PDF/OCR extraction -> safe prompt -> streamed local LLM answer.
    """
    if not message or not message.strip():
        yield "Please enter a question."
        return

    try:
        user_query = message.strip()

        yield "Searching India Code and reading relevant documents..."

        # India Code lookup.
        _raw_context, source_docs = build_indiacode_context(user_query)

        yield "Preparing a detailed answer from India Code sources..."

        # Load local model lazily.
        model = load_model()

        # Build context-safe prompt.
        prompt, used_docs, prompt_tokens = build_safe_prompt(
            model=model,
            message=user_query,
            history=history,
            source_docs=source_docs,
        )

        generation_tokens = safe_generation_max_tokens(prompt_tokens)

        print(
            f"Prompt tokens: {prompt_tokens}, "
            f"generation tokens: {generation_tokens}, "
            f"context window: {N_CTX}"
        )

        streamed_answer = ""

        stream = model(
            prompt,
            max_tokens=generation_tokens,
            temperature=TEMPERATURE,
            top_p=TOP_P,
            echo=False,
            stream=True,
            stop=["User:", "\nUser:", "\n\n\n"],
        )

        for chunk in stream:
            try:
                token = chunk["choices"][0].get("text", "")
            except Exception:
                token = ""

            if token:
                streamed_answer += token
                yield streamed_answer

        if not streamed_answer.strip():
            streamed_answer = "I could not generate a response. Please try rephrasing your question."

        streamed_answer = streamed_answer.strip()
        streamed_answer += format_sources(used_docs)

        yield streamed_answer

    except Exception as error:
        print("Error during India Code lookup, OCR, or generation:")
        traceback.print_exc()

        yield (
            "The app encountered an error while searching India Code, reading a PDF, "
            "running OCR, or generating the response.\n\n"
            f"Error details: {str(error)}"
        )


# =================================================
# Gradio UI
# =================================================

CUSTOM_CSS = """
#footer {
    text-align: center;
    margin-top: 24px;
    padding: 14px;
    font-size: 14px;
    color: #555;
}

#footer a {
    color: #1e40af;
    text-decoration: none;
    font-weight: 600;
}

#footer a:hover {
    text-decoration: underline;
}
"""

description = """
# 🏛️ Indian Legal AI Assistant

Ask questions about Indian laws, Acts, legal sections, rules, and follow-up questions.

This app uses:

- local GGUF inference through `llama-cpp-python`
- integrated India Code lookup
- HTML extraction
- text-based PDF extraction
- lightweight OCR fallback for scanned PDFs
- recent chat history for follow-up questions
- streamed answers for faster response display

The India Code lookup runs inside each answer.

---

**Primary source:** https://www.indiacode.nic.in/

**Disclaimer:** This assistant provides general legal information only.  
It is not a substitute for advice from a qualified legal professional.  
For specific legal matters, please consult a lawyer.
"""

with gr.Blocks(css=CUSTOM_CSS, title="Indian Legal AI Assistant") as demo:
    gr.ChatInterface(
        fn=chat,
        title="Indian Legal AI Assistant",
        description=description,
        textbox=gr.Textbox(
            placeholder="Ask about Indian laws, Acts, sections, rules, or follow-up questions...",
            lines=3,
            label="Your Question",
        ),
        examples=[
            "What is the current status of Section 377 under Indian law?",
            "Explain the Bharatiya Nyaya Sanhita in detail.",
            "What are the grounds for divorce under the Hindu Marriage Act?",
            "Find the latest India Code position on the Right to Information Act.",
        ],
        cache_examples=False,
    )

    gr.HTML(
        """
        <div id="footer">
            Developed by
            <a href="https://rotsl.github.io" target="_blank" rel="noopener noreferrer">
                Rohan R
            </a>
        </div>
        """
    )


# =================================================
# App entry point
# =================================================

if __name__ == "__main__":
    print("Starting Indian Legal AI Assistant...")
    print(f"Using model: {MODEL_REPO}/{MODEL_FILE}")
    print("CPU-only mode enabled.")
    print("Integrated India Code lookup enabled.")
    print("PDF OCR fallback enabled.")
    print(
        f"N_CTX={N_CTX}, "
        f"N_THREADS={N_THREADS}, "
        f"N_THREADS_BATCH={N_THREADS_BATCH}, "
        f"N_BATCH={N_BATCH}, "
        f"MAX_TOKENS={MAX_TOKENS}"
    )

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )