import json
from collections import defaultdict
from sqlalchemy import create_engine, text

# 🔧 Adjust this to your environment
DB_URL = "postgresql+psycopg2://postgres:postgres@192.168.1.86:5432/digital_aeros_fixed"

def export_rbd_json(project_id, output_file="rbd.json"):
    engine = create_engine(DB_URL)

    with engine.connect() as conn:
        # --- Project info ---
        project = conn.execute(
            text('SELECT "ProjectId","ProjectName" FROM public."Projects" WHERE "ProjectId" = :pid'),
            {"pid": project_id}
        ).mappings().first()
        if not project:
            raise ValueError(f"❌ Project {project_id} not found")

        print(f"🔍 Loaded project: {project['ProjectName']} (ID={project_id})")

        # --- Schematics ---
        schematics = conn.execute(
            text('SELECT "SchematicId","SchematicName" FROM public."Schematics" WHERE "ProjectId" = :pid'),
            {"pid": project_id}
        ).mappings().all()

        all_schematics_json = {}

        for s in schematics:
            sid = s["SchematicId"]
            sname = s["SchematicName"]
            print(f"\n=== Processing schematic {sid} ({sname}) ===")

            # --- Nodes ---
            nodes_data = conn.execute(
                text('SELECT "NodeId","NodeUIId","NodeName","NodeTypeCode" '
                     'FROM public."RegularNodes" WHERE "SchematicId" = :sid'),
                {"sid": sid}
            ).mappings().all()

            # dict keyed by NodeId
            nodes = {str(n["NodeId"]): dict(n) for n in nodes_data}
            # map UIId -> NodeId
            uiid_to_id = {str(n["NodeUIId"]): str(n["NodeId"]) for n in nodes_data}

            print(f"🟦 Nodes: {len(nodes)}")

            # --- Connectors (edges) ---
            edges = conn.execute(
                text('SELECT "SourceUIId","SinkUIId" FROM public."Connectors" WHERE "SchematicId" = :sid'),
                {"sid": sid}
            ).mappings().all()

            children_map = defaultdict(list)
            for e in edges:
                src_ui = str(e["SourceUIId"])
                sink_ui = str(e["SinkUIId"])
                if src_ui in uiid_to_id and sink_ui in uiid_to_id:
                    src_id = uiid_to_id[src_ui]
                    sink_id = uiid_to_id[sink_ui]
                    children_map[src_id].append(sink_id)

            print(f"🔗 Connectors: {len(edges)}")

            # --- Redundancies ---
            redundancies = conn.execute(
                text('SELECT * FROM public."Redundancies" WHERE "SchematicId" = :sid'),
                {"sid": sid}
            ).mappings().all()

            states = conn.execute(
                text('SELECT * FROM public."RedundancyStates" WHERE "RedundancyId" IN '
                     '(SELECT "RedundancyId" FROM public."Redundancies" WHERE "SchematicId" = :sid)'),
                {"sid": sid}
            ).mappings().all()

            state_cells = conn.execute(
                text('SELECT * FROM public."RedundancyStateCells" WHERE "RedundancyStateId" IN '
                     '(SELECT "StateId" FROM public."RedundancyStates" WHERE "RedundancyId" IN '
                     '(SELECT "RedundancyId" FROM public."Redundancies" WHERE "SchematicId" = :sid))'),
                {"sid": sid}
            ).mappings().all()

            # Build redundancy structure
            redundancy_map = {}
            for r in redundancies:
                rid = str(r["RedundancyId"])
                redundancy_map[rid] = {
                    "name": r["RedundancyName"],
                    "type": "standby" if r["IsKNStandby"] else "kofn",
                    "min_required": r["MinRequiredNodes"],
                    "switch_delay": r["SwitchDelay"],
                    "states": {}
                }

            for st in states:
                sid2 = str(st["StateId"])
                rid = str(st["RedundancyId"])
                redundancy_map[rid]["states"][sid2] = {
                    "priority": st["StatePriority"],
                    "duration": st["StateDuration"],
                    "cells": []
                }

            for c in state_cells:
                sid2 = str(c["RedundancyStateId"])
                for r_id, r in redundancy_map.items():
                    if sid2 in r["states"]:
                        node_id = str(c["InvolvedNodeId"])
                        if node_id in nodes:  # only link if valid
                            r["states"][sid2]["cells"].append({
                                "node_id": node_id,
                                "status": c["NodeStatusCode"]
                            })
                        break

            # --- Build JSON nodes ---
            json_nodes = {}
            for node_id, n in nodes.items():
                node_type = str(n["NodeTypeCode"]).lower()
                node_json = {
                    "type": node_type,
                    "name": n["NodeName"],
                    "children": children_map.get(node_id, []),  # now NodeId-based
                    "availability": 1.0   # ✅ default availability
                }
                json_nodes[node_id] = node_json

            # --- Root detection ---
            all_children = {c for lst in children_map.values() for c in lst}
            root_candidates = [nid for nid in nodes.keys() if nid not in all_children]
            root = root_candidates[0] if root_candidates else None
            print(f"🌳 Root node detected: {root}")

            # --- Assemble schematic JSON ---
            schematic_json = {
                "schematic_id": sid,
                "schematic_name": sname,
                "root": str(root) if root else None,
                "nodes": json_nodes,
                "redundancies": redundancy_map
            }
            all_schematics_json[str(sid)] = schematic_json

            print(f"✅ Finished schematic {sid}")

        # --- Final JSON ---
        final_json = {
            "project_id": project["ProjectId"],
            "project_name": project["ProjectName"],
            "schematics": all_schematics_json
        }

        with open(output_file, "w") as f:
            json.dump(final_json, f, indent=2)

        print(f"\n🎉 Exported simplified RBD JSON (NodeId-based, availability=1.0) "
              f"for Project {project_id} → {output_file}")


if __name__ == "__main__":
    export_rbd_json(project_id=70, output_file="project_70_rbd.json")