import { useMemo } from "react";
import type { DirectiveStep } from "../../lib/api";
import { StepNode } from "./StepNode";
import {
  OrchestratorStepNode,
  type OrchestratorStepType,
  type OrchestratorStepStatus,
} from "./OrchestratorStepNode";

export interface VirtualStep {
  type: OrchestratorStepType;
  taskId: string;
  status: OrchestratorStepStatus;
  label: string;
  hasQuestions?: boolean;
}

export interface SpecializedStep {
  id: string;
  name: string;
  type: "orchestrator" | "completion";
  taskId: string;
  status: "running" | "completed";
}

interface DirectiveDAGProps {
  steps: DirectiveStep[];
  specializedSteps?: SpecializedStep[];
  onComplete?: (stepId: string) => void;
  onFail?: (stepId: string) => void;
  onSkip?: (stepId: string) => void;
}

interface Layer {
  steps: DirectiveStep[];
}

/** Types that should appear before the regular DAG steps */
const BEFORE_TYPES = new Set<OrchestratorStepType>([
  "planning",
  "replanning",
  "plan-orders",
]);

function topoSort(steps: DirectiveStep[]): Layer[] {
  if (steps.length === 0) return [];

  // Group steps by orderIndex — each unique orderIndex is one execution phase
  const byOrder = new Map<number, DirectiveStep[]>();
  for (const step of steps) {
    const group = byOrder.get(step.orderIndex) ?? [];
    group.push(step);
    byOrder.set(step.orderIndex, group);
  }

  // Sort groups by ascending orderIndex
  const sortedKeys = [...byOrder.keys()].sort((a, b) => a - b);
  return sortedKeys.map((key) => ({
    steps: byOrder.get(key)!.sort((a, b) => a.name.localeCompare(b.name)),
  }));
}

export function DirectiveDAG({ steps, specializedSteps, onComplete, onFail, onSkip }: DirectiveDAGProps) {
  const layers = useMemo(() => topoSort(steps), [steps]);

  const orchestratorSteps = specializedSteps?.filter(s => s.type === "orchestrator") ?? [];
  const completionSteps = specializedSteps?.filter(s => s.type === "completion") ?? [];

  if (steps.length === 0 && orchestratorSteps.length === 0 && completionSteps.length === 0) {
    return (
      <div className="text-center py-8 text-[#7788aa] font-mono text-sm">
        No steps yet. Add steps to build the DAG.
      </div>
    );
  }

  return (
    <div className="flex flex-col gap-4 items-center py-4">
      {/* Orchestrator steps (Planning/Cleanup/Orders) - rendered above regular steps */}
      {orchestratorSteps.map(step => (
        <SpecializedStepNode key={step.id} step={step} />
      ))}

      {/* Connector line if both orchestrator step and regular steps exist */}
      {orchestratorSteps.length > 0 && layers.length > 0 && (
        <div className="flex justify-center py-1">
          <div className="w-px h-4 bg-[rgba(117,170,252,0.2)]" />
        </div>
      )}

      {/* Regular step layers */}
      {layers.map((layer, layerIdx) => (
        <div key={layerIdx}>
          {layerIdx > 0 && (
            <div className="flex justify-center py-1">
              <div className="w-px h-4 bg-[#2a3a5a]" />
            </div>
          )}
          <div className="flex flex-wrap gap-3 justify-center">
            {afterSteps.map((vs) => (
              <OrchestratorStepNode
                key={`${vs.type}-${vs.taskId}`}
                type={vs.type}
                taskId={vs.taskId}
                status={vs.status}
                label={vs.label}
                hasQuestions={vs.hasQuestions}
              />
            ))}
          </div>
        </div>
      ))}

      {/* Connector line if both regular steps and completion step exist */}
      {completionSteps.length > 0 && layers.length > 0 && (
        <div className="flex justify-center py-1">
          <div className="w-px h-4 bg-[rgba(117,170,252,0.2)]" />
        </div>
      )}

      {/* Completion steps (PR creation) - rendered below regular steps */}
      {completionSteps.map(step => (
        <SpecializedStepNode key={step.id} step={step} />
      ))}
    </div>
  );
}

function SpecializedStepNode({ step }: { step: SpecializedStep }) {
  const themeColors = step.type === "orchestrator"
    ? {
        bg: "bg-[#1a1a30]",
        border: "border-[rgba(117,170,252,0.3)]",
        text: "text-[#75aafc]",
        dot: "bg-[#75aafc]",
        label: step.name.startsWith("Cleanup") ? "CLEANUP"
             : step.name.startsWith("Pick up") ? "ORDERS"
             : "PLANNING",
      }
    : {
        bg: "bg-[#1a1a10]",
        border: "border-yellow-900/50",
        text: "text-yellow-400",
        dot: "bg-yellow-400",
        label: "PR",
      };

  return (
    <div className={`flex items-center gap-2 px-3 py-2 ${themeColors.bg} border ${themeColors.border} rounded-lg mx-2`}>
      <span className={`inline-block w-2 h-2 rounded-full ${themeColors.dot} animate-pulse`} />
      <span className={`text-[9px] font-mono uppercase tracking-wide ${themeColors.text} opacity-60`}>
        {themeColors.label}
      </span>
      <span className={`text-[11px] font-mono ${themeColors.text} flex-1 truncate`}>
        {step.name}
      </span>
      <a
        href={`/mesh/${step.taskId}`}
        className="text-[9px] font-mono text-[#556677] hover:text-white underline"
      >
        View task
      </a>
    </div>
  );
}