import { useState, useMemo, useEffect } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
  X,
  Search,
  Plus,
  Loader2,
  CheckCircle,
  FlaskConical,
  BookOpen,
  ArrowLeft,
  ChevronRight,
  Filter,
} from "lucide-react";
import { getTests } from "../api/tests";
import { addTestToCampaign } from "../api/campaigns";
import { getTemplates, getTemplateById } from "../api/test-templates";
import { createTestFromTemplate } from "../api/tests";
import type { Test, TestState, TestTemplateSummary } from "../types/models";

/* ── helpers ─────────────────────────────────────────────────────── */

const stateBadge: Record<TestState, string> = {
  draft: "bg-gray-800/50 text-gray-400 border-gray-600/30",
  red_executing: "bg-orange-900/50 text-orange-400 border-orange-500/30",
  blue_evaluating: "bg-indigo-900/50 text-indigo-400 border-indigo-500/30",
  in_review: "bg-blue-900/50 text-blue-400 border-blue-500/30",
  validated: "bg-green-900/50 text-green-400 border-green-500/30",
  rejected: "bg-red-900/50 text-red-400 border-red-500/30",
};

const SEVERITY_COLORS: Record<string, string> = {
  critical: "text-red-400 border-red-500/30 bg-red-500/10",
  high:     "text-orange-400 border-orange-500/30 bg-orange-500/10",
  medium:   "text-yellow-400 border-yellow-500/30 bg-yellow-500/10",
  low:      "text-blue-400 border-blue-500/30 bg-blue-500/10",
};

const SOURCE_LABELS: Record<string, string> = {
  atomic_red_team: "Atomic Red Team",
  mitre:           "MITRE",
  custom:          "Custom",
};

type Tab = "existing" | "template";
type TemplateStep = "list" | "form";

/* ── props ───────────────────────────────────────────────────────── */

interface AddTestToCampaignModalProps {
  campaignId: string;
  existingTestIds: string[];
  open: boolean;
  onClose: () => void;
  onSuccess: () => void;
}

/* ── component ───────────────────────────────────────────────────── */

export default function AddTestToCampaignModal({
  campaignId,
  existingTestIds,
  open,
  onClose,
  onSuccess,
}: AddTestToCampaignModalProps) {
  const queryClient = useQueryClient();

  // ── shared state ────────────────────────────────────────────────
  const [tab, setTab] = useState<Tab>("template");
  const [addedCount, setAddedCount] = useState(0);

  // ── existing-test tab ───────────────────────────────────────────
  const [existingSearch, setExistingSearch] = useState("");
  const [addedIds, setAddedIds] = useState<Set<string>>(new Set());

  // ── template tab ────────────────────────────────────────────────
  const [templateStep, setTemplateStep] = useState<TemplateStep>("list");
  const [templateSearch, setTemplateSearch] = useState("");
  const [filterPlatform, setFilterPlatform] = useState("");
  const [filterSeverity, setFilterSeverity] = useState("");
  const [selectedTemplateId, setSelectedTemplateId] = useState<string | null>(null);

  // form fields (pre-filled from selected template)
  const [formName, setFormName] = useState("");
  const [formDescription, setFormDescription] = useState("");
  const [formPlatform, setFormPlatform] = useState("");
  const [formProcedure, setFormProcedure] = useState("");
  const [formTool, setFormTool] = useState("");

  // ── reset when closed ────────────────────────────────────────────
  useEffect(() => {
    if (!open) {
      setTab("template");
      setExistingSearch("");
      setAddedIds(new Set());
      setAddedCount(0);
      setTemplateStep("list");
      setTemplateSearch("");
      setFilterPlatform("");
      setFilterSeverity("");
      setSelectedTemplateId(null);
    }
  }, [open]);

  // ── queries ──────────────────────────────────────────────────────

  const { data: allTests, isLoading: testsLoading } = useQuery({
    queryKey: ["tests", "for-campaign-picker"],
    queryFn: () => getTests({ limit: 200 }),
    enabled: open && tab === "existing",
  });

  const { data: templates, isLoading: templatesLoading } = useQuery({
    queryKey: ["templates", "picker", filterPlatform, filterSeverity, templateSearch],
    queryFn: () =>
      getTemplates({
        search: templateSearch || undefined,
        platform: filterPlatform || undefined,
        severity: filterSeverity || undefined,
        limit: 100,
      }),
    enabled: open && tab === "template",
    staleTime: 60_000,
  });

  const { data: fullTemplate, isLoading: fullTemplateLoading } = useQuery({
    queryKey: ["template-detail", selectedTemplateId],
    queryFn: () => getTemplateById(selectedTemplateId!),
    enabled: !!selectedTemplateId,
  });

  // Pre-fill form when full template loads
  useEffect(() => {
    if (fullTemplate) {
      setFormName(fullTemplate.name);
      setFormDescription(fullTemplate.description || "");
      setFormPlatform(fullTemplate.platform || "");
      setFormProcedure(fullTemplate.attack_procedure || "");
      setFormTool(fullTemplate.tool_suggested || "");
    }
  }, [fullTemplate]);

  // ── mutations ─────────────────────────────────────────────────────

  /** Add an existing test directly to the campaign */
  const addExistingMutation = useMutation({
    mutationFn: (testId: string) =>
      addTestToCampaign(campaignId, { test_id: testId }),
    onSuccess: (_data, testId) => {
      setAddedIds((prev) => new Set(prev).add(testId));
      setAddedCount((n) => n + 1);
      queryClient.invalidateQueries({ queryKey: ["campaign", campaignId] });
      onSuccess();
    },
  });

  /** Create test from template, then add to campaign */
  const createAndAddMutation = useMutation({
    mutationFn: async () => {
      if (!fullTemplate) throw new Error("No template loaded");
      const test = await createTestFromTemplate(
        selectedTemplateId!,
        fullTemplate.mitre_technique_id,
        {
          name: formName.trim() || undefined,
          description: formDescription.trim() || undefined,
          platform: formPlatform.trim() || undefined,
          procedure_text: formProcedure.trim() || undefined,
          tool_used: formTool.trim() || undefined,
        },
      );
      await addTestToCampaign(campaignId, { test_id: test.id });
      return test;
    },
    onSuccess: () => {
      setAddedCount((n) => n + 1);
      queryClient.invalidateQueries({ queryKey: ["campaign", campaignId] });
      onSuccess();
      // Back to template list, ready to add another
      setTemplateStep("list");
      setSelectedTemplateId(null);
    },
  });

  // ── derived data ──────────────────────────────────────────────────

  const filteredExisting = useMemo(() => {
    if (!allTests) return [];
    const alreadyIn = new Set([...existingTestIds, ...addedIds]);
    let list = allTests.filter((t) => !alreadyIn.has(t.id));
    if (existingSearch.trim()) {
      const q = existingSearch.toLowerCase();
      list = list.filter(
        (t) =>
          t.name.toLowerCase().includes(q) ||
          (t.technique_mitre_id && t.technique_mitre_id.toLowerCase().includes(q)) ||
          (t.technique_name && t.technique_name.toLowerCase().includes(q)),
      );
    }
    return list;
  }, [allTests, existingSearch, existingTestIds, addedIds]);

  const canSubmitForm =
    formName.trim().length > 0 && !!fullTemplate && !createAndAddMutation.isPending;

  if (!open) return null;

  /* ── render ──────────────────────────────────────────────────── */

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center bg-black/60 backdrop-blur-sm">
      <div className="flex w-full max-w-2xl flex-col rounded-xl border border-gray-700 bg-gray-900 shadow-2xl"
        style={{ maxHeight: "90vh" }}>

        {/* ── Header ── */}
        <div className="flex items-center justify-between border-b border-gray-800 px-6 py-4">
          {templateStep === "form" ? (
            <button
              onClick={() => { setTemplateStep("list"); setSelectedTemplateId(null); }}
              className="flex items-center gap-2 text-sm text-gray-400 hover:text-white transition-colors"
            >
              <ArrowLeft className="h-4 w-4" />
              Back to templates
            </button>
          ) : (
            <h2 className="text-lg font-semibold text-white">Add Test to Campaign</h2>
          )}
          <button
            onClick={onClose}
            className="rounded-lg p-1.5 text-gray-400 hover:bg-gray-800 hover:text-white transition-colors"
          >
            <X className="h-5 w-5" />
          </button>
        </div>

        {/* ── Tab bar (only on list step) ── */}
        {templateStep === "list" && (
          <div className="flex border-b border-gray-800">
            {(["template", "existing"] as Tab[]).map((t) => (
              <button
                key={t}
                onClick={() => setTab(t)}
                className={`flex items-center gap-2 px-5 py-3 text-sm font-medium transition-colors border-b-2 ${
                  tab === t
                    ? "border-cyan-500 text-cyan-400"
                    : "border-transparent text-gray-500 hover:text-gray-300"
                }`}
              >
                {t === "template" ? (
                  <><BookOpen className="h-4 w-4" /> From Template</>
                ) : (
                  <><FlaskConical className="h-4 w-4" /> Existing Test</>
                )}
              </button>
            ))}
          </div>
        )}

        {/* ── Body ── */}
        <div className="flex-1 overflow-y-auto">

          {/* ═══ TEMPLATE TAB — LIST ══════════════════════════════ */}
          {tab === "template" && templateStep === "list" && (
            <>
              {/* Filters */}
              <div className="border-b border-gray-800 px-5 py-3 space-y-2">
                <div className="relative">
                  <Search className="absolute left-3 top-1/2 h-4 w-4 -translate-y-1/2 text-gray-500" />
                  <input
                    type="text"
                    value={templateSearch}
                    onChange={(e) => setTemplateSearch(e.target.value)}
                    placeholder="Search templates by name or technique…"
                    className="w-full rounded-lg border border-gray-700 bg-gray-800 pl-9 pr-3 py-2 text-sm text-gray-300 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                  />
                </div>
                <div className="flex gap-2">
                  <div className="flex items-center gap-1.5 text-gray-500">
                    <Filter className="h-3.5 w-3.5" />
                  </div>
                  <select
                    value={filterPlatform}
                    onChange={(e) => setFilterPlatform(e.target.value)}
                    className="rounded-lg border border-gray-700 bg-gray-800 px-2.5 py-1.5 text-xs text-gray-300 focus:border-cyan-500 focus:outline-none"
                  >
                    <option value="">All platforms</option>
                    <option value="windows">Windows</option>
                    <option value="linux">Linux</option>
                    <option value="macos">macOS</option>
                  </select>
                  <select
                    value={filterSeverity}
                    onChange={(e) => setFilterSeverity(e.target.value)}
                    className="rounded-lg border border-gray-700 bg-gray-800 px-2.5 py-1.5 text-xs text-gray-300 focus:border-cyan-500 focus:outline-none"
                  >
                    <option value="">All severities</option>
                    <option value="critical">Critical</option>
                    <option value="high">High</option>
                    <option value="medium">Medium</option>
                    <option value="low">Low</option>
                  </select>
                </div>
              </div>

              {/* Template list */}
              <div className="px-5 py-3">
                {templatesLoading ? (
                  <div className="flex items-center justify-center py-12">
                    <Loader2 className="h-6 w-6 animate-spin text-cyan-400" />
                  </div>
                ) : !templates?.length ? (
                  <div className="flex flex-col items-center justify-center py-12 text-gray-500">
                    <BookOpen className="mb-2 h-8 w-8 text-gray-700" />
                    <p className="text-sm">No templates match your filters.</p>
                  </div>
                ) : (
                  <div className="space-y-1">
                    {templates.map((tmpl: TestTemplateSummary) => (
                      <button
                        key={tmpl.id}
                        onClick={() => {
                          setSelectedTemplateId(tmpl.id);
                          setTemplateStep("form");
                        }}
                        className="w-full flex items-center justify-between rounded-lg border border-gray-800 px-4 py-3 text-left hover:border-gray-700 hover:bg-gray-800/50 transition-colors"
                      >
                        <div className="flex-1 min-w-0">
                          <div className="flex items-center gap-2 flex-wrap">
                            <span className="text-sm font-medium text-gray-200 truncate">
                              {tmpl.name}
                            </span>
                            {tmpl.severity && (
                              <span className={`inline-flex rounded-full border px-1.5 py-0.5 text-[10px] font-medium capitalize ${SEVERITY_COLORS[tmpl.severity] ?? ""}`}>
                                {tmpl.severity}
                              </span>
                            )}
                          </div>
                          <div className="mt-0.5 flex items-center gap-3 text-xs text-gray-500">
                            <span className="font-mono text-cyan-400/70">
                              {tmpl.mitre_technique_id}
                            </span>
                            {tmpl.platform && (
                              <span className="capitalize">{tmpl.platform}</span>
                            )}
                            {tmpl.source && (
                              <span>{SOURCE_LABELS[tmpl.source] ?? tmpl.source}</span>
                            )}
                          </div>
                        </div>
                        <ChevronRight className="h-4 w-4 shrink-0 text-gray-600 ml-2" />
                      </button>
                    ))}
                  </div>
                )}
              </div>
            </>
          )}

          {/* ═══ TEMPLATE TAB — FORM ══════════════════════════════ */}
          {tab === "template" && templateStep === "form" && (
            <div className="px-6 py-5 space-y-4">
              {fullTemplateLoading ? (
                <div className="flex items-center justify-center py-12">
                  <Loader2 className="h-6 w-6 animate-spin text-cyan-400" />
                </div>
              ) : (
                <>
                  {/* Template info banner */}
                  {fullTemplate && (
                    <div className="rounded-lg border border-cyan-500/20 bg-cyan-900/10 px-4 py-3 text-xs text-cyan-400">
                      Template: <strong>{fullTemplate.name}</strong>
                      <span className="ml-2 font-mono text-gray-500">
                        {fullTemplate.mitre_technique_id}
                      </span>
                      {fullTemplate.source && (
                        <span className="ml-2 rounded-full border border-gray-600/30 bg-gray-800/50 px-2 py-0.5 text-gray-400">
                          {SOURCE_LABELS[fullTemplate.source] ?? fullTemplate.source}
                        </span>
                      )}
                    </div>
                  )}

                  {/* Name */}
                  <div>
                    <label className="mb-1.5 block text-sm font-medium text-gray-300">
                      Test Name <span className="text-red-400">*</span>
                    </label>
                    <input
                      value={formName}
                      onChange={(e) => setFormName(e.target.value)}
                      className="w-full rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-sm text-gray-200 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                      placeholder="Test name"
                    />
                  </div>

                  {/* Platform */}
                  <div>
                    <label className="mb-1.5 block text-sm font-medium text-gray-300">Platform</label>
                    <input
                      value={formPlatform}
                      onChange={(e) => setFormPlatform(e.target.value)}
                      className="w-full rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-sm text-gray-200 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                      placeholder="e.g. windows, linux, macos"
                    />
                  </div>

                  {/* Description */}
                  <div>
                    <label className="mb-1.5 block text-sm font-medium text-gray-300">Description</label>
                    <textarea
                      value={formDescription}
                      onChange={(e) => setFormDescription(e.target.value)}
                      rows={2}
                      className="w-full rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-sm text-gray-200 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                      placeholder="Optional description…"
                    />
                  </div>

                  {/* Attack Procedure */}
                  <div>
                    <label className="mb-1.5 block text-sm font-medium text-gray-300">
                      Attack Procedure
                    </label>
                    <textarea
                      value={formProcedure}
                      onChange={(e) => setFormProcedure(e.target.value)}
                      rows={4}
                      className="w-full rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 font-mono text-sm text-gray-200 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                      placeholder="Steps to execute the attack…"
                    />
                  </div>

                  {/* Tool */}
                  <div>
                    <label className="mb-1.5 block text-sm font-medium text-gray-300">
                      Suggested Tool
                    </label>
                    <input
                      value={formTool}
                      onChange={(e) => setFormTool(e.target.value)}
                      className="w-full rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-sm text-gray-200 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                      placeholder="e.g. Atomic Red Team, Cobalt Strike"
                    />
                  </div>

                  {/* Expected detection — read-only reference */}
                  {fullTemplate?.expected_detection && (
                    <div>
                      <label className="mb-1.5 block text-sm font-medium text-gray-300">
                        Expected Detection
                        <span className="ml-2 text-xs text-gray-500">(reference — Blue Team)</span>
                      </label>
                      <div className="rounded-lg border border-gray-700 bg-gray-800/50 p-3">
                        <p className="whitespace-pre-wrap font-mono text-xs text-gray-400">
                          {fullTemplate.expected_detection}
                        </p>
                      </div>
                    </div>
                  )}

                  {/* Error */}
                  {createAndAddMutation.isError && (
                    <div className="rounded-lg border border-red-500/30 bg-red-900/20 p-3 text-sm text-red-400">
                      {(createAndAddMutation.error as Error)?.message || "Failed to create test"}
                    </div>
                  )}
                </>
              )}
            </div>
          )}

          {/* ═══ EXISTING TESTS TAB ═══════════════════════════════ */}
          {tab === "existing" && (
            <>
              <div className="border-b border-gray-800 px-6 py-3">
                <div className="relative">
                  <Search className="absolute left-3 top-1/2 h-4 w-4 -translate-y-1/2 text-gray-500" />
                  <input
                    type="text"
                    value={existingSearch}
                    onChange={(e) => setExistingSearch(e.target.value)}
                    placeholder="Search tests by name or technique…"
                    autoFocus
                    className="w-full rounded-lg border border-gray-700 bg-gray-800 pl-9 pr-3 py-2.5 text-sm text-gray-300 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
                  />
                </div>
              </div>

              <div className="px-6 py-3">
                {testsLoading ? (
                  <div className="flex items-center justify-center py-12">
                    <Loader2 className="h-6 w-6 animate-spin text-cyan-400" />
                  </div>
                ) : filteredExisting.length === 0 ? (
                  <div className="flex flex-col items-center justify-center py-12 text-gray-400">
                    <FlaskConical className="mb-2 h-8 w-8 text-gray-600" />
                    <p className="text-sm">
                      {existingSearch
                        ? "No tests match your search."
                        : "No available tests to add."}
                    </p>
                  </div>
                ) : (
                  <div className="space-y-1">
                    {filteredExisting.map((test: Test) => (
                      <div
                        key={test.id}
                        className="flex items-center justify-between rounded-lg border border-gray-800 px-4 py-3 hover:border-gray-700 hover:bg-gray-800/50 transition-colors"
                      >
                        <div className="flex-1 min-w-0">
                          <div className="flex items-center gap-2">
                            <span className="text-sm font-medium text-gray-200 truncate">
                              {test.name}
                            </span>
                            <span className={`inline-flex shrink-0 rounded-full border px-2 py-0.5 text-[10px] font-medium ${stateBadge[test.state]}`}>
                              {test.state.replace(/_/g, " ")}
                            </span>
                          </div>
                          <div className="mt-0.5 flex items-center gap-3 text-xs text-gray-500">
                            {test.technique_mitre_id && (
                              <span className="font-mono text-cyan-400/70">{test.technique_mitre_id}</span>
                            )}
                            {test.technique_name && (
                              <span className="truncate">{test.technique_name}</span>
                            )}
                            {test.platform && <span className="capitalize">{test.platform}</span>}
                          </div>
                        </div>
                        <button
                          onClick={() => addExistingMutation.mutate(test.id)}
                          disabled={addExistingMutation.isPending && addExistingMutation.variables === test.id}
                          className="ml-3 flex shrink-0 items-center gap-1.5 rounded-lg border border-cyan-500/30 bg-cyan-500/10 px-3 py-1.5 text-xs font-medium text-cyan-400 hover:bg-cyan-500/20 disabled:opacity-50 transition-colors"
                        >
                          {addExistingMutation.isPending && addExistingMutation.variables === test.id ? (
                            <Loader2 className="h-3.5 w-3.5 animate-spin" />
                          ) : (
                            <Plus className="h-3.5 w-3.5" />
                          )}
                          Add
                        </button>
                      </div>
                    ))}
                  </div>
                )}
              </div>
            </>
          )}
        </div>

        {/* ── Footer ── */}
        <div className="flex items-center justify-between border-t border-gray-800 px-6 py-4">
          <div className="text-xs text-gray-500">
            {addedCount > 0 && (
              <span className="flex items-center gap-1 text-green-400">
                <CheckCircle className="h-3.5 w-3.5" />
                {addedCount} test{addedCount !== 1 ? "s" : ""} added
              </span>
            )}
          </div>

          {/* Template form submit */}
          {tab === "template" && templateStep === "form" ? (
            <div className="flex items-center gap-2">
              <button
                onClick={() => { setTemplateStep("list"); setSelectedTemplateId(null); }}
                className="rounded-lg border border-gray-700 px-4 py-2 text-sm text-gray-400 hover:bg-gray-800 transition-colors"
              >
                Back
              </button>
              <button
                onClick={() => createAndAddMutation.mutate()}
                disabled={!canSubmitForm}
                className="flex items-center gap-1.5 rounded-lg bg-cyan-600 px-4 py-2 text-sm font-medium text-white hover:bg-cyan-500 disabled:opacity-50 transition-colors"
              >
                {createAndAddMutation.isPending ? (
                  <Loader2 className="h-4 w-4 animate-spin" />
                ) : (
                  <Plus className="h-4 w-4" />
                )}
                Create &amp; Add to Campaign
              </button>
            </div>
          ) : (
            <button
              onClick={onClose}
              className="rounded-lg border border-gray-700 bg-gray-800 px-4 py-2 text-sm font-medium text-gray-300 hover:bg-gray-700 transition-colors"
            >
              Done
            </button>
          )}
        </div>
      </div>
    </div>
  );
}