frontend/src/app/screening/page.tsx

"use client";

import { useState, useEffect, useRef } from "react";
import { startWorkflow, streamWorkflow, screenPatient, getGraphPatients } from "@/lib/api";
import { ClipboardCheck, CheckCircle, XCircle, AlertCircle, Loader2, ChevronRight, GitBranch } from "lucide-react";
import { clsx } from "clsx";

const FALLBACK_PATIENTS = ["P001", "P002", "P003", "P004", "P005"];
const QUICK_TRIALS = [
  { id: "NCT04889131", label: "NCT04889131 – HER2+ Breast Cancer" },
  { id: "NCT05123456", label: "NCT05123456 – Immunotherapy Combo" },
  { id: "NCT05456789", label: "NCT05456789 – BRCA2 Prostate" },
  { id: "NCT06112233", label: "NCT06112233 – EGFR NSCLC" },
  { id: "NCT05334455", label: "NCT05334455 – MSI-H Colorectal" },
];

const WORKFLOW_STATES = [
  "PENDING", "INGESTING", "PARSING_PROTOCOL", "MATCHING", "SCORING", "RECRUITING", "COMPLETED",
];

function StateTracker({ events }: { events: any[] }) {
  const stateSet = new Set(events.map((e: any) => e.state));
  return (
    <div className="flex items-center gap-1 flex-wrap mb-6">
      {WORKFLOW_STATES.map((state, i) => {
        const done = stateSet.has(state);
        const isCurrent = events[events.length - 1]?.state === state;
        return (
          <div key={state} className="flex items-center gap-1">
            <span className={clsx(
              "text-xs px-2.5 py-1 rounded-full font-medium",
              isCurrent ? "bg-indigo-600 text-white" :
              done ? "bg-indigo-100 text-indigo-700" :
              "bg-slate-100 text-slate-400"
            )}>
              {state}
            </span>
            {i < WORKFLOW_STATES.length - 1 && <ChevronRight className="w-3 h-3 text-slate-300" />}
          </div>
        );
      })}
    </div>
  );
}

function ScoreBar({ score }: { score: number }) {
  const pct = Math.round(score * 100);
  const color = pct >= 80 ? "bg-emerald-500" : pct >= 60 ? "bg-amber-500" : "bg-red-500";
  return (
    <div className="flex items-center gap-3">
      <div className="flex-1 h-2 bg-slate-200 rounded-full overflow-hidden">
        <div className={clsx("h-full rounded-full transition-all", color)} style={{ width: `${pct}%` }} />
      </div>
      <span className={clsx("text-sm font-bold", pct >= 80 ? "text-emerald-600" : pct >= 60 ? "text-amber-600" : "text-red-600")}>
        {pct}%
      </span>
    </div>
  );
}

function CriterionRow({ criterion, met, triggered, confidence, note }: any) {
  const pass = met === true || triggered === false;
  const fail = met === false || triggered === true;
  return (
    <div className={clsx("flex items-start gap-3 px-4 py-2.5 rounded-lg", pass ? "bg-emerald-50" : fail ? "bg-red-50" : "bg-amber-50")}>
      {pass ? <CheckCircle className="w-4 h-4 text-emerald-500 shrink-0 mt-0.5" />
             : fail ? <XCircle className="w-4 h-4 text-red-500 shrink-0 mt-0.5" />
             : <AlertCircle className="w-4 h-4 text-amber-500 shrink-0 mt-0.5" />}
      <div className="flex-1 min-w-0">
        <div className="text-sm text-slate-700">{criterion}</div>
        {note && <div className="text-xs text-slate-500 mt-0.5">{note}</div>}
      </div>
      {confidence && (
        <span className="text-xs text-slate-400 shrink-0 mt-0.5">{confidence}</span>
      )}
    </div>
  );
}

function MatchPath({ path }: { path: any[] }) {
  if (!path?.length) return null;
  return (
    <div className="mt-3 p-3 bg-slate-50 rounded-lg border border-slate-200">
      <div className="flex items-center gap-1.5 mb-2">
        <GitBranch className="w-3.5 h-3.5 text-indigo-500" />
        <span className="text-xs font-semibold text-slate-600">Graph Match Path</span>
      </div>
      <div className="flex flex-wrap items-center gap-1 text-xs">
        {path.map((node: any, i: number) => (
          <span key={i} className="flex items-center gap-1">
            <span className="bg-indigo-100 text-indigo-700 px-1.5 py-0.5 rounded font-mono">
              {node.from?.split(":")[1] ?? node.from}
            </span>
            <span className="text-slate-400">—[{node.rel}]→</span>
            <span className="bg-emerald-100 text-emerald-700 px-1.5 py-0.5 rounded font-mono">
              {node.to?.split(":")[1] ?? node.to}
            </span>
            {node.note && <span className="text-slate-400 italic">({node.note})</span>}
          </span>
        ))}
      </div>
    </div>
  );
}

export default function ScreeningPage() {
  const [patientId, setPatientId] = useState("P001");
  const [nctId, setNctId] = useState(QUICK_TRIALS[0].id);
  const [customNct, setCustomNct] = useState("");
  const [useCustom, setUseCustom] = useState(false);
  const [loading, setLoading] = useState(false);
  const [mode, setMode] = useState<"screen" | "workflow">("screen");
  const [result, setResult] = useState<any>(null);
  const [workflowResult, setWorkflowResult] = useState<any>(null);
  const [streamEvents, setStreamEvents] = useState<any[]>([]);
  const [streamDone, setStreamDone] = useState(false);
  const [error, setError] = useState("");
  const [graphPatients, setGraphPatients] = useState<any[]>([]);
  const cleanupRef = useRef<(() => void) | null>(null);

  useEffect(() => {
    getGraphPatients(undefined, 500).then((d) => setGraphPatients(d.patients)).catch(() => {});
    return () => { cleanupRef.current?.(); };
  }, []);

  const effectiveNct = useCustom ? customNct : nctId;

  const handleScreen = async () => {
    if (!patientId.trim()) { setError("Patient ID is required"); return; }
    if (!effectiveNct.trim()) { setError("NCT ID is required"); return; }
    setLoading(true);
    setError("");
    setResult(null);
    try {
      const data = await screenPatient(patientId.trim(), effectiveNct.trim());
      setResult(data);
    } catch (e: any) {
      setError(e.message);
    }
    setLoading(false);
  };

  const handleWorkflow = async () => {
    if (!patientId.trim()) { setError("Patient ID is required"); return; }
    cleanupRef.current?.();
    setLoading(true);
    setError("");
    setWorkflowResult(null);
    setStreamEvents([]);
    setStreamDone(false);
    try {
      const { workflow_id } = await startWorkflow(patientId.trim());
      const stop = streamWorkflow(
        workflow_id,
        (evt) => setStreamEvents((prev) => [...prev, evt]),
        () => { setStreamDone(true); setLoading(false); },
      );
      cleanupRef.current = stop;
    } catch (e: any) {
      setError(e.message);
      setLoading(false);
    }
  };

  return (
    <div className="p-6 max-w-4xl mx-auto">
      <div className="mb-6">
        <h1 className="text-2xl font-bold text-slate-900 mb-1">Patient Screening</h1>
        <p className="text-slate-500 text-sm">AI-powered eligibility assessment using FHIR R4 patient data and A2A orchestration</p>
      </div>

      {/* Mode tabs */}
      <div className="flex gap-2 mb-6">
        {(["screen", "workflow"] as const).map((m) => (
          <button
            key={m}
            onClick={() => setMode(m)}
            className={clsx("px-4 py-2 rounded-lg text-sm font-medium transition-colors",
              mode === m ? "bg-indigo-600 text-white" : "bg-white border border-slate-200 text-slate-600 hover:bg-slate-50"
            )}
          >
            {m === "screen" ? "Single Trial Screen" : "A2A Full Pipeline"}
          </button>
        ))}
      </div>

      <div className="bg-white rounded-xl border border-slate-200 p-5 mb-6">
        <div className="grid grid-cols-2 gap-4 mb-4">
          <div>
            <label className="block text-xs font-semibold text-slate-600 mb-1.5">
              Patient ID
              {graphPatients.length > 0 && <span className="ml-1 text-slate-400 font-normal">({graphPatients.length} loaded)</span>}
            </label>
            <input
              list="patient-list"
              value={patientId}
              onChange={(e) => setPatientId(e.target.value)}
              placeholder="e.g. P_C50_0001 or P001"
              className="w-full border border-slate-200 rounded-lg px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-indigo-500"
            />
            <datalist id="patient-list">
              {(graphPatients.length > 0 ? graphPatients : FALLBACK_PATIENTS.map((id) => ({ id }))).map((p: any) => (
                <option key={p.id} value={p.id}>
                  {p.name ? `${p.id} — ${p.name}${p.condition ? ` (${p.condition})` : ""}` : p.id}
                </option>
              ))}
            </datalist>
          </div>
          {mode === "screen" && (
            <div>
              <label className="block text-xs font-semibold text-slate-600 mb-1.5">NCT ID</label>
              <input
                list="trial-list"
                value={useCustom ? customNct : nctId}
                onChange={(e) => {
                  const val = e.target.value;
                  const quick = QUICK_TRIALS.find((t) => t.id === val);
                  if (quick) { setUseCustom(false); setNctId(val); }
                  else { setUseCustom(true); setCustomNct(val); }
                }}
                placeholder="e.g. NCT04889131"
                className="w-full border border-slate-200 rounded-lg px-3 py-2 text-sm focus:outline-none focus:ring-2 focus:ring-indigo-500"
              />
              <datalist id="trial-list">
                {QUICK_TRIALS.map((t) => <option key={t.id} value={t.id}>{t.label}</option>)}
              </datalist>
            </div>
          )}
        </div>

        <button
          onClick={mode === "screen" ? handleScreen : handleWorkflow}
          disabled={loading}
          className="flex items-center gap-2 bg-indigo-600 text-white px-5 py-2.5 rounded-lg text-sm font-medium hover:bg-indigo-700 disabled:opacity-50 transition-colors"
        >
          {loading ? <><Loader2 className="w-4 h-4 animate-spin" />Running...</> : <><ClipboardCheck className="w-4 h-4" />{mode === "screen" ? "Screen Patient" : "Run A2A Pipeline"}</>}
        </button>
      </div>

      {error && (
        <div className="bg-red-50 border border-red-200 text-red-700 rounded-lg px-4 py-3 text-sm mb-4">{error}</div>
      )}

      {/* Single screen result */}
      {result && (
        <div className="space-y-4">
          <div className="bg-white rounded-xl border border-slate-200 p-5">
            <div className="flex items-center justify-between mb-3">
              <h2 className="font-semibold text-slate-900">Eligibility Score</h2>
              <span className={clsx("text-sm font-bold px-3 py-1 rounded-full",
                result.eligible ? "bg-emerald-100 text-emerald-700" : "bg-red-100 text-red-700"
              )}>
                {result.eligible ? "ELIGIBLE" : "NOT ELIGIBLE"}
              </span>
            </div>
            <ScoreBar score={result.overall_score || 0} />
            {result.summary && (
              <p className="text-sm text-slate-600 mt-3 leading-relaxed">{result.summary}</p>
            )}
            {result.risk_flags?.length > 0 && (
              <div className="mt-3 flex flex-wrap gap-2">
                {result.risk_flags.map((f: string, i: number) => (
                  <span key={i} className="text-xs bg-amber-50 border border-amber-200 text-amber-700 px-2 py-0.5 rounded">{f}</span>
                ))}
              </div>
            )}
          </div>

          {result.inclusion_results?.length > 0 && (
            <div className="bg-white rounded-xl border border-slate-200 p-5">
              <h3 className="font-semibold text-slate-900 text-sm mb-3">Inclusion Criteria</h3>
              <div className="space-y-2">
                {result.inclusion_results.map((c: any, i: number) => <CriterionRow key={i} {...c} />)}
              </div>
            </div>
          )}

          {result.exclusion_results?.length > 0 && (
            <div className="bg-white rounded-xl border border-slate-200 p-5">
              <h3 className="font-semibold text-slate-900 text-sm mb-3">Exclusion Criteria</h3>
              <div className="space-y-2">
                {result.exclusion_results.map((c: any, i: number) => <CriterionRow key={i} {...c} />)}
              </div>
            </div>
          )}

          {result.match_path?.length > 0 && (
            <div className="bg-white rounded-xl border border-slate-200 p-5">
              <h3 className="font-semibold text-slate-900 text-sm mb-2">Graph Explainability</h3>
              <MatchPath path={result.match_path} />
            </div>
          )}
        </div>
      )}

      {/* Streaming A2A workflow */}
      {(streamEvents.length > 0 || loading) && mode === "workflow" && (
        <div className="space-y-4">
          <div className="bg-white rounded-xl border border-slate-200 p-5">
            <div className="flex items-center justify-between mb-3">
              <h2 className="font-semibold text-slate-900">A2A Pipeline — Live</h2>
              {!streamDone && <Loader2 className="w-4 h-4 animate-spin text-indigo-500" />}
              {streamDone && <span className="text-xs text-emerald-600 font-medium">Complete</span>}
            </div>
            <StateTracker events={streamEvents} />
            <div className="space-y-1.5 mt-2 max-h-48 overflow-y-auto">
              {streamEvents.map((evt: any, i: number) => (
                <div key={i} className="flex gap-3 text-xs items-start">
                  <span className="text-indigo-600 font-mono shrink-0 pt-0.5">{evt.state}</span>
                  <span className="text-slate-600">{evt.message}</span>
                  {evt.data?.eligible_count !== undefined && (
                    <span className="ml-auto text-emerald-600 font-medium shrink-0">
                      {evt.data.eligible_count} eligible
                    </span>
                  )}
                </div>
              ))}
            </div>
          </div>

          {/* Final summary when complete */}
          {streamDone && (() => {
            const final = streamEvents[streamEvents.length - 1];
            if (!final || !["COMPLETED", "FAILED"].includes(final.state)) return null;
            return (
              <div className={clsx("bg-white rounded-xl border p-5",
                final.state === "COMPLETED" ? "border-emerald-200" : "border-red-200"
              )}>
                <h3 className="font-semibold text-slate-900 text-sm mb-2">
                  {final.state === "COMPLETED" ? "Pipeline Complete" : "Pipeline Failed"}
                </h3>
                {final.state === "COMPLETED" && (
                  <div className="grid grid-cols-3 gap-4 text-center">
                    <div className="bg-emerald-50 rounded-lg p-3">
                      <div className="text-2xl font-bold text-emerald-600">{final.eligible_trials ?? 0}</div>
                      <div className="text-xs text-slate-500 mt-0.5">Eligible Trials</div>
                    </div>
                    <div className="bg-indigo-50 rounded-lg p-3">
                      <div className="text-2xl font-bold text-indigo-600">{final.total_evaluated ?? 0}</div>
                      <div className="text-xs text-slate-500 mt-0.5">Evaluated</div>
                    </div>
                    <div className="bg-amber-50 rounded-lg p-3">
                      <div className="text-2xl font-bold text-amber-600">{final.recruitment_records ?? 0}</div>
                      <div className="text-xs text-slate-500 mt-0.5">Outreach Generated</div>
                    </div>
                  </div>
                )}
                {final.error && (
                  <p className="text-sm text-red-600 mt-2">{final.error}</p>
                )}
              </div>
            );
          })()}
        </div>
      )}
    </div>
  );
}