"""
Tree Builder - Document Tree Assembly

This module builds a hierarchical document tree from classified spans.
Uses a stack-based parser to:
1. Process blocks in reading order (top-to-bottom, left-to-right)
2. Assign body text to the nearest preceding header
3. Output nested DocumentNode structure

The tree structure enables:
- Hierarchical navigation by the agent
- Section-based summarization
- Context-aware retrieval

Multi-Document Support:
- Detects document boundaries in combined PDFs
- Creates separate subtrees for each logical document
- Preserves internal hierarchy within each document
"""

from __future__ import annotations

from typing import TYPE_CHECKING
from uuid import uuid4

import structlog

from rnsr.models import ClassifiedSpan, DocumentNode, DocumentTree

if TYPE_CHECKING:
    from rnsr.ingestion.document_boundary import DocumentSegment

logger = structlog.get_logger(__name__)


class TreeBuilder:
    """
    Builds a hierarchical document tree from classified spans.
    
    Uses a stack-based approach where:
    - Headers push new nodes onto the stack (at appropriate level)
    - Body text appends to the current node's content
    - The stack maintains the current path in the hierarchy
    """

    def __init__(self):
        """Initialize the Tree Builder."""
        self._current_page = -1

    def build_tree(
        self,
        spans: list[ClassifiedSpan],
        document_title: str = "",
    ) -> DocumentTree:
        """
        Build a document tree from classified spans.
        
        Args:
            spans: List of ClassifiedSpan (headers and body text).
            document_title: Optional title for the document.
            
        Returns:
            DocumentTree with hierarchical structure.
        """
        if not spans:
            logger.warning("empty_spans_list")
            root = DocumentNode(id="root", level=0, header="Document")
            return DocumentTree(
                title=document_title or "Untitled",
                root=root,
                total_nodes=1,
            )
        
        # Sort spans by reading order (page, then top-to-bottom, left-to-right)
        sorted_spans = self._sort_by_reading_order(spans)
        
        # Initialize root node
        root = DocumentNode(
            id="root",
            level=0,
            header=document_title or self._extract_title(sorted_spans),
        )
        
        # Stack tracks current path: [(level, node), ...]
        # Start with root at level 0
        stack: list[tuple[int, DocumentNode]] = [(0, root)]
        
        # Process each span
        for span in sorted_spans:
            if span.role == "header":
                self._process_header(span, stack)
            else:
                self._process_body(span, stack)
        
        # Count total nodes
        total_nodes = self._count_nodes(root)
        
        logger.info(
            "tree_built",
            total_nodes=total_nodes,
            max_depth=self._get_max_depth(root),
        )
        
        return DocumentTree(
            title=root.header,
            root=root,
            total_nodes=total_nodes,
        )

    def _sort_by_reading_order(
        self, 
        spans: list[ClassifiedSpan],
    ) -> list[ClassifiedSpan]:
        """
        Sort spans by reading order: page number, then y position, then x position.
        """
        return sorted(
            spans,
            key=lambda s: (
                s.page_num,
                s.bbox.y0,  # Top to bottom
                s.bbox.x0,  # Left to right
            ),
        )

    def _extract_title(self, spans: list[ClassifiedSpan]) -> str:
        """
        Extract document title from the first H1 header.
        """
        for span in spans:
            if span.role == "header" and span.header_level == 1:
                return span.text.strip()
        
        # Fallback: use first header of any level
        for span in spans:
            if span.role == "header":
                return span.text.strip()
        
        return "Untitled Document"

    def _process_header(
        self,
        span: ClassifiedSpan,
        stack: list[tuple[int, DocumentNode]],
    ) -> None:
        """
        Process a header span by adding a new node to the tree.
        
        The stack is adjusted so that:
        - Headers pop nodes until finding a parent with lower level
        - Then push the new header node
        """
        level = span.header_level
        
        # Pop stack until we find a parent with level < this header
        while len(stack) > 1 and stack[-1][0] >= level:
            stack.pop()
        
        # Create new node for this header
        new_node = DocumentNode(
            id=f"sec_{str(uuid4())[:6]}",
            level=level,
            header=span.text.strip(),
            page_num=span.page_num,
            bbox=span.bbox,
        )
        
        # Add as child of current top of stack
        parent_node = stack[-1][1]
        parent_node.children.append(new_node)
        
        # Push onto stack
        stack.append((level, new_node))
        
        logger.debug(
            "header_added",
            level=level,
            header=span.text[:50],
            parent=parent_node.header[:30] if parent_node.header else "root",
        )

    def _process_body(
        self,
        span: ClassifiedSpan,
        stack: list[tuple[int, DocumentNode]],
    ) -> None:
        """
        Process body text by appending to the current node's content.
        """
        if not stack:
            return
        
        current_node = stack[-1][1]
        
        # Append text with appropriate spacing
        if current_node.content:
            # Check if we need a new paragraph (different page or significant y gap)
            current_node.content += " " + span.text.strip()
        else:
            current_node.content = span.text.strip()
        
        # Update page number if not set
        if current_node.page_num is None:
            current_node.page_num = span.page_num

    def _count_nodes(self, node: DocumentNode) -> int:
        """Recursively count all nodes in the tree."""
        count = 1  # Count this node
        for child in node.children:
            count += self._count_nodes(child)
        return count

    def _get_max_depth(self, node: DocumentNode, current_depth: int = 0) -> int:
        """Get the maximum depth of the tree."""
        if not node.children:
            return current_depth
        
        return max(
            self._get_max_depth(child, current_depth + 1)
            for child in node.children
        )


def build_document_tree(
    spans: list[ClassifiedSpan],
    title: str = "",
) -> DocumentTree:
    """
    Convenience function to build a document tree from classified spans.
    
    Args:
        spans: List of ClassifiedSpan from header classification.
        title: Optional document title.
        
    Returns:
        DocumentTree with hierarchical structure.
        
    Example:
        analysis, raw_spans = analyze_font_histogram("doc.pdf")
        classified = classify_headers(raw_spans, analysis)
        tree = build_document_tree(classified)
    """
    builder = TreeBuilder()
    return builder.build_tree(spans, title)


def build_multi_document_tree(
    segments: list[DocumentSegment],
    container_title: str = "",
) -> DocumentTree:
    """
    Build a tree from multiple document segments.
    
    Creates a structure where:
    - Root node represents the container (e.g., the PDF file)
    - Each document segment becomes a level-1 child
    - Internal structure of each document is preserved below that
    
    Args:
        segments: List of DocumentSegment from boundary detection
        container_title: Title for the root container node
        
    Returns:
        DocumentTree with multi-document structure
    """
    from rnsr.ingestion.header_classifier import classify_headers
    from rnsr.ingestion.font_histogram import FontHistogramAnalyzer
    
    if not segments:
        root = DocumentNode(id="root", level=0, header=container_title or "Documents")
        return DocumentTree(
            title=container_title or "Documents",
            root=root,
            total_nodes=1,
        )
    
    # If only one segment, build a regular tree
    if len(segments) == 1:
        builder = TreeBuilder()
        # Need to classify spans first
        analyzer = FontHistogramAnalyzer()
        analysis = analyzer.analyze_spans(segments[0].spans)
        classified = classify_headers(segments[0].spans, analysis)
        return builder.build_tree(classified, segments[0].title or container_title)
    
    # Create container root node
    root = DocumentNode(
        id="root",
        level=0,
        header=container_title or f"{len(segments)} Documents",
    )
    
    total_nodes = 1
    builder = TreeBuilder()
    analyzer = FontHistogramAnalyzer()
    
    for i, segment in enumerate(segments):
        # Create a document node for this segment
        doc_title = segment.title or f"Document {i + 1}"
        
        logger.debug(
            "building_document_subtree",
            doc_index=i,
            title=doc_title[:50],
            span_count=len(segment.spans),
            pages=f"{segment.start_page}-{segment.end_page}",
        )
        
        # Analyze and classify spans for this segment
        if segment.spans:
            analysis = analyzer.analyze_spans(segment.spans)
            classified = classify_headers(segment.spans, analysis)
            
            # Build subtree for this document
            subtree = builder.build_tree(classified, doc_title)
            
            # The subtree's root becomes a child of our container root
            # But we need to shift levels down by 1
            # Pass doc_index to generate unique IDs and avoid collision with container root
            doc_node = _shift_node_levels(subtree.root, level_shift=1, doc_index=i)
            doc_node.header = doc_title  # Ensure title is preserved
            
            root.children.append(doc_node)
            total_nodes += subtree.total_nodes
        else:
            # Empty segment - just create a placeholder
            doc_node = DocumentNode(
                id=f"doc_{i:03d}",
                level=1,
                header=doc_title,
                page_num=segment.start_page,
            )
            root.children.append(doc_node)
            total_nodes += 1
    
    logger.info(
        "multi_document_tree_built",
        documents=len(segments),
        total_nodes=total_nodes,
    )
    
    return DocumentTree(
        title=container_title or f"{len(segments)} Documents",
        root=root,
        total_nodes=total_nodes,
    )


def _shift_node_levels(node: DocumentNode, level_shift: int, doc_index: int = 0) -> DocumentNode:
    """
    Recursively shift all node levels by a fixed amount.
    
    Used when embedding a subtree under a higher-level node.
    Generates new unique IDs to avoid collisions with the container root.
    """
    # Generate new unique ID to avoid collision with container root
    if node.id == "root":
        new_id = f"doc_{doc_index:03d}"
    else:
        new_id = f"d{doc_index}_{node.id}"
    
    new_node = DocumentNode(
        id=new_id,
        level=node.level + level_shift,
        header=node.header,
        content=node.content,
        page_num=node.page_num,
        bbox=node.bbox,
        children=[],
    )
    
    for child in node.children:
        new_node.children.append(_shift_node_levels(child, level_shift, doc_index))
    
    return new_node


def tree_to_dict(tree: DocumentTree) -> dict:
    """
    Convert a DocumentTree to a nested dictionary for serialization.
    
    This produces the JSON structure specified in the plan.
    """
    def node_to_dict(node: DocumentNode) -> dict:
        return {
            "id": node.id,
            "type": "section" if node.level > 0 else "document",
            "level": node.level,
            "header": node.header,
            "content": node.content,
            "page_num": node.page_num,
            "children": [node_to_dict(child) for child in node.children],
        }
    
    return {
        "id": tree.id,
        "title": tree.title,
        "total_nodes": tree.total_nodes,
        "ingestion_tier": tree.ingestion_tier,
        "ingestion_method": tree.ingestion_method,
        "root": node_to_dict(tree.root),
    }
