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, **attrs):
        self.graph.add_node(node_id, **attrs)

    # --- 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