"""
Marziel PageIndex RAG — Reasoning-Based Document Retrieval.

Inspired by VectifyAI/PageIndex: builds a hierarchical tree index
from documents and uses LLM reasoning to navigate the tree for
context-aware retrieval. No vector similarity — pure reasoning.

Integrated into `pip install marziel && marziel serve`.
"""

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


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 via OpenAI-compatible API."""
    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.1,
    }).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}]"


class PageIndexRAG:
    """
    Reasoning-based RAG using hierarchical document tree indexes.

    1. index_document() — builds a tree structure from document pages
    2. tree_search() — LLM navigates the tree to find relevant sections
    3. retrieve() — extracts content from relevant nodes
    """

    def __init__(self, max_pages_per_node: int = 10):
        self.max_pages_per_node = max_pages_per_node

    def build_tree(self, pages: list[str], filename: str = "document") -> dict:
        """
        Build a PageIndex tree structure from document pages.
        Uses LLM to generate a hierarchical table-of-contents.
        """
        # Build page summaries first (batch pages for efficiency)
        page_summaries = []
        batch_size = 5
        for i in range(0, len(pages), batch_size):
            batch = pages[i:i + batch_size]
            batch_text = "\n\n".join([
                f"--- PAGE {i + j + 1} ---\n{p[:1500]}"
                for j, p in enumerate(batch) if p.strip()
            ])

            prompt = f"""Analyze these document pages and provide a brief summary (1-2 sentences) for each page.
Format: one line per page as "Page X: summary"

{batch_text}"""

            result = _llm_call(prompt, system="You are a document analyst. Be concise.")
            for line in result.split("\n"):
                line = line.strip()
                if line and re.match(r"Page \d+", line):
                    page_summaries.append(line)

        # Generate tree structure from summaries
        summary_text = "\n".join(page_summaries[:100])  # Cap for context

        tree_prompt = f"""Given these page summaries from "{filename}" ({len(pages)} pages), create a hierarchical table of contents as JSON.

Page summaries:
{summary_text}

Create a tree structure JSON with this format:
{{
  "title": "Document Title",
  "summary": "Brief overall summary",
  "nodes": [
    {{
      "title": "Section Name",
      "node_id": "0001",
      "start_page": 1,
      "end_page": 5,
      "summary": "What this section covers",
      "nodes": [...]
    }}
  ]
}}

Rules:
- Each node spans a range of pages
- Leaf nodes should cover at most {self.max_pages_per_node} pages
- Use meaningful section titles
- Include summaries for each node
- Return ONLY valid JSON, no markdown."""

        result = _llm_call(tree_prompt, max_tokens=2048)

        # Parse JSON from response
        try:
            # Try to extract JSON from response
            json_match = re.search(r'\{.*\}', result, re.DOTALL)
            if json_match:
                tree = json.loads(json_match.group())
            else:
                tree = json.loads(result)
        except json.JSONDecodeError:
            # Fallback: simple tree from pages
            tree = self._fallback_tree(pages, filename)

        return tree

    def _fallback_tree(self, pages: list[str], filename: str) -> dict:
        """Create a simple tree when LLM fails to generate one."""
        nodes = []
        chunk_size = self.max_pages_per_node
        for i in range(0, len(pages), chunk_size):
            end = min(i + chunk_size, len(pages))
            nodes.append({
                "title": f"Section {i // chunk_size + 1}",
                "node_id": f"{i // chunk_size + 1:04d}",
                "start_page": i + 1,
                "end_page": end,
                "summary": f"Pages {i + 1}-{end}",
                "nodes": [],
            })
        return {
            "title": filename,
            "summary": f"Document with {len(pages)} pages",
            "nodes": nodes,
        }

    def tree_search(self, tree: dict, query: str, pages: list[str],
                    max_depth: int = 3) -> list[dict]:
        """
        Navigate the tree using LLM reasoning to find relevant sections.
        Returns list of relevant nodes with page ranges.
        """
        relevant = []
        self._search_node(tree, query, pages, relevant, depth=0,
                          max_depth=max_depth)
        return relevant

    def _search_node(self, node: dict, query: str, pages: list[str],
                     relevant: list, depth: int, max_depth: int):
        """Recursively search tree nodes using LLM reasoning."""
        children = node.get("nodes", [])

        if not children or depth >= max_depth:
            # Leaf node — check relevance
            start = node.get("start_page", 1) - 1
            end = node.get("end_page", start + 1)
            page_content = "\n".join(pages[start:end])[:3000]

            if page_content.strip():
                relevance_prompt = f"""Is this section relevant to the question?
Question: {query}

Section: {node.get('title', 'Unknown')}
Summary: {node.get('summary', 'No summary')}
Content preview: {page_content[:500]}

Answer with a relevance score 0-10 and brief reason.
Format: SCORE: X | REASON: ..."""

                result = _llm_call(relevance_prompt,
                                   system="Rate relevance 0-10. Be strict.")
                score = 0
                try:
                    score_match = re.search(r'SCORE:\s*(\d+)', result)
                    if score_match:
                        score = int(score_match.group(1))
                except Exception:
                    pass

                if score >= 5:
                    relevant.append({
                        "title": node.get("title", ""),
                        "start_page": node.get("start_page", 1),
                        "end_page": node.get("end_page", 1),
                        "score": score / 10.0,
                        "content": page_content,
                        "source": "pageindex",
                    })
            return

        # Internal node — ask LLM which children to explore
        children_desc = "\n".join([
            f"  [{i}] {c.get('title', 'Unknown')} "
            f"(pages {c.get('start_page', '?')}-{c.get('end_page', '?')}): "
            f"{c.get('summary', 'No summary')}"
            for i, c in enumerate(children)
        ])

        nav_prompt = f"""You are navigating a document tree to find sections relevant to:
Question: {query}

Current section: {node.get('title', 'Document')}
Children:
{children_desc}

Which children are potentially relevant? List their indices (comma-separated).
If none are relevant, say NONE.
Format: INDICES: 0, 2, 3"""

        result = _llm_call(nav_prompt,
                           system="Select relevant sections. Be selective.")

        if "NONE" in result.upper():
            return

        # Parse indices
        indices = set()
        idx_match = re.search(r'INDICES:\s*([\d,\s]+)', result)
        if idx_match:
            for idx_str in idx_match.group(1).split(","):
                try:
                    idx = int(idx_str.strip())
                    if 0 <= idx < len(children):
                        indices.add(idx)
                except ValueError:
                    pass

        # If no indices parsed, try all
        if not indices:
            for m in re.finditer(r'\b(\d+)\b', result):
                idx = int(m.group(1))
                if 0 <= idx < len(children):
                    indices.add(idx)

        # Recurse into selected children
        for idx in sorted(indices):
            self._search_node(children[idx], query, pages, relevant,
                              depth + 1, max_depth)
