import os
import pickle
import networkx as nx
from typing import Optional

from enhanced_toon import EnhancedToon

class LocalGraph:
    def __init__(self, root_path: str):
        self.root_path = root_path
        self.graph_dir = os.path.join(root_path, ".mcp_cache", "graph")
        os.makedirs(self.graph_dir, exist_ok=True)
        self.graph_path = os.path.join(self.graph_dir, "graph.pkl")
        self.graph = nx.DiGraph()

    # --- Creation ---
    def add_node(self, node_id: str | dict, **attrs):
        """
        Enhanced add_node that can return the generated node_id
        """
        if isinstance(node_id, dict):
            lang = node_id.get("lang", "unknown")
            kind = node_id.get("type", "Symbol")
            name = node_id.get("name", "unknown")
            fpath = node_id.get("file", "")
            node_id = f"{lang}::{kind}::{fpath}::{name}"

        # Merge any additional attributes
        final_attrs = {}
        if isinstance(node_id, dict):
            final_attrs.update(node_id)
        final_attrs.update(attrs)

        self.graph.add_node(node_id, **final_attrs)
        return node_id  # Return the ID for reference

    def _add_import_edges(self, src: str, imports: list[str]):
        for imp in imports:
            # Resolve import to a file node if present
            target = f"file::{imp}"
            self.add_edge(src, target, "imports")

    def _add_call_edges(self, src: str, calls: list[str], file_path: str):
        """Enhanced call edges that work cross-file by searching the entire graph."""
        for cal in calls:
            if not cal or not isinstance(cal, str):
                continue

            # Try to find the called function in the graph
            target_found = False

            # Search for functions with this name across ALL files
            for node_id, data in self.graph.nodes(data=True):
                if (data.get('type') in ['function', 'method', 'async_function', 'constructor'] and
                    data.get('name') == cal):
                    # Found a matching function - create call edge
                    self.add_edge(src, node_id, "calls")
                    target_found = True
                    # Don't break - a function might be called from multiple places

            # If not found, create a symbolic node for the call target
            if not target_found:
                # Create a canonical function node (file-less) for cross-file references
                symbolic_id = f"rust::function::{cal}"  # Assuming Rust for now
                if symbolic_id not in self.graph.nodes:
                    self.add_node(
                        symbolic_id,
                        type="function",
                        name=cal,
                        file="",  # No specific file
                        lang="rust",
                        is_symbolic=True
                    )
                self.add_edge(src, symbolic_id, "calls")



    # --- Persistence ---
    def save(self):
        with open(self.graph_path, "wb") as f:
            pickle.dump(self.graph, f, protocol=pickle.HIGHEST_PROTOCOL)

    def load(self) -> bool:
        if not os.path.exists(self.graph_path):
            return False
        with open(self.graph_path, "rb") as f:
            self.graph = pickle.load(f)
        return True

    def clear(self):
        self.graph = nx.DiGraph()
        if os.path.exists(self.graph_path):
            os.remove(self.graph_path)

    # --- Query helpers ---
    def find_nodes(self, **filters):
        for node, data in self.graph.nodes(data=True):
            if all(data.get(k) == v for k, v in filters.items()):
                yield node, data

    def neighbors(self, node_id, edge_type: Optional[str] = None):
        for n in self.graph.successors(node_id):
            if not edge_type or self.graph.edges[node_id, n]["type"] == edge_type:
                yield n, self.graph.nodes[n]

    def to_json(self):
        out = {
            "nodes": [
                {"id": n, **data} for n, data in self.graph.nodes(data=True)
            ],
            "edges": [
                {"src": s, "dst": d, **data} for s, d, data in self.graph.edges(data=True)
            ],
        }
        path = os.path.join(self.graph_dir, "graph.json")
        with open(path, "w") as f:
            import json
            json.dump(out, f, indent=2)
        return path

    def to_toon(self):
        """
        Export graph to compact TOON (Token-Oriented Object Notation) format.
        This reduces token count by 30-60% for LLM consumption.
        """
        node_list = []
        for node_id, data in self.graph.nodes(data=True):
            # Build a flattened row
            row = {
                "id": node_id,
                "type": data.get("type", ""),
                "name": data.get("name", ""),
                "lang": data.get("lang", ""),
                "file": data.get("file", ""),
            }
            node_list.append(row)

        edge_list = []
        for src, dst, data in self.graph.edges(data=True):
            edge_list.append({
                "src": src,
                "dst": dst,
                "type": data.get("type", "")
            })

        lines = [
            "# Graph Export",
            f"nodes[{len(node_list)}]{{id,type,name,lang,file}}:"
        ]
        for n in node_list:
            row = [n["id"], n["type"], n["name"], n["lang"], n["file"]]
            escaped = [EnhancedToon._escape_toon_field(str(f)) for f in row]
            lines.append("  "   ",".join(escaped))

        lines.append(f"edges[{len(edge_list)}]{{src,dst,type}}:")
        for e in edge_list:
            row = [e["src"], e["dst"], e["type"]]
            escaped = [EnhancedToon._escape_toon_field(str(f)) for f in row]
            lines.append("  "   ",".join(escaped))

        toon_text = "\n".join(lines)
        path = os.path.join(self.graph_dir, "graph.toon")
        with open(path, "w", encoding="utf-8") as f:
            f.write(toon_text)
        return path

    def _parse_node_id(self, node_id: str):
        """
        Parse node_id into components and return a dict of canonical attributes.

        Supported formats:
        1) lang::node_type::file_path::name    (full chunk/file tied)
        2) lang::node_type::name                (file-less canonical form - Option B)
        3) file::<file_path>                    (file nodes)
        otherwise -> fallback to Symbol
        """
        attrs = {
            "type": "Symbol",
            "name": node_id,
            "lang": "",
            "file": ""
        }
        if not isinstance(node_id, str):
            return attrs
        parts = node_id.split("::")
        if len(parts) == 4:
            lang, node_type, file_path, name = parts
            attrs["type"] = node_type
            attrs["name"] = name
            attrs["lang"] = lang
            attrs["file"] = file_path
        elif len(parts) == 3:
            # Option B canonical file-less representation: lang::kind::entity_name
            lang, node_type, name = parts
            attrs["type"] = node_type
            attrs["name"] = name
            attrs["lang"] = lang
            attrs["file"] = ""
        elif len(parts) == 2 and parts[0] == "file":
            # file::<file_path>
            attrs["type"] = "File"
            attrs["name"] = parts[1]
            attrs["file"] = parts[1]
            attrs["lang"] = ""
        else:
            # fallback: try to be helpful by guessing the name
            if len(parts) >= 1:
                attrs["name"] = parts[-1]
        return attrs

    def _ensure_node_exists(self, node_id: str):
        """
        Ensure a node with node_id exists in the graph.
        If it doesn't, create it using sensible attributes derived from the node_id format.
        """
        if node_id in self.graph.nodes:
            return
        parsed = self._parse_node_id(node_id)
        # Use add_node (keeps behaviour consistent)
        self.add_node(node_id, **{
            "type": parsed.get("type", "Symbol"),
            "name": parsed.get("name", node_id),
            "lang": parsed.get("lang", ""),
            "file": parsed.get("file", "")
        })

    def add_edge(self, src: str, dst: str, edge_type: str, **attrs):
        """
        Add an edge but first ensure both source and destination nodes exist and have basic attributes.
        This prevents the creation of attribute-less nodes and makes semantic edges meaningful.
        """
        try:
            # Ensure source node exists (create minimal entry if missing)
            if src not in self.graph.nodes:
                self._ensure_node_exists(src)

            # Ensure destination node exists (create minimal entry if missing)
            if dst not in self.graph.nodes:
                self._ensure_node_exists(dst)

            # Finally add the edge with type and any extra attrs
            self.graph.add_edge(src, dst, type=edge_type, **attrs)
        except Exception:
            # Keep behaviour non-fatal for indexing runs - log if you have a logger available
            # fallback: still attempt to add the edge
            try:
                self.graph.add_edge(src, dst, type=edge_type, **attrs)
            except Exception:
                # swallow; graph should remain usable
                pass

    def get_node_relationships(self, node_id: str) -> dict:
        """
        Get comprehensive relationships for a node in a structured format.
        Returns: {
            'outgoing': [(neighbor_id, edge_type, neighbor_data)],
            'incoming': [(predecessor_id, edge_type, predecessor_data)],
            'cross_file': list of cross-file relationships,
            'same_file': list of same-file relationships
        }
        """
        if node_id not in self.graph:
            return {}

        result = {
            'outgoing': [],
            'incoming': [],
            'cross_file': [],
            'same_file': []
        }

        node_file = self.graph.nodes[node_id].get('file', '')

        # Outgoing relationships
        for neighbor_id in self.graph.successors(node_id):
            edge_data = self.graph.edges[node_id, neighbor_id]
            neighbor_data = self.graph.nodes[neighbor_id]
            edge_type = edge_data.get('type', 'related')

            result['outgoing'].append((neighbor_id, edge_type, neighbor_data))

            # Cross-file classification
            neighbor_file = neighbor_data.get('file', '')
            rel_info = f"{edge_type}:{neighbor_data.get('type', '')}:{neighbor_data.get('name', '')}"

            if neighbor_file and neighbor_file != node_file:
                result['cross_file'].append(f"{rel_info}@{neighbor_file}")
            else:
                result['same_file'].append(rel_info)

        # Incoming relationships
        for predecessor_id in self.graph.predecessors(node_id):
            edge_data = self.graph.edges[predecessor_id, node_id]
            predecessor_data = self.graph.nodes[predecessor_id]
            edge_type = edge_data.get('type', 'related')

            result['incoming'].append((predecessor_id, edge_type, predecessor_data))

        return result

    def get_file_entities(self, file_path: str) -> list:
        """Get all entities in a file in compact format."""
        file_node_id = f"file::{file_path}"
        if file_node_id not in self.graph:
            return []

        entities = []
        for neighbor_id, neighbor_data in self.neighbors(file_node_id):
            entity_type = neighbor_data.get('type', '')
            entity_name = neighbor_data.get('name', '')
            if entity_type and entity_name:
                entities.append(f"{entity_type}:{entity_name}")

        return entities

    def find_nodes_by_attributes(self, **filters):
        """Find nodes by attributes (not just exact matches)."""
        for node_id, data in self.graph.nodes(data=True):
            if all(data.get(k) == v for k, v in filters.items()):
                yield node_id, data

    def get_node_by_name_and_file(self, name: str, file_path: str, node_type: str = None, lang: str = None):
        """Find a node by name and file path with optional type/language filtering."""
        candidates = []
        for node_id, data in self.graph.nodes(data=True):
            if (data.get('name') == name and
                data.get('file') == file_path and
                (node_type is None or data.get('type') == node_type) and
                (lang is None or data.get('lang') == lang)):
                candidates.append((node_id, data))
        return candidates

    def get_cross_file_relationships(self, node_id: str, current_file: str):
        """Get relationships that cross file boundaries."""
        if node_id not in self.graph.nodes:
            return [], []

        cross_file_out = []
        cross_file_in = []

        # Outgoing cross-file relationships
        for neighbor_id in self.graph.successors(node_id):
            if neighbor_id in self.graph.nodes:
                neighbor_file = self.graph.nodes[neighbor_id].get('file', '')
                if neighbor_file and neighbor_file != current_file:
                    edge_data = self.graph.edges[node_id, neighbor_id]
                    rel_type = edge_data.get('type', 'related')
                    neighbor_name = self.graph.nodes[neighbor_id].get('name', '')
                    neighbor_type = self.graph.nodes[neighbor_id].get('type', '')
                    cross_file_out.append(f"{rel_type}:{neighbor_type}:{neighbor_name}@{neighbor_file}")

        # Incoming cross-file relationships
        for predecessor_id in self.graph.predecessors(node_id):
            if predecessor_id in self.graph.nodes:
                pred_file = self.graph.nodes[predecessor_id].get('file', '')
                if pred_file and pred_file != current_file:
                    edge_data = self.graph.edges[predecessor_id, node_id]
                    rel_type = edge_data.get('type', 'related')
                    pred_name = self.graph.nodes[predecessor_id].get('name', '')
                    pred_type = self.graph.nodes[predecessor_id].get('type', '')
                    cross_file_in.append(f"{rel_type}:{pred_type}:{pred_name}@{pred_file}")

        return cross_file_out, cross_file_in
