document_parser.py

import logging
import re
from typing import List, Dict, Any
from pathlib import Path

from llama_index.core import SimpleDirectoryReader, Document
from llama_index.core.node_parser import MarkdownNodeParser, CodeSplitter, SentenceSplitter
from llama_index.core.schema import BaseNode

from config import DATA_DIR, DEFAULT_METADATA

logger = logging.getLogger(__name__)

def get_metadata_hook(filename: str) -> Dict[str, Any]:
    """
    Hook to attach base key-value metadata to documents before embedding.
    This serves as the default fallback before chunk-level intelligence is applied.
    """
    metadata = DEFAULT_METADATA.copy()
    
    try:
        file_path = Path(filename)
        metadata["source_file"] = file_path.name
        
        # Determine base doc_type based on extension
        ext = file_path.suffix.lower()
        if ext == ".json":
            metadata["doc_type"] = "api_reference"
        elif ext == ".rst":
            metadata["doc_type"] = "documentation"
        elif ext == ".py":
            metadata["doc_type"] = "source_code"
            
    except Exception as e:
        logger.warning(f"Failed to extract base metadata for {filename}: {e}")
        
    return metadata

def enrich_chunk_metadata(node: BaseNode) -> BaseNode:
    """
    Deep Content Analysis (The Intelligence Layer): 
    Inspects the text and headers of a specific chunk to dynamically assign 
    highly accurate language, framework, and OS tags.
    """
    text = node.get_content().lower()
    
    # Inherit context from Markdown Headers (LlamaIndex stores these as Header_1, Header_2, etc.)
    headers = " ".join([str(val).lower() for key, val in node.metadata.items() if key.startswith("Header_")])
    combined_context = text + " " + headers

    # 1. OS / Platform Detection
    if any(kw in combined_context for kw in ["windows", "chocolatey", "winget", "powershell"]):
        node.metadata["os"] = "windows"
    elif any(kw in combined_context for kw in ["macos", "homebrew", "brew", "darwin", "macports"]):
        node.metadata["os"] = "macos"
    elif any(kw in combined_context for kw in ["ubuntu", "debian", "apt", "linux", "apt-get"]):
        node.metadata["os"] = "ubuntu"

    # 2. Language Detection
    if "```typescript" in text or "```ts" in text or "npm install" in text or "yarn add" in text:
        node.metadata["language"] = "typescript"
    elif "```python" in text or "```py" in text or "pip install" in text or "requirements.txt" in text:
        node.metadata["language"] = "python"
    elif "```bash" in text or "```sh" in text or "sudo " in text or "choco install" in text:
        node.metadata["language"] = "bash"

    # 3. Framework Detection
    if "react" in combined_context or "jsx" in combined_context or "tsx" in combined_context:
        node.metadata["framework"] = "react"
    elif "fastapi" in combined_context or "uvicorn" in combined_context:
        node.metadata["framework"] = "fastapi"
    elif "docker" in combined_context or "docker-compose" in combined_context or "dockerfile" in combined_context:
        node.metadata["framework"] = "docker"

    # 4. Universal Topic Tagging (Fixes the Installation Filter Bug)
    # If a chunk is about installation/setup, we relax the strict framework 
    # requirements so it is universally retrievable by Agent B.
    if any(kw in combined_context for kw in ["setup", "installation", "prerequisite", "getting started", "clone"]):
        node.metadata["doc_type"] = "installation_guide"
        
        # If the chunk didn't explicitly trigger a specific framework above, 
        # set it to 'none' so Agent B drops the strict framework filter.
        if "framework" not in node.metadata or node.metadata.get("framework") == DEFAULT_METADATA.get("framework"):
            node.metadata["framework"] = "none"
            
    return node

def load_and_chunk_documents() -> List[BaseNode]:
    """
    Ingests files, applies context-aware chunking, and passes them through 
    the intelligence layer for dynamic metadata tagging.
    """
    if not DATA_DIR.exists():
        logger.error(f"Data directory {DATA_DIR} does not exist.")
        return []

    logger.info(f"Scanning {DATA_DIR} for documents...")

    try:
        # Load raw documents
        reader = SimpleDirectoryReader(
            input_dir=str(DATA_DIR),
            required_exts=[".md", ".rst", ".json", ".py"], 
            file_metadata=get_metadata_hook
        )
        documents = reader.load_data()
        logger.info(f"Successfully loaded {len(documents)} raw document entities.")
    except Exception as e:
        logger.critical(f"Failed to load documents: {e}")
        return []

    if not documents:
        return []

    # Initialize Chunkers
    md_parser = MarkdownNodeParser()
    code_parser = CodeSplitter(language="python", chunk_lines=40, chunk_lines_overlap=15)
    fallback_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)

    raw_nodes: List[BaseNode] = []
    
    for doc in documents:
        file_name = doc.metadata.get("source_file", "")
        
        try:
            # Context-Aware Splitting
            if file_name.endswith(".md"):
                doc_nodes = md_parser.get_nodes_from_documents([doc])
            elif file_name.endswith(".py"):
                doc_nodes = code_parser.get_nodes_from_documents([doc])
            else:
                doc_nodes = fallback_parser.get_nodes_from_documents([doc])
                
            raw_nodes.extend(doc_nodes)
        except Exception as e:
            logger.error(f"Failed to parse document {file_name}: {e}")

    # Pass all nodes through the Deep Content Analysis layer
    enriched_nodes = []
    for node in raw_nodes:
        enriched_nodes.append(enrich_chunk_metadata(node))

    logger.info(f"Intelligence Layer applied. Generated {len(enriched_nodes)} context-aware vector chunks.")
    
    return enriched_nodes