"""
Marziel Hybrid RAG — Combines PageIndex reasoning + vector similarity.

Score fusion of two retrieval paths:
1. PageIndex: LLM tree search for structured document reasoning
2. Vector: cosine similarity via fastembed for fast recall

Everything auto-installs with `pip install marziel`.
"""

import json
import os
import tempfile
import urllib.request
from typing import Optional

from marziel.vector_store import VectorStore
from marziel.pageindex_rag import PageIndexRAG


LLM_PORT = int(os.environ.get("LLM_PORT", "8000"))


def _llm_call(prompt: str, system: str = "", max_tokens: int = 1024) -> str:
    """Call the local Marziel LLM."""
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    data = json.dumps({
        "model": "default",
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.3,
    }).encode()

    req = urllib.request.Request(
        f"http://localhost:{LLM_PORT}/v1/chat/completions",
        data=data,
        headers={"Content-Type": "application/json"},
    )
    try:
        resp = urllib.request.urlopen(req, timeout=120)
        result = json.loads(resp.read())
        return result["choices"][0]["message"]["content"]
    except Exception as e:
        return f"[LLM Error: {e}]"


def _extract_pages_from_pdf(pdf_path: str) -> list[str]:
    """Extract text from each page of a PDF."""
    try:
        import fitz  # PyMuPDF
    except ImportError:
        import subprocess, sys
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "PyMuPDF", "-q"],
            check=True, capture_output=True, timeout=120
        )
        import fitz

    doc = fitz.open(pdf_path)
    pages = []
    for page in doc:
        pages.append(page.get_text())
    doc.close()
    return pages


class HybridRAG:
    """
    Hybrid RAG system combining vector search + PageIndex reasoning.

    Usage:
        rag = HybridRAG()
        doc_id = rag.index_document("/path/to/doc.pdf")
        result = rag.query("What is the EEXI calculation?")
    """

    def __init__(self, vector_weight: float = 0.4,
                 pageindex_weight: float = 0.6):
        self.vector_store = VectorStore()
        self.pageindex = PageIndexRAG()
        self.vector_weight = vector_weight
        self.pageindex_weight = pageindex_weight
        # Cache pages in memory for tree search
        self._page_cache: dict[int, list[str]] = {}

    def index_document(self, file_path: str,
                       use_pageindex: bool = True) -> dict:
        """
        Index a document for hybrid RAG.

        1. Extract pages from PDF
        2. Build PageIndex tree (reasoning-based)
        3. Embed chunks into vector store
        """
        filename = os.path.basename(file_path)

        # Extract pages
        if file_path.lower().endswith(".pdf"):
            pages = _extract_pages_from_pdf(file_path)
        elif file_path.lower().endswith((".txt", ".md")):
            with open(file_path) as f:
                content = f.read()
            # Split by double newlines or headers for text files
            pages = [p for p in content.split("\n\n") if p.strip()]
        else:
            raise ValueError(f"Unsupported file type: {file_path}")

        if not pages:
            return {"error": "No content extracted from document"}

        # Build PageIndex tree
        tree = None
        if use_pageindex:
            try:
                tree = self.pageindex.build_tree(pages, filename)
            except Exception as e:
                tree = None
                print(f"  PageIndex tree generation failed: {e}")

        # Generate document summary
        preview = "\n".join(pages[:3])[:2000]
        summary = _llm_call(
            f"Summarize this document in 2-3 sentences:\n\n{preview}",
            system="Be concise."
        )

        # Add to vector store (embeddings + metadata)
        doc_id = self.vector_store.add_document(
            filename=filename,
            pages=pages,
            pageindex_tree=tree,
            summary=summary,
        )

        # Cache pages for tree search
        self._page_cache[doc_id] = pages

        return {
            "document_id": doc_id,
            "filename": filename,
            "pages": len(pages),
            "has_tree": tree is not None,
            "summary": summary,
        }

    def query(self, question: str, doc_id: Optional[int] = None,
              top_k: int = 5) -> dict:
        """
        Hybrid query: combines vector search + PageIndex reasoning.

        Returns answer with citations.
        """
        all_results = []

        # Path 1: Vector similarity search
        vector_results = self.vector_store.search(
            query=question, top_k=top_k, doc_id=doc_id
        )
        for r in vector_results:
            r["score"] *= self.vector_weight
        all_results.extend(vector_results)

        # Path 2: PageIndex tree search (if tree exists)
        if doc_id:
            doc = self.vector_store.get_document(doc_id)
            if doc and doc.get("pageindex_tree"):
                pages = self._page_cache.get(doc_id)
                if not pages:
                    # Reload pages from chunks
                    chunks = self.vector_store.conn.execute(
                        """SELECT content, page_num FROM document_chunks
                           WHERE document_id = ? ORDER BY page_num""",
                        (doc_id,)
                    ).fetchall()
                    pages = [c["content"] for c in chunks]
                    self._page_cache[doc_id] = pages

                if pages:
                    tree_results = self.pageindex.tree_search(
                        doc["pageindex_tree"], question, pages
                    )
                    for r in tree_results:
                        r["score"] *= self.pageindex_weight
                    all_results.extend(tree_results)
        else:
            # Search across all documents with PageIndex trees
            docs = self.vector_store.list_documents()
            for doc_info in docs[:3]:  # Limit to recent 3 docs
                doc = self.vector_store.get_document(doc_info["id"])
                if doc and doc.get("pageindex_tree"):
                    pages = self._page_cache.get(doc_info["id"])
                    if pages:
                        tree_results = self.pageindex.tree_search(
                            doc["pageindex_tree"], question, pages
                        )
                        for r in tree_results:
                            r["score"] *= self.pageindex_weight
                            r["document_id"] = doc_info["id"]
                        all_results.extend(tree_results)

        # Deduplicate and rank
        seen = set()
        unique_results = []
        for r in sorted(all_results, key=lambda x: x["score"], reverse=True):
            key = (r.get("document_id"), r.get("page_num"))
            if key not in seen:
                seen.add(key)
                unique_results.append(r)

        top_results = unique_results[:top_k]

        # Generate answer with context
        if not top_results:
            answer = _llm_call(
                question,
                system="You are Marziel AI. Answer based on your knowledge."
            )
            return {
                "answer": answer,
                "sources": [],
                "retrieval_method": "none",
            }

        context = "\n\n---\n\n".join([
            f"[Source: {r.get('section', 'Page ' + str(r.get('page_num', '?')))}] "
            f"(score: {r['score']:.2f})\n{r['content'][:1500]}"
            for r in top_results
        ])

        answer = _llm_call(
            f"""Answer the question using the provided context. Cite sources.

Context:
{context}

Question: {question}""",
            system="You are Marziel AI. Answer using the provided context. "
                   "Cite page numbers and sections when possible."
        )

        sources = [
            {
                "page": r.get("page_num"),
                "section": r.get("section", r.get("title", "")),
                "score": round(r["score"], 3),
                "method": r.get("source", "unknown"),
            }
            for r in top_results
        ]

        return {
            "answer": answer,
            "sources": sources,
            "retrieval_method": "hybrid" if any(
                r.get("source") == "pageindex" for r in top_results
            ) else "vector",
        }

    def list_documents(self) -> list[dict]:
        """List all indexed documents."""
        return self.vector_store.list_documents()

    def delete_document(self, doc_id: int):
        """Delete a document from the RAG system."""
        self.vector_store.delete_document(doc_id)
        self._page_cache.pop(doc_id, None)
