"use client";

import { useCallback, useEffect, useRef, useState } from "react";
import { useDiagramStore, type EntropykNodeData } from "@/store/diagramStore";
import { fetchModule } from "@/lib/api";
import type { Node, Edge } from "@xyflow/react";

export interface Doc {
  id: string; // "workspace" or `module:${name}`
  name: string; // scenario name, or class name (e.g. "BaseChiller")
  kind: "workspace" | "module";
  nodes: Node<EntropykNodeData>[];
  edges: Edge[];
  dirty: boolean;
}

/**
 * Multi-document manager — Modelica / OMEdit style.
 *
 * Classes (.ekmod) are first-class: each open class is its own tab named
 * `ClassName.ekmod`. The workspace tab is the solvable scenario assembly
 * (CLI ScenarioConfig) where class instances are placed.
 *
 *   - openModule(name)      — open / focus a class tab (loads .ekmod)
 *   - openEmptyModule(name) — open a freshly created empty class
 *   - renameWorkspace(name) — rename the scenario without clearing
 *   - newScenario(name, clear) — name (and optionally clear) the workspace
 */
export function useDocuments() {
  const [docs, setDocs] = useState<Doc[]>([
    {
      id: "workspace",
      name: "Workspace",
      kind: "workspace",
      nodes: [],
      edges: [],
      dirty: false,
    },
  ]);
  const [activeId, setActiveId] = useState<string>("workspace");

  // Refs to read latest store state inside async callbacks without re-subscribing.
  const storeNodes = useRef(useDiagramStore.getState().nodes);
  const storeEdges = useRef(useDiagramStore.getState().edges);
  useEffect(() => {
    return useDiagramStore.subscribe((s) => {
      storeNodes.current = s.nodes;
      storeEdges.current = s.edges;
    });
  }, []);

  const activeDoc = docs.find((d) => d.id === activeId) ?? docs[0];

  /** Snapshot the live canvas state into a doc record. */
  const snapshotCurrent = useCallback(() => {
    setDocs((prev) =>
      prev.map((d) =>
        d.id === activeId
          ? { ...d, nodes: storeNodes.current as Node<EntropykNodeData>[], edges: storeEdges.current }
          : d,
      ),
    );
  }, [activeId]);

  /** Switch the active tab — save current canvas, restore target canvas. */
  const switchTo = useCallback(
    (id: string) => {
      if (id === activeId) return;
      // Save current canvas into its doc.
      setDocs((prev) =>
        prev.map((d) =>
          d.id === activeId
            ? { ...d, nodes: storeNodes.current as Node<EntropykNodeData>[], edges: storeEdges.current }
            : d,
        ),
      );
      const target = docs.find((d) => d.id === id);
      if (!target) return;
      useDiagramStore.setState({
        nodes: target.nodes,
        edges: target.edges,
        selectedNodeId: null,
      });
      setActiveId(id);
    },
    [activeId, docs],
  );

  /** Snapshot current canvas, then open a class tab with given graph. */
  const pushClassTab = useCallback(
    (moduleName: string, nodes: Node<EntropykNodeData>[], edges: Edge[], dirty: boolean) => {
      const id = `module:${moduleName}`;
      setDocs((prev) => {
        const saved = prev.map((d) =>
          d.id === activeId
            ? {
                ...d,
                nodes: storeNodes.current as Node<EntropykNodeData>[],
                edges: storeEdges.current,
              }
            : d,
        );
        if (saved.some((d) => d.id === id)) {
          return saved.map((d) =>
            d.id === id ? { ...d, nodes, edges, dirty, name: moduleName } : d,
          );
        }
        return [
          ...saved,
          { id, name: moduleName, kind: "module" as const, nodes, edges, dirty },
        ];
      });
      useDiagramStore.setState({ nodes, edges, selectedNodeId: null });
      setActiveId(id);
    },
    [activeId],
  );

  /** Open a class (.ekmod) in its own tab. If already open, just focus it. */
  const openModule = useCallback(
    async (moduleName: string) => {
      const id = `module:${moduleName}`;
      if (docs.some((d) => d.id === id)) {
        switchTo(id);
        return;
      }
      const modData = await fetchModule(moduleName);
      const { nodes: newNodes, edges: newEdges } = parseModuleData(modData);
      pushClassTab(moduleName, newNodes, newEdges, false);
    },
    [docs, switchTo, pushClassTab],
  );

  /** Open a newly created empty class without a round-trip fetch. */
  const openEmptyModule = useCallback(
    (moduleName: string) => {
      const id = `module:${moduleName}`;
      if (docs.some((d) => d.id === id)) {
        switchTo(id);
        return;
      }
      pushClassTab(moduleName, [], [], false);
    },
    [docs, switchTo, pushClassTab],
  );

  /** Rename the workspace scenario (CLI config name). */
  const renameWorkspace = useCallback((name: string) => {
    const clean = name.trim();
    if (!clean) return;
    setDocs((prev) =>
      prev.map((d) => (d.id === "workspace" ? { ...d, name: clean } : d)),
    );
  }, []);

  /** Create / rename scenario; optionally clear the workspace canvas. */
  const newScenario = useCallback(
    (name: string, clearCanvas: boolean) => {
      const clean = name.trim();
      if (!clean) return;

      // Persist whatever tab we are leaving.
      setDocs((prev) => {
        const withSnap = prev.map((d) =>
          d.id === activeId
            ? {
                ...d,
                nodes: storeNodes.current as Node<EntropykNodeData>[],
                edges: storeEdges.current,
              }
            : d,
        );
        return withSnap.map((d) =>
          d.id === "workspace"
            ? {
                ...d,
                name: clean,
                nodes: clearCanvas ? [] : d.nodes,
                edges: clearCanvas ? [] : d.edges,
                dirty: false,
              }
            : d,
        );
      });

      if (clearCanvas) {
        useDiagramStore.getState().clear();
      } else if (activeId !== "workspace") {
        const ws = docs.find((d) => d.id === "workspace");
        useDiagramStore.setState({
          nodes: ws?.nodes ?? [],
          edges: ws?.edges ?? [],
          selectedNodeId: null,
        });
      }
      setActiveId("workspace");
    },
    [activeId, docs],
  );

  /** Close a tab. Workspace cannot be closed. */
  const closeDoc = useCallback(
    (id: string) => {
      if (id === "workspace") return;

      const idx = docs.findIndex((d) => d.id === id);
      if (idx < 0) return;

      // Snapshot live canvas into the tab being closed (if it's active).
      const closing: Doc =
        id === activeId
          ? {
              ...docs[idx],
              nodes: storeNodes.current as Node<EntropykNodeData>[],
              edges: storeEdges.current,
            }
          : docs[idx];

      const next = docs
        .map((d) => (d.id === id ? closing : d))
        .filter((d) => d.id !== id);

      const switchingAway = id === activeId;
      const fallback = switchingAway
        ? (next[Math.max(0, idx - 1)] ?? next[0])
        : null;

      // Pure state updates only — never call another setState inside a setDocs updater.
      setDocs(next);
      if (fallback) {
        setActiveId(fallback.id);
        useDiagramStore.setState({
          nodes: fallback.nodes,
          edges: fallback.edges,
          selectedNodeId: null,
        });
      }
    },
    [activeId, docs],
  );

  /** Mark the active doc dirty (after edits). */
  const markDirty = useCallback(() => {
    setDocs((prev) =>
      prev.map((d) => (d.id === activeId ? { ...d, dirty: true } : d)),
    );
  }, [activeId]);

  /** Mark a doc clean (after save). */
  const markClean = useCallback((id: string) => {
    setDocs((prev) => prev.map((d) => (d.id === id ? { ...d, dirty: false } : d)));
  }, []);

  /** Reload a doc's content from server (after save). */
  const reloadDoc = useCallback(
    async (id: string) => {
      if (!id.startsWith("module:")) return;
      const name = id.slice("module:".length);
      const modData = await fetchModule(name);
      const { nodes: nn, edges: ne } = parseModuleData(modData);
      setDocs((prev) =>
        prev.map((d) => (d.id === id ? { ...d, nodes: nn, edges: ne, dirty: false } : d)),
      );
      if (id === activeId) {
        useDiagramStore.setState({ nodes: nn, edges: ne, selectedNodeId: null });
      }
    },
    [activeId],
  );

  // Periodically snapshot the canvas into the active doc so switching tabs
  // always restores the very latest state (covers drag/edit ops that don't
  // go through explicit save handlers).
  useEffect(() => {
    const t = setInterval(() => {
      snapshotCurrent();
    }, 1500);
    return () => clearInterval(t);
  }, [snapshotCurrent]);

  const workspaceName =
    docs.find((d) => d.id === "workspace")?.name ?? "Workspace";

  return {
    docs,
    activeDoc,
    activeId,
    workspaceName,
    switchTo,
    openModule,
    openEmptyModule,
    closeDoc,
    markDirty,
    markClean,
    reloadDoc,
    renameWorkspace,
    newScenario,
  };
}

/* ── Helpers ──────────────────────────────────────────────────────── */

function parseModuleData(modData: unknown): {
  nodes: Node<EntropykNodeData>[];
  edges: Edge[];
} {
  const newNodes: Node<EntropykNodeData>[] = [];
  const newEdges: Edge[] = [];
  const nameToId: Record<string, string> = {};

  if (!modData || typeof modData !== "object") {
    return { nodes: [], edges: [] };
  }

  const modObj = modData as {
    components?: Array<Record<string, unknown>>;
    instances?: Array<{
      of: string;
      name: string;
      circuit?: number;
      params?: Record<string, unknown>;
    }>;
    edges?: Array<{ from: string; to: string }>;
    connections?: Array<{ from: string; to: string }>;
    ports?: Record<string, string>;
  };

  const atomicComps = modObj.components ?? [];
  atomicComps.forEach((c, idx) => {
    const cName = String(c.name ?? `comp_${idx}`);
    const cType = String(c.type ?? "Generic");
    const nodeId = `mod_node_${Date.now()}_${idx}`;
    nameToId[cName] = nodeId;
    const { type: _, name: __, ...params } = c;
    newNodes.push({
      id: nodeId,
      type: "entropykNode",
      position: {
        x: 100 + (idx % 3) * 260,
        y: 100 + Math.floor(idx / 3) * 180,
      },
      data: {
        type: cType,
        kind: cType,
        name: cName,
        circuit: 0,
        rotation: 0,
        flipH: false,
        flipV: false,
        params: params as Record<string, number | string | boolean>,
        status: "ok",
      },
    });
  });

  const moduleInsts = modObj.instances ?? [];
  moduleInsts.forEach((inst, idx) => {
    const instName = String(inst.name ?? `inst_${idx}`);
    const ofModule = String(inst.of ?? "BaseChiller");
    const nodeId = `mod_inst_${Date.now()}_${idx}`;
    nameToId[instName] = nodeId;
    newNodes.push({
      id: nodeId,
      type: "entropykNode",
      position: {
        x: 100 + ((idx + atomicComps.length) % 3) * 280,
        y: 100 + Math.floor((idx + atomicComps.length) / 3) * 200,
      },
      data: {
        type: "ModuleInstance",
        kind: "ModuleInstance",
        name: instName,
        moduleName: ofModule,
        circuit: inst.circuit ?? 0,
        rotation: 0,
        flipH: false,
        flipV: false,
        params: (inst.params as Record<string, number | string | boolean>) ?? {},
        status: "ok",
      },
    });
  });

  const atomicEdg = modObj.edges ?? [];
  atomicEdg.forEach((e, idx) => {
    const [fromComp, fromPort] = e.from.split(":");
    const [toComp, toPort] = e.to.split(":");
    const sourceId = nameToId[fromComp];
    const targetId = nameToId[toComp];
    if (sourceId && targetId) {
      newEdges.push({
        id: `mod_edge_${Date.now()}_${idx}`,
        source: sourceId,
        target: targetId,
        sourceHandle: fromPort,
        targetHandle: toPort,
        type: "entropykEdge",
      });
    }
  });

  const instConn = modObj.connections ?? [];
  instConn.forEach((c, idx) => {
    const fromParts = c.from.split(".");
    const toParts = c.to.split(".");
    const fromComp = fromParts[0];
    const fromPort = fromParts.slice(1).join(".");
    const toComp = toParts[0];
    const toPort = toParts.slice(1).join(".");
    const sourceId = nameToId[fromComp];
    const targetId = nameToId[toComp];
    if (sourceId && targetId) {
      newEdges.push({
        id: `mod_conn_${Date.now()}_${idx}`,
        source: sourceId,
        target: targetId,
        sourceHandle: fromPort,
        targetHandle: toPort,
        type: "entropykEdge",
      });
    }
  });

  return { nodes: newNodes, edges: newEdges };
}
