src/app_hf.py

"""
Hugging Face Spaces compatible RAG application.
Uses Hugging Face Inference API instead of local Ollama.
"""

import hashlib
import json
import os
from pathlib import Path
from typing import Optional, List
from dotenv import load_dotenv
import requests
from bs4 import BeautifulSoup
from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader, WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_huggingface import ChatHuggingFace
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.documents import Document

load_dotenv()

try:
    import google.generativeai as genai
except ImportError:
    genai = None

PERSIST_ROOT = Path("./chroma_db")
PERSIST_ROOT.mkdir(exist_ok=True)

_RAG_CHAIN = None
_URLS = []
_VECTORSTORE = None
_EMBEDDINGS = None

# Check for HF_TOKEN
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN environment variable not set. "
        "Please add your Hugging Face API token in the Spaces settings."
    )

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GGEMINI_MODEL = os.getenv("GGEMINI_MODEL", "gemini-1.5-mini")

# Ensure Hugging Face SDK sees the token
os.environ.setdefault("HUGGINGFACEHUB_API_TOKEN", HF_TOKEN)
os.environ.setdefault("HF_TOKEN", HF_TOKEN)

# Allow overriding the model and task via environment variables.
# Default uses the requested DeepSeek model and falls back to gpt2 only if it fails.
HF_MODEL = os.getenv("HF_MODEL", "deepseek-ai/DeepSeek-V4-Pro:novita")
HF_TASK = os.getenv("HF_TASK", "conversational")

API_URL = "https://router.huggingface.co/v1/chat/completions"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}


def query(payload, timeout: int = 30):
    response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=timeout)
    try:
        return response.json()
    except ValueError:
        response.raise_for_status()


def _get_persist_directory(doc_dir: str, urls: Optional[List[str]]):
    source_key = doc_dir
    if urls:
        source_key += "|" + "|".join(sorted(urls))
    digest = hashlib.md5(source_key.encode("utf-8")).hexdigest()
    persist_dir = PERSIST_ROOT / digest
    persist_dir.mkdir(parents=True, exist_ok=True)
    return persist_dir


def _fetch_web_text(url: str) -> str:
    """Fetch the rendered text of a webpage as a fallback when WebBaseLoader fails."""
    try:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
        }
        response = requests.get(url, headers=headers, timeout=30)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")

        for element in soup(["script", "style", "noscript", "header", "footer", "nav"]):
            element.extract()

        text = " ".join(soup.stripped_strings)
        return text.strip()
    except Exception as e:
        print(f"  ✗ Fallback HTML fetch failed for {url}: {e}")
        return ""


def get_first_result_url(query: str) -> Optional[str]:
    """Return the top search result URL from DuckDuckGo HTML search."""
    url = "https://html.duckduckgo.com/html/"
    headers = {"User-Agent": "Mozilla/5.0"}
    data = {"q": query}
    try:
        response = requests.post(url, headers=headers, data=data, timeout=30)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")
        first_result = soup.find('a', class_='result__a')
        if first_result and first_result.has_attr('href'):
            return first_result['href']
    except Exception as e:
        print(f"✗ Search query failed for '{query}': {e}")
    return None


def crawl_website(url: str, max_chars: int = 5000) -> str:
    """Scrape the main text content from a URL, removing common boilerplate tags."""
    print(f"Crawling: {url} ...", flush=True)
    try:
        response = requests.get(url, timeout=30, headers={"User-Agent": "Mozilla/5.0"})
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")
        for tag in soup(["nav", "footer", "script", "style", "aside", "header", "noscript"]):
            tag.decompose()
        return soup.get_text(separator=' ', strip=True)[:max_chars]
    except Exception as e:
        return f"Error crawling {url}: {e}"


def search_and_crawl(query: str, max_chars: int = 3000, model_name: Optional[str] = None, task: Optional[str] = None) -> str:
    """Perform a search, crawl the first result, and generate an answer with the configured HF model."""
    target_url = get_first_result_url(query)
    if not target_url:
        return "❌ No search results found."

    content = crawl_website(target_url, max_chars=max_chars)
    if content.startswith("Error crawling"):
        return content

    prompt_text = (
        "Answer the user query based on the following text.\n\n"
        f"Query: {query}\n\n"
        f"Text: {content}"
    )

    llm = get_llm(model_name=model_name, task=task)
    response = llm.invoke(prompt_text)
    answer_text = extract_response_text(response)
    return f"{answer_text}\n\nSource: {target_url}"

system_prompt = (
    "You are an assistant for question-answering tasks. "
    "Use the following pieces of retrieved context to answer the question. "
    "If you don't know the answer, say that you don't know. "
    "Always cite the source URL if available.\n\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])

class GeminiLLMWrapper(RunnablePassthrough):
    """Gemini LLM wrapper for use with the LangChain runtime."""

    def __init__(self, model_name: str = GGEMINI_MODEL, temperature: float = 0.1, max_tokens: int = 512):
        if genai is None:
            raise ImportError(
                "google-generativeai is not installed. Please add it to requirements.txt and install it."
            )
        if not GEMINI_API_KEY:
            raise ValueError("GEMINI_API_KEY is not set.")
        self.model_name = model_name
        self.temperature = temperature
        self.max_tokens = max_tokens
        genai.configure(api_key=GEMINI_API_KEY)

    def invoke(self, prompt_text: str):
        try:
            if hasattr(genai, "generate"):
                response = genai.generate(
                    model=self.model_name,
                    prompt=prompt_text,
                    temperature=self.temperature,
                    max_output_tokens=self.max_tokens,
                )
            elif hasattr(genai, "responses") and hasattr(genai.responses, "generate"):
                response = genai.responses.generate(
                    model=self.model_name,
                    prompt=prompt_text,
                    temperature=self.temperature,
                    max_output_tokens=self.max_tokens,
                )
            else:
                raise RuntimeError("Unsupported google-generativeai package version")
            return extract_response_text(response)
        except Exception as exc:
            raise RuntimeError(f"Gemini generation failed: {exc}") from exc


# Helper to initialize the Hugging Face LLM with a safe fallback.
def get_llm(model_name: str | None = None, task: str | None = None):
    model = model_name or HF_MODEL
    task = task or HF_TASK
    if GEMINI_API_KEY is not None:
        try:
            print(f"Initializing Gemini LLM model={GGEMINI_MODEL}", flush=True)
            return GeminiLLMWrapper(model_name=GGEMINI_MODEL)
        except Exception as gemini_exc:
            print(f"Warning: Gemini LLM init failed: {gemini_exc}", flush=True)

    try:
        print(f"Initializing HF endpoint model={model} task={task}", flush=True)
        endpoint = HuggingFaceEndpoint(
            repo_id=model,
            task=task,
            max_new_tokens=512,
            temperature=0.1,
            top_p=0.95,
            huggingfacehub_api_token=HF_TOKEN,
        )
        llm = ChatHuggingFace(
            llm=endpoint,
            model_id=model,
            temperature=0.1,
            top_p=0.95,
            max_tokens=512,
        )
        return llm
    except Exception as e:
        print(f"Warning: failed to initialize model {model}: {e}", flush=True)
        # Try a small, widely-supported fallback model to ensure functionality
        if model != "gpt2":
            try:
                print("Falling back to gpt2", flush=True)
                endpoint = HuggingFaceEndpoint(
                    repo_id="gpt2",
                    task="text-generation",
                    max_new_tokens=512,
                    temperature=0.1,
                    top_p=0.95,
                    huggingfacehub_api_token=HF_TOKEN,
                )
                llm = ChatHuggingFace(
                    llm=endpoint,
                    model_id="gpt2",
                    temperature=0.1,
                    top_p=0.95,
                    max_tokens=512,
                )
                return llm
            except Exception as e2:
                print(f"Fallback to gpt2 failed: {e2}", flush=True)
        raise


def format_docs(docs):
    """Format retrieved documents with source citations."""
    formatted = []
    for doc in docs:
        source = doc.metadata.get('source', 'Unknown source')
        formatted.append(f"Source: {source}\n{doc.page_content}")
    return "\n\n---\n\n".join(formatted)


def _emit_progress(progress, current: float, total: float, message: str):
    """Best-effort progress reporting for terminal and UI."""
    if progress is None:
        return

    try:
        progress(current / total, desc=message)
    except TypeError:
        try:
            progress(current, total, message)
        except Exception:
            pass
    except Exception:
        pass


class GeminiEmbeddings:
    """Embed text using Google Gemini / Generative Language embeddings via API key."""

    def __init__(self, api_key: str, model_name: str = "text-embedding-004"):
        if not api_key:
            raise ValueError("GEMINI_API_KEY must be set to use Gemini embeddings.")
        self.api_key = api_key
        self.model_name = model_name
        self.model_resource = self.model_name if self.model_name.startswith("models/") else f"models/{self.model_name}"
        self.endpoint = os.getenv(
            "GGEMINI_EMBEDDING_ENDPOINT",
            f"https://generativelanguage.googleapis.com/v1beta/{self.model_resource}:embedContent?key={self.api_key}"
        )

    def _embed_text(self, text: str) -> List[float]:
        payload = {
            "model": self.model_resource,
            "content": {"parts": [{"text": text}]},
            "task_type": "retrieval_document",
        }
        headers = {"Content-Type": "application/json"}
        response = requests.post(self.endpoint, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        data = response.json()
        embedding = data.get("embedding")
        if embedding is None:
            raise ValueError(f"Gemini response missing embedding data: {data}")
        return embedding

    def _call_api(self, texts: List[str]) -> List[List[float]]:
        return [self._embed_text(text) for text in texts]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self._call_api(texts)

    def embed_query(self, text: str) -> List[float]:
        return self._call_api([text])[0]


def _get_embeddings():
    """Reuse the embedding model so persisted-store loads are faster."""
    global _EMBEDDINGS
    if _EMBEDDINGS is None:
        gemini_api_key = os.getenv("GEMINI_API_KEY")
        if gemini_api_key:
            try:
                print("Initializing Gemini embeddings...", flush=True)
                _EMBEDDINGS = GeminiEmbeddings(api_key=gemini_api_key)
            except Exception as e:
                print(f"Warning: Gemini embeddings initialization failed: {e}", flush=True)
                _EMBEDDINGS = None

        if _EMBEDDINGS is None:
            print("Initializing Hugging Face embedding model...", flush=True)
            _EMBEDDINGS = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    return _EMBEDDINGS


def load_documents_from_sources(doc_dir: str = "./my_docs", urls: Optional[List[str]] = None):
    """
    Load documents from multiple sources: PDFs and URLs.
    """
    docs = []
    
    # Load PDFs from directory
    if os.path.exists(doc_dir):
        try:
            loader = DirectoryLoader(doc_dir, glob="./*.pdf", loader_cls=PyPDFLoader)
            pdf_docs = loader.load()
            docs.extend(pdf_docs)
            if pdf_docs:
                print(f"✓ Loaded {len(pdf_docs)} PDF(s) from {doc_dir}")
        except Exception as e:
            print(f"✗ Error loading PDFs: {e}")
    
    # Load from URLs if provided
    if urls:
        print(f"Loading content from {len(urls)} URL(s)...")
        for url in urls:
            web_docs = []
            try:
                web_loader = WebBaseLoader(url)
                web_docs = web_loader.load()
                if web_docs:
                    docs.extend(web_docs)
                    print(f"  ✓ Loaded: {url} (via WebBaseLoader)")
                    continue
                print(f"  ⚠️ WebBaseLoader loaded no content for: {url}")
            except Exception as e:
                print(f"  ✗ WebBaseLoader failed for {url}: {str(e)}")

            # Fallback to direct HTML scraping if WebBaseLoader fails or returns no docs.
            fallback_text = _fetch_web_text(url)
            if fallback_text:
                docs.append(Document(page_content=fallback_text, metadata={"source": url}))
                print(f"  ✓ Loaded: {url} (via HTML fallback)")
            else:
                print(f"  ✗ No text extracted from {url}")

    return docs


def load_documents_from_crawler_cache(crawler_json: str = "./crawler_docs.json"):
    """Load pre-crawled documents from the persisted crawler cache."""
    docs = []
    if os.path.exists(crawler_json):
        try:
            with open(crawler_json, "r", encoding="utf-8") as f:
                cached_docs = json.load(f)
            for item in cached_docs:
                content = item.get("content")
                source = item.get("url", "Unknown source")
                if content:
                    docs.append(Document(page_content=content, metadata={"source": source}))
            if docs:
                print(f"✓ Loaded {len(docs)} documents from crawler cache: {crawler_json}")
        except Exception as e:
            print(f"✗ Failed to load crawler cache: {e}")
    return docs


def build_rag_chain(doc_dir: str = "./my_docs", urls: Optional[List[str]] = None, progress=None):
    """
    Build the RAG chain with documents from multiple sources.
    Optimized for Hugging Face Spaces.
    
    Args:
        doc_dir: Directory containing PDFs
        urls: List of URLs to scrape
        progress: Callable(current, total, message) for progress updates, or Gradio Progress object
    """
    global _VECTORSTORE
    
    docs = load_documents_from_sources(doc_dir, urls)
    docs.extend(load_documents_from_crawler_cache())

    if not docs:
        raise ValueError(
            f"No documents found. Add PDFs to '{doc_dir}', provide URLs, or run the crawler first."
        )

    # Code-aware text splitting
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    splits = text_splitter.split_documents(docs)
    
    if not splits:
        raise ValueError("No text content could be extracted from documents.")

    print(f"Processing {len(splits)} text chunks...", flush=True)
    
    # Local vectorization (works on HF Spaces)
    hf_embeddings = _get_embeddings()

    persist_directory = _get_persist_directory(doc_dir, urls)
    if any(persist_directory.iterdir()):
        _emit_progress(progress, 0.05, 1.0, "Found cached index. Loading persisted vector store...")
        _VECTORSTORE = Chroma(persist_directory=str(persist_directory), embedding_function=hf_embeddings)
        _emit_progress(progress, 1.0, 1.0, f"Cached vector store loaded from {persist_directory}")
    else:
        _emit_progress(progress, 0.01, 1.0, "No cached index found. Building vector store now...")
        # Create empty vectorstore for progressive adding
        _VECTORSTORE = Chroma(
            persist_directory=str(persist_directory),
            embedding_function=hf_embeddings
        )
        
        # Add documents in batches with progress tracking
        batch_size = 10
        total_chunks = len(splits)
        
        for batch_idx in range(0, total_chunks, batch_size):
            batch = splits[batch_idx:batch_idx + batch_size]
            
            # Add batch to vectorstore
            _VECTORSTORE.add_documents(batch)
            
            # Update progress
            processed = min(batch_idx + batch_size, total_chunks)
            msg = f"Indexing documents {processed}/{total_chunks}"
            _emit_progress(progress, processed, total_chunks, msg)
        
        _emit_progress(progress, 1.0, 1.0, f"Indexing complete. Saved vector store to {persist_directory}")

    # Semantic search
    retriever = _VECTORSTORE.as_retriever(search_kwargs={"k": 3})

    # Initialize or obtain the LLM (may fallback to a supported model)
    llm_instance = get_llm()

    return (
        {"context": retriever | format_docs, "input": RunnablePassthrough()}
        | prompt
        | llm_instance
    )


def get_rag_chain(doc_dir: str = "./my_docs", urls: Optional[List[str]] = None, progress=None):
    """Get or create the RAG chain."""
    global _RAG_CHAIN, _URLS

    # Preserve previously added URLs when urls=None is passed explicitly.
    urls = urls if urls is not None else (_URLS or None)
    normalized_urls = urls or []

    # Rebuild if URLs changed or if the chain is empty.
    if _RAG_CHAIN is None or _URLS != normalized_urls:
        _RAG_CHAIN = build_rag_chain(doc_dir, urls, progress=progress)
        _URLS = normalized_urls

    return _RAG_CHAIN


# Low-level router query helper for models/providers not supported by the local wrapper.

def hf_router_query(messages, model: str, timeout: int = 30):
    """Query the Hugging Face router endpoint directly.

    messages: list of chat message dicts matching the HF router schema.
    model: model id string (e.g., 'Qwen/Qwen3.6-27B:featherless-ai')
    Returns the assistant text or raises on error.
    """
    if not HF_TOKEN:
        raise RuntimeError("HF_TOKEN is not set; cannot call HF router.")

    payload = {"model": model, "messages": messages}
    data = query(payload, timeout=timeout)

    if isinstance(data, dict) and data.get("error"):
        raise RuntimeError(f"HF router error: {data}")

    choices = data.get("choices")
    if choices and isinstance(choices, list) and len(choices) > 0:
        msg = choices[0].get("message")

        def extract_text_from_item(item):
            if isinstance(item, dict):
                if item.get("type") == "text":
                    return item.get("text", "")
                return item.get("text", "") or item.get("content", "")
            if hasattr(item, "type") and getattr(item, "type") == "text":
                return getattr(item, "text", "")
            if hasattr(item, "text"):
                return getattr(item, "text", "")
            if hasattr(item, "content") and isinstance(getattr(item, "content"), str):
                return getattr(item, "content")
            return str(item)

        if isinstance(msg, dict):
            if "content" in msg:
                content = msg["content"]
                if isinstance(content, list):
                    parts = [extract_text_from_item(c) for c in content]
                    return " ".join([p for p in parts if p]).strip()
                if isinstance(content, str):
                    return content.strip()
            if "text" in msg:
                return msg["text"].strip()
        else:
            if hasattr(msg, "content"):
                content = getattr(msg, "content")
                if isinstance(content, list):
                    parts = [extract_text_from_item(c) for c in content]
                    return " ".join([p for p in parts if p]).strip()
                if isinstance(content, str):
                    return content.strip()
            if hasattr(msg, "text"):
                return getattr(msg, "text").strip()
        return str(msg)

    return str(data)


def extract_response_text(response):
    """Normalize LangChain or provider responses into plain text."""
    if isinstance(response, str):
        return response

    if hasattr(response, "content"):
        content = getattr(response, "content")
        if isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        parts.append(item.get("text", ""))
                    else:
                        parts.append(item.get("text", "") or item.get("content", ""))
                elif hasattr(item, "text"):
                    parts.append(getattr(item, "text", ""))
                elif hasattr(item, "content"):
                    parts.append(str(getattr(item, "content")))
                else:
                    parts.append(str(item))
            return " ".join(part for part in parts if part).strip()
        if isinstance(content, str):
            return content.strip()
        return str(content)

    if isinstance(response, dict):
        if "text" in response and isinstance(response["text"], str):
            return response["text"].strip()
        if "content" in response:
            return extract_response_text(response["content"])

    return str(response)


def answer_question(question: str, doc_dir: str = "./my_docs", urls: Optional[List[str]] = None, progress=None) -> str:
    """
    Answer a question using the RAG system.
    
    Args:
        question: The question to answer
        doc_dir: Directory containing PDFs
        urls: List of URLs to scrape
        progress: Callable or Gradio Progress object for tracking vector store creation
    """
    if urls is None and _URLS:
        urls = _URLS

    print(f"answer_question: doc_dir={doc_dir} urls={urls} _URLS={_URLS}", flush=True)

    try:
        rag_chain = get_rag_chain(doc_dir, urls, progress=progress)
        response = rag_chain.invoke(question)
        return extract_response_text(response)
    
    except Exception as e:
        error_msg = str(e)

        # If the error looks like a model/provider issue, try a direct HF router call as a fallback.
        low = error_msg.lower()
        if any(k in low for k in ("not supported", "model", "inference", "router", "forbidden", "403")):
            try:
                # Build a simple chat-style message payload using the question
                messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": question}
                        ]
                    }
                ]
                model_to_call = os.getenv("HF_MODEL", HF_MODEL)
                resp = hf_router_query(messages, model_to_call)
                if resp:
                    return resp
            except Exception:
                # fall through to standard error handling
                pass

        # Check for common errors and provide helpful messages
        if "HF_TOKEN" in error_msg or "unauthorized" in error_msg:
            return "❌ Error: Hugging Face API token not configured.\n\nPlease add your HF_TOKEN in the Spaces settings."

        elif "no documents found" in error_msg or "no text content" in error_msg:
            # No local documents available; fall back to a live web search + crawl.
            fallback = search_and_crawl(question)
            if fallback.startswith("❌ Error") or fallback.startswith("Error"):
                return f"❌ Error: No documents available and live web search failed.\n\n{fallback}"
            return fallback

        else:
            return f"❌ Error: {error_msg}"


def add_urls(url_list: List[str]):
    """Add URLs to scrape. Clears existing chain to rebuild."""
    global _RAG_CHAIN, _URLS
    _URLS = url_list or []
    _RAG_CHAIN = None  # Reset to rebuild
    print(f"add_urls: stored {_URLS} (count={len(_URLS)})", flush=True)


def get_url_state() -> dict:
    """Return current internal URL state for debugging."""
    return {"urls": list(_URLS), "rag_chain_set": _RAG_CHAIN is not None}


def force_rebuild(doc_dir: str = "./my_docs", urls: Optional[List[str]] = None, progress=None):
    """Force rebuilding the RAG chain and return the document count.

    Use this to verify that documents are being loaded and indexed.
    """
    global _RAG_CHAIN
    _RAG_CHAIN = None
    chain = get_rag_chain(doc_dir, urls, progress=progress)
    # Try to trigger a load to validate documents exist
    try:
        # Accessing the retriever should force any lazy initialization
        if isinstance(chain, dict) and "context" in chain:
            return {"status": "ok"}
    except Exception as e:
        return {"status": "error", "error": str(e)}
    return {"status": "ok"}


if __name__ == "__main__":
    print("Hugging Face RAG System initialized")