# Copyright 2025-2026 Gregorio Elias Roecker Momm and nxCypher contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Knowledge Base Loader for NXCypher.

Loads import bundle JSONL files into a NetworkX graph for Cypher querying.
"""

import json
import gzip
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
import networkx as nx


class KnowledgeBaseLoader:
    """Loads import bundle data into a NetworkX graph for NXCypher queries."""

    # Entity type to label mapping
    ENTITY_LABELS = {
        "Program": "Program",
        "CodeBlock": "CodeBlock",
        "Symbol": "Symbol",
        "DataStructure": "DataStructure",
        "Copybook": "Copybook",
        "PseudoCode": "PseudoCode",
        "ProgramContract": "ProgramContract",
        "BlockContract": "BlockContract",
        "Cluster": "Cluster",
        "Community": "Community",
    }

    def __init__(self, bundle_path: Union[str, Path]):
        """
        Initialize the loader with a bundle path.

        Args:
            bundle_path: Path to the import_bundle directory
        """
        self.bundle_path = Path(bundle_path)
        self.graph = nx.MultiDiGraph()
        self.stats = {
            "nodes": {},
            "relationships": 0,
            "errors": []
        }

    def _open_jsonl(self, filepath: Path):
        """Open a JSONL file (handles .gz compression)."""
        if filepath.suffix == '.gz':
            return gzip.open(filepath, 'rt', encoding='utf-8')
        return open(filepath, 'r', encoding='utf-8')

    def _load_jsonl_file(self, filepath: Path) -> List[Dict]:
        """Load all records from a JSONL file."""
        records = []
        if not filepath.exists():
            return records

        try:
            with self._open_jsonl(filepath) as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError as e:
                        self.stats["errors"].append(f"{filepath}:{line_num}: {e}")
        except Exception as e:
            self.stats["errors"].append(f"Error reading {filepath}: {e}")

        return records

    def _extract_node_id(self, record: Dict) -> str:
        """Extract the node ID from a record."""
        meta = record.get("_meta", {})
        data = record.get("data", {})

        # Try global_id first (most reliable)
        if meta.get("global_id"):
            return meta["global_id"]

        # Fall back to entity-specific ID fields
        id_fields = [
            "program_id", "block_id", "symbol_id", "structure_id",
            "copybook_id", "pseudo_id", "contract_id", "cluster_id",
            "community_id", "id"
        ]
        for field in id_fields:
            if data.get(field):
                return data[field]

        # Last resort: content hash
        return meta.get("content_hash", str(hash(json.dumps(record))))

    def _build_node_properties(self, record: Dict) -> Dict[str, Any]:
        """Build node properties from a record."""
        meta = record.get("_meta", {})
        data = record.get("data", {})

        props = {}

        # Add all data fields
        for key, value in data.items():
            # Skip complex nested structures for now (can be enabled if needed)
            if isinstance(value, (str, int, float, bool)) or value is None:
                props[key] = value
            elif isinstance(value, list) and all(isinstance(x, (str, int, float)) for x in value):
                props[key] = value

        # Add metadata fields with underscore prefix
        props["_global_id"] = meta.get("global_id")
        props["_content_hash"] = meta.get("content_hash")
        props["_knowledge_space"] = meta.get("knowledge_space")
        props["_extraction_source"] = meta.get("extraction_source")
        props["_source_file"] = meta.get("source_file")
        props["_version"] = meta.get("version")

        # Set __labels__ for NXCypher
        entity_type = meta.get("entity_type", "Unknown")
        props["__labels__"] = {self.ENTITY_LABELS.get(entity_type, entity_type)}

        return props

    def load_nodes(self, node_type: str, filename: str) -> int:
        """
        Load nodes of a specific type from a JSONL file.

        Args:
            node_type: The entity type (e.g., "Program", "CodeBlock")
            filename: The filename (e.g., "programs.jsonl")

        Returns:
            Number of nodes loaded
        """
        filepath = self.bundle_path / "nodes" / filename

        # Also check for .gz version
        if not filepath.exists():
            filepath = self.bundle_path / "nodes" / (filename + ".gz")

        records = self._load_jsonl_file(filepath)
        count = 0

        for record in records:
            node_id = self._extract_node_id(record)
            props = self._build_node_properties(record)
            self.graph.add_node(node_id, **props)
            count += 1

        self.stats["nodes"][node_type] = count
        return count

    def load_relationships(self, filename: str = "all_relationships.jsonl") -> int:
        """
        Load relationships from a JSONL file.

        Args:
            filename: The filename (default: "all_relationships.jsonl")

        Returns:
            Number of relationships loaded
        """
        filepath = self.bundle_path / "relationships" / filename

        # Also check for .gz version
        if not filepath.exists():
            filepath = self.bundle_path / "relationships" / (filename + ".gz")

        records = self._load_jsonl_file(filepath)
        count = 0

        for record in records:
            data = record.get("data", {})
            meta = record.get("_meta", {})

            # Get source and target
            source_id = data.get("source_global_id") or data.get("source_id")
            target_id = data.get("target_global_id") or data.get("target_id")
            rel_type = data.get("relationship_type", "RELATED_TO")

            if not source_id or not target_id:
                continue

            # Build edge properties
            edge_props = {
                "__labels__": {rel_type},
                "_relationship_type": rel_type,
            }

            # Add any additional properties
            properties = data.get("properties", {})
            if isinstance(properties, dict):
                edge_props.update(properties)

            self.graph.add_edge(source_id, target_id, **edge_props)
            count += 1

        self.stats["relationships"] = count
        return count

    def load_all(self) -> "KnowledgeBaseLoader":
        """
        Load all node types and relationships from the bundle.

        Returns:
            self for chaining
        """
        # Define node files to load
        node_files = [
            ("Program", "programs.jsonl"),
            ("CodeBlock", "code_blocks.jsonl"),
            ("Symbol", "symbols.jsonl"),
            ("DataStructure", "data_structures.jsonl"),
            ("Copybook", "copybooks.jsonl"),
            ("PseudoCode", "pseudo_codes.jsonl"),
            ("ProgramContract", "program_contracts.jsonl"),
            ("BlockContract", "block_contracts.jsonl"),
            ("Cluster", "clusters.jsonl"),
            ("Community", "communities.jsonl"),
        ]

        print(f"Loading knowledge base from: {self.bundle_path}")

        # Load all node types
        for node_type, filename in node_files:
            count = self.load_nodes(node_type, filename)
            if count > 0:
                print(f"  Loaded {count:,} {node_type} nodes")

        # Load relationships
        rel_count = self.load_relationships()
        print(f"  Loaded {rel_count:,} relationships")

        # Summary
        total_nodes = sum(self.stats["nodes"].values())
        print(f"\nTotal: {total_nodes:,} nodes, {rel_count:,} relationships")

        if self.stats["errors"]:
            print(f"Warnings: {len(self.stats['errors'])} errors encountered")

        return self

    def get_graph(self) -> nx.MultiDiGraph:
        """Return the loaded graph."""
        return self.graph

    def get_stats(self) -> Dict:
        """Return loading statistics."""
        return self.stats


def load_knowledge_base(bundle_path: Union[str, Path]) -> nx.MultiDiGraph:
    """
    Convenience function to load a knowledge base.

    Args:
        bundle_path: Path to the import_bundle directory

    Returns:
        NetworkX MultiDiGraph with all nodes and relationships
    """
    loader = KnowledgeBaseLoader(bundle_path)
    loader.load_all()
    return loader.get_graph()


# Example usage and testing
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python kb_loader.py <path_to_import_bundle>")
        print("\nExample:")
        print("  python kb_loader.py /path/to/pipeline_output/genapp/import_bundle")
        sys.exit(1)

    bundle_path = sys.argv[1]

    # Load the knowledge base
    G = load_knowledge_base(bundle_path)

    # Test with NXCypher
    from nxcypher import NXCypher

    gc = NXCypher(G)

    print("\n" + "="*60)
    print("Sample Queries")
    print("="*60)

    # Query 1: List all programs
    print("\n1. All Programs:")
    result = gc.run("""
        MATCH (p:Program)
        RETURN p.program_name, p.language, p.lines_of_code
        LIMIT 10
    """)
    for i, name in enumerate(result.get("p.program_name", [])):
        lang = result.get("p.language", [None])[i]
        loc = result.get("p.lines_of_code", [None])[i]
        print(f"   {name}: {lang}, {loc} lines")

    # Query 2: Program calls
    print("\n2. Program Calls:")
    result = gc.run("""
        MATCH (p1:Program)-[:CALLS]->(p2:Program)
        RETURN p1.program_name AS caller, p2.program_name AS callee
        LIMIT 10
    """)
    for i, caller in enumerate(result.get("caller", [])):
        callee = result.get("callee", [None])[i]
        print(f"   {caller} -> {callee}")

    # Query 3: Code blocks in a program
    print("\n3. Code Blocks (first program):")
    result = gc.run("""
        MATCH (p:Program)
        RETURN p.program_id
        LIMIT 1
    """)
    if result.get("p.program_id"):
        prog_id = result["p.program_id"][0]
        result = gc.run(f"""
            MATCH (b:CodeBlock)
            WHERE b.unit_id = "{prog_id}"
            RETURN b.block_name, b.block_type
            LIMIT 10
        """)
        for i, name in enumerate(result.get("b.block_name", [])):
            btype = result.get("b.block_type", [None])[i]
            print(f"   {name} ({btype})")

    # Query 4: Data structures accessed
    print("\n4. Table Accesses:")
    result = gc.run("""
        MATCH (p:Program)-[:ACCESSES_TABLE]->(d:DataStructure)
        RETURN p.program_name, d.structure_name
        LIMIT 10
    """)
    for i, prog in enumerate(result.get("p.program_name", [])):
        struct = result.get("d.structure_name", [None])[i]
        print(f"   {prog} accesses {struct}")

    print("\n" + "="*60)
    print("Knowledge base loaded and ready for queries!")
    print("="*60)
