import { useMemo } from "react";
import TechniqueCell from "./TechniqueCell";
import type { TechniqueSummary } from "../api/techniques";

interface AttackMatrixProps {
  techniques: TechniqueSummary[];
}

// MITRE ATT&CK Enterprise Tactics in order
const TACTIC_ORDER = [
  "reconnaissance",
  "resource-development",
  "initial-access",
  "execution",
  "persistence",
  "privilege-escalation",
  "defense-evasion",
  "credential-access",
  "discovery",
  "lateral-movement",
  "collection",
  "command-and-control",
  "exfiltration",
  "impact",
];

const formatTacticName = (tactic: string): string => {
  return tactic
    .split("-")
    .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
    .join(" ");
};

export default function AttackMatrix({ techniques }: AttackMatrixProps) {
  // Group techniques by tactic
  const groupedByTactic = useMemo(() => {
    const groups: Record<string, TechniqueSummary[]> = {};

    for (const tech of techniques) {
      // A technique can belong to multiple tactics (comma-separated)
      const tactics = tech.tactic
        ? tech.tactic.split(",").map((t) => t.trim().toLowerCase())
        : ["unknown"];

      for (const tactic of tactics) {
        if (!groups[tactic]) {
          groups[tactic] = [];
        }
        groups[tactic].push(tech);
      }
    }

    // Sort techniques within each tactic by mitre_id
    for (const tactic of Object.keys(groups)) {
      groups[tactic].sort((a, b) => a.mitre_id.localeCompare(b.mitre_id));
    }

    return groups;
  }, [techniques]);

  // Get ordered tactics that have techniques
  const orderedTactics = useMemo(() => {
    const tacticSet = new Set(Object.keys(groupedByTactic));
    const ordered = TACTIC_ORDER.filter((t) => tacticSet.has(t));
    // Add any unknown tactics at the end
    const remaining = Array.from(tacticSet).filter((t) => !TACTIC_ORDER.includes(t));
    return [...ordered, ...remaining];
  }, [groupedByTactic]);

  if (techniques.length === 0) {
    return (
      <div className="rounded-xl border border-gray-800 bg-gray-900 p-8 text-center">
        <p className="text-gray-400">No techniques found matching your filters</p>
      </div>
    );
  }

  return (
    <div className="overflow-x-auto rounded-xl border border-gray-800 bg-gray-900">
      <div className="min-w-max p-4">
        <div className="flex gap-3">
          {orderedTactics.map((tactic) => (
            <div key={tactic} className="w-48 flex-shrink-0">
              {/* Tactic header */}
              <div className="mb-3 rounded-lg bg-gray-800 px-3 py-2">
                <h3 className="text-center text-sm font-semibold text-cyan-400">
                  {formatTacticName(tactic)}
                </h3>
                <p className="mt-0.5 text-center text-xs text-gray-500">
                  {groupedByTactic[tactic]?.length || 0} techniques
                </p>
              </div>

              {/* Technique cells */}
              <div className="space-y-2">
                {groupedByTactic[tactic]?.map((tech) => (
                  <TechniqueCell
                    key={`${tactic}-${tech.mitre_id}`}
                    mitreId={tech.mitre_id}
                    name={tech.name}
                    status={tech.status_global}
                    reviewRequired={tech.review_required}
                  />
                ))}
              </div>
            </div>
          ))}
        </div>
      </div>
    </div>
  );
}