"""
Lógica de las 4 tools de RAG sobre ESL e ISLP (v2).

Diferencia con v1: la base ChromaDB se obtiene de un dataset publicado en
HF Hub vía `snapshot_download`. La primera invocación tarda lo que tarde
la descarga (~40 MB); las siguientes son cache hit.

Variables de entorno:
- RAG_CHROMA_DIR        Si está set y apunta a una carpeta existente, se usa
                        en lugar del dataset (útil para dev local con índice
                        recién regenerado por `ingest.py`).
- RAG_CHROMA_DATASET    Repo del dataset HF a descargar.
                        Default: gusdelact/rag-esl-islp-chromadb
- RAG_CHROMA_REVISION   Revision (branch/tag/commit) del dataset.
                        Default: main
- RAG_CHROMA_CACHE_DIR  Directorio cache para el snapshot_download.
                        Default: ~/.cache/rag-books-mcp/chroma_db (o /data/chroma_db
                        si existe /data, como en HF Spaces con persistent storage).
- HF_TOKEN              Solo si el dataset es privado.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Optional

import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction


# --- Configuración ---
EMBEDDING_MODEL = "all-MiniLM-L6-v2"

DEFAULT_DATASET = "gusdelact/rag-esl-islp-chromadb"
DEFAULT_REVISION = "main"


def _resolve_cache_dir() -> Path:
    """Decide dónde guardar el snapshot del dataset.

    Prioridad:
      1. RAG_CHROMA_CACHE_DIR si está set.
      2. /data/chroma_db si existe /data (HF Spaces con persistent storage).
      3. ~/.cache/rag-books-mcp/chroma_db.
    """
    explicit = os.environ.get("RAG_CHROMA_CACHE_DIR")
    if explicit:
        return Path(explicit)
    if Path("/data").is_dir():
        return Path("/data/chroma_db")
    return Path.home() / ".cache" / "rag-books-mcp" / "chroma_db"


# Singletons por proceso
_client: Optional[chromadb.ClientAPI] = None
_embedding_fn = None
_chroma_path_resolved: Optional[str] = None


def _resolve_chroma_path() -> str:
    """Resuelve la ruta a usar como ChromaDB persistente.

    Si `RAG_CHROMA_DIR` apunta a una carpeta existente, la usa tal cual.
    En caso contrario, baja `RAG_CHROMA_DATASET@RAG_CHROMA_REVISION` desde HF
    Hub y devuelve la ruta al snapshot.
    """
    global _chroma_path_resolved
    if _chroma_path_resolved is not None:
        return _chroma_path_resolved

    override = os.environ.get("RAG_CHROMA_DIR")
    if override and Path(override).is_dir():
        print(f"[rag-books-mcp v2] Using local RAG_CHROMA_DIR: {override}", file=sys.stderr)
        _chroma_path_resolved = override
        return override

    repo_id = os.environ.get("RAG_CHROMA_DATASET", DEFAULT_DATASET)
    revision = os.environ.get("RAG_CHROMA_REVISION", DEFAULT_REVISION)
    cache_dir = _resolve_cache_dir()
    cache_dir.mkdir(parents=True, exist_ok=True)

    print(
        f"[rag-books-mcp v2] Downloading ChromaDB from HF dataset "
        f"{repo_id}@{revision} into {cache_dir} ...",
        file=sys.stderr,
    )

    # Import perezoso para no pagar el costo si está cacheado vía RAG_CHROMA_DIR.
    from huggingface_hub import snapshot_download

    snapshot = snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        revision=revision,
        cache_dir=str(cache_dir),
        token=os.environ.get("HF_TOKEN"),  # solo si es privado
    )

    print(f"[rag-books-mcp v2] ChromaDB ready at {snapshot}", file=sys.stderr)
    _chroma_path_resolved = snapshot
    return snapshot


def get_client() -> chromadb.ClientAPI:
    """Cliente ChromaDB persistente (singleton)."""
    global _client
    if _client is None:
        _client = chromadb.PersistentClient(path=_resolve_chroma_path())
    return _client


def get_embedding_fn():
    """Función de embedding `sentence-transformers/all-MiniLM-L6-v2` (singleton)."""
    global _embedding_fn
    if _embedding_fn is None:
        _embedding_fn = SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)
    return _embedding_fn


def get_collection(name: str):
    """Obtiene una colección de ChromaDB por nombre."""
    return get_client().get_collection(name=name, embedding_function=get_embedding_fn())


# --- Tools (idénticas en comportamiento a v1) ---

def search_theory(
    query: str,
    book: str = "both",
    top_k: int = 5,
) -> str:
    """
    Busca fragmentos relevantes en los libros ESL e ISLP usando búsqueda semántica.

    Args:
        query (str): Consulta en lenguaje natural (ej: "bias-variance tradeoff",
                     "regularización L1 vs L2", "random forest out-of-bag error").
        book (str): Libro donde buscar. Opciones: "esl", "islp", "both" (default: "both").
        top_k (int): Número de resultados a devolver (default: 5, máximo: 10).

    Returns:
        str: Fragmentos relevantes con metadata (libro, capítulo, sección, similitud).
    """
    top_k = min(max(int(top_k), 1), 10)

    collections_to_search = []
    if book in ("esl", "both"):
        try:
            collections_to_search.append(("ESL", get_collection("esl_chapters")))
        except Exception:
            pass
    if book in ("islp", "both"):
        try:
            collections_to_search.append(("ISLP", get_collection("islp_chapters")))
        except Exception:
            pass

    if not collections_to_search:
        return (
            "❌ No se encontraron colecciones. Verifica que el dataset HF "
            "esté disponible o ejecuta la ingesta local."
        )

    results = []
    for book_label, collection in collections_to_search:
        res = collection.query(query_texts=[query], n_results=top_k)

        if res["documents"] and res["documents"][0]:
            for doc, meta, dist in zip(
                res["documents"][0], res["metadatas"][0], res["distances"][0]
            ):
                similarity = 1 - dist
                results.append({
                    "book": book_label,
                    "chapter": meta.get("chapter", ""),
                    "section": meta.get("section", ""),
                    "similarity": similarity,
                    "content": doc,
                })

    results.sort(key=lambda x: x["similarity"], reverse=True)
    results = results[:top_k]

    if not results:
        return f"No se encontraron resultados para: '{query}'"

    output_parts = [f"## Resultados para: \"{query}\"\n"]
    for i, r in enumerate(results, 1):
        output_parts.append(
            f"### [{i}] {r['book']} — {r['chapter']} § {r['section']}\n"
            f"**Similitud:** {r['similarity']:.3f}\n\n"
            f"{r['content'][:1500]}\n\n---\n"
        )
    return "\n".join(output_parts)


def get_section(
    book: str,
    chapter: str,
    section: str = "",
    max_chunks: int = 5,
) -> str:
    """
    Recupera una sección específica de un libro por referencia exacta.

    Args:
        book (str): Libro a consultar. Opciones: "esl" o "islp".
        chapter (str): Nombre del capítulo (búsqueda parcial soportada).
        section (str): Nombre de la sección dentro del capítulo (opcional).
        max_chunks (int): Máximo de chunks a devolver (default: 5).

    Returns:
        str: Contenido de la sección con metadata.
    """
    max_chunks = int(max_chunks)
    collection_name = f"{book}_chapters"
    try:
        collection = get_collection(collection_name)
    except Exception:
        return f"❌ Colección '{collection_name}' no encontrada. Opciones: esl, islp"

    try:
        if section:
            results = collection.get(
                where={"$and": [
                    {"chapter": {"$contains": chapter}},
                    {"section": {"$contains": section}},
                ]},
                limit=max_chunks,
            )
        else:
            results = collection.get(
                where={"chapter": {"$contains": chapter}},
                limit=max_chunks,
            )
    except Exception:
        search_query = f"{chapter} {section}".strip()
        results = collection.query(query_texts=[search_query], n_results=max_chunks)
        if results["documents"] and results["documents"][0]:
            output_parts = [f"## {book.upper()} — {chapter}\n"]
            for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
                output_parts.append(
                    f"### § {meta.get('section', 'N/A')}\n\n{doc}\n\n---\n"
                )
            return "\n".join(output_parts)
        return f"No se encontró el capítulo '{chapter}' en {book.upper()}"

    if not results["documents"]:
        search_query = f"{chapter} {section}".strip()
        results = collection.query(query_texts=[search_query], n_results=max_chunks)
        if results["documents"] and results["documents"][0]:
            output_parts = [f"## {book.upper()} — {chapter}\n"]
            for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
                output_parts.append(
                    f"### § {meta.get('section', 'N/A')}\n\n{doc}\n\n---\n"
                )
            return "\n".join(output_parts)
        return f"No se encontró el capítulo '{chapter}' en {book.upper()}"

    output_parts = [f"## {book.upper()} — {chapter}"]
    if section:
        output_parts[0] += f" § {section}"
    output_parts[0] += "\n"

    for doc, meta in zip(results["documents"], results["metadatas"]):
        sec_title = meta.get("section", "")
        chunk_idx = meta.get("chunk_index", 0)
        total = meta.get("total_chunks_in_section", 1)
        output_parts.append(
            f"### § {sec_title} (parte {chunk_idx + 1}/{total})\n\n{doc}\n\n---\n"
        )
    return "\n".join(output_parts)


def cite_foundation(
    topic: str,
    detail_level: str = "medium",
) -> str:
    """
    Devuelve la fundamentación teórica para un tema citando ambos libros (ESL + ISLP).

    Args:
        topic (str): Tema a fundamentar (ej: "ridge regression", "bagging").
        detail_level (str): "brief" (1-2), "medium" (3-4) o "deep" (6-8).

    Returns:
        str: Fundamentación teórica con citas, organizada de intuitivo (ISLP)
             a riguroso (ESL).
    """
    top_k_map = {"brief": 2, "medium": 4, "deep": 8}
    top_k = top_k_map.get(detail_level, 4)

    islp_results = []
    try:
        islp_col = get_collection("islp_chapters")
        res = islp_col.query(query_texts=[topic], n_results=top_k)
        if res["documents"] and res["documents"][0]:
            for doc, meta, dist in zip(
                res["documents"][0], res["metadatas"][0], res["distances"][0]
            ):
                islp_results.append({
                    "content": doc,
                    "chapter": meta.get("chapter", ""),
                    "section": meta.get("section", ""),
                    "similarity": 1 - dist,
                })
    except Exception:
        pass

    esl_results = []
    try:
        esl_col = get_collection("esl_chapters")
        res = esl_col.query(query_texts=[topic], n_results=top_k)
        if res["documents"] and res["documents"][0]:
            for doc, meta, dist in zip(
                res["documents"][0], res["metadatas"][0], res["distances"][0]
            ):
                esl_results.append({
                    "content": doc,
                    "chapter": meta.get("chapter", ""),
                    "section": meta.get("section", ""),
                    "similarity": 1 - dist,
                })
    except Exception:
        pass

    if not islp_results and not esl_results:
        return (
            f"❌ No se encontró fundamentación para '{topic}'. "
            "Verifica que la ingesta se haya ejecutado correctamente."
        )

    output_parts = [
        f"# Fundamentación Teórica: {topic}\n",
        f"**Nivel de detalle:** {detail_level}\n",
    ]

    if islp_results:
        output_parts.append("\n## 📘 ISLP (Explicación Intuitiva)\n")
        for i, r in enumerate(islp_results, 1):
            output_parts.append(
                f"### [{i}] Cap. {r['chapter']} § {r['section']} "
                f"(sim: {r['similarity']:.3f})\n\n"
                f"{r['content'][:1200]}\n\n---\n"
            )

    if esl_results:
        output_parts.append("\n## 📗 ESL (Tratamiento Riguroso)\n")
        for i, r in enumerate(esl_results, 1):
            output_parts.append(
                f"### [{i}] Cap. {r['chapter']} § {r['section']} "
                f"(sim: {r['similarity']:.3f})\n\n"
                f"{r['content'][:1200]}\n\n---\n"
            )

    output_parts.append("\n## 📚 Referencias\n")
    if islp_results:
        chapters = set(r["chapter"] for r in islp_results)
        output_parts.append(f"- **ISLP:** {', '.join(chapters)}\n")
    if esl_results:
        chapters = set(r["chapter"] for r in esl_results)
        output_parts.append(f"- **ESL:** {', '.join(chapters)}\n")

    return "\n".join(output_parts)


def list_available_topics() -> str:
    """
    Lista los capítulos y temas indexados en la base de conocimiento.

    Returns:
        str: Lista organizada de capítulos por libro con sus secciones principales.
    """
    output_parts = ["# 📚 Contenido Disponible en la Base de Conocimiento\n"]

    for book_key, collection_name in [("ESL", "esl_chapters"), ("ISLP", "islp_chapters")]:
        try:
            collection = get_collection(collection_name)
            all_data = collection.get(include=["metadatas"])

            if not all_data["metadatas"]:
                output_parts.append(f"\n## {book_key}: Sin datos\n")
                continue

            chapters = {}
            for meta in all_data["metadatas"]:
                chapter = meta.get("chapter", "Unknown")
                section = meta.get("section", "")
                if chapter not in chapters:
                    chapters[chapter] = set()
                if section:
                    chapters[chapter].add(section)

            output_parts.append(f"\n## 📗 {book_key}\n")
            for chapter in sorted(chapters.keys()):
                sections = sorted(chapters[chapter])
                output_parts.append(f"\n### {chapter}\n")
                if sections:
                    for sec in sections[:8]:
                        output_parts.append(f"  - {sec}\n")
                    if len(sections) > 8:
                        output_parts.append(f"  - ... y {len(sections) - 8} secciones más\n")

            total = collection.count()
            output_parts.append(f"\n**Total chunks indexados:** {total}\n")

        except Exception as e:
            output_parts.append(f"\n## {book_key}: ❌ Error ({e})\n")

    return "\n".join(output_parts)