proof-of-work/proof-of-work/components/ai/GenerateClient.tsx

'use client';

import { useEffect, useMemo, useRef, useState } from 'react';
import { useRouter } from 'next/navigation';
import { Loader2, Sparkles, Square } from 'lucide-react';
import { lenientJsonParse } from '@/lib/ai/lenientJson';

const DAY_LABELS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat'];

interface Template {
  id: string;
  name: string;
  description: string | null;
  isBuiltIn: boolean;
}
interface LibraryExercise {
  id: string;
  name: string;
  type: string;
}

// AI output shape (matches lib/ai/programSchema.ts)
interface AIExercise {
  exerciseId: string | null;
  exerciseName: string;
  order: number;
  sets?: number | null;
  repsMin?: number | null;
  repsMax?: number | null;
  rpe?: number | null;
  restSeconds?: number | null;
  notes?: string | null;
}
interface AIDay {
  dayOfWeek: number;
  name?: string | null;
  description?: string | null;
  exercises: AIExercise[];
}
interface AIWeek {
  weekNumber: number;
  phase?: string | null;
  description?: string | null;
  days: AIDay[];
}
interface AIProgram {
  name: string;
  description?: string | null;
  type: string;
  durationWeeks: number;
  weeks: AIWeek[];
}

type Phase =
  | { kind: 'idle' }
  | { kind: 'streaming'; raw: string; partial: Partial<AIProgram> | null }
  | { kind: 'parsed'; raw: string; program: AIProgram }
  | { kind: 'failed'; raw: string; message: string };

export default function GenerateClient({
  templates,
  exercises,
  providerLabel,
  modelLabel,
  workoutCount,
}: {
  templates: Template[];
  exercises: LibraryExercise[];
  providerLabel: string;
  modelLabel: string;
  workoutCount: number;
}) {
  const router = useRouter();
  const [templateId, setTemplateId] = useState(templates[0]?.id ?? '');
  const [userInput, setUserInput] = useState('');
  const [includeHistory, setIncludeHistory] = useState(workoutCount >= 10);
  const [generationId, setGenerationId] = useState<string | null>(null);
  const [phase, setPhase] = useState<Phase>({ kind: 'idle' });
  const [tokens, setTokens] = useState<{ in?: number; out?: number }>({});
  const abortRef = useRef<AbortController | null>(null);

  const selectedTemplate = useMemo(
    () => templates.find((t) => t.id === templateId),
    [templates, templateId],
  );

  const handleGenerate = async () => {
    if (!userInput.trim()) return;
    setPhase({ kind: 'streaming', raw: '', partial: null });
    setGenerationId(null);
    setTokens({});

    abortRef.current = new AbortController();
    let raw = '';
    try {
      const res = await fetch('/api/ai/generate', {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          templateId: templateId || null,
          userInput,
          includeHistory,
        }),
        signal: abortRef.current.signal,
      });
      if (!res.ok) {
        const body = await res.json().catch(() => ({}));
        setPhase({
          kind: 'failed',
          raw: '',
          message: body.error ?? `HTTP ${res.status}`,
        });
        return;
      }
      if (!res.body) {
        setPhase({ kind: 'failed', raw: '', message: 'No response body.' });
        return;
      }
      // Parse SSE stream
      const reader = res.body.getReader();
      const decoder = new TextDecoder();
      let buf = '';
      let done = false;
      while (!done) {
        const { value, done: d } = await reader.read();
        if (d) {
          done = true;
          break;
        }
        buf += decoder.decode(value, { stream: true });
        let idx;
        while ((idx = buf.indexOf('\n\n')) >= 0) {
          const event = buf.slice(0, idx);
          buf = buf.slice(idx + 2);
          let evtName = 'message';
          const dataLines: string[] = [];
          for (const line of event.split('\n')) {
            if (line.startsWith('event:')) evtName = line.slice(6).trim();
            else if (line.startsWith('data:'))
              dataLines.push(line.slice(5).trimStart());
          }
          if (!dataLines.length) continue;
          const data = dataLines.join('\n');
          let parsed: any;
          try {
            parsed = JSON.parse(data);
          } catch {
            continue;
          }
          if (evtName === 'generation') {
            setGenerationId(parsed.id);
          } else if (evtName === 'text') {
            raw += parsed.delta;
            const partial = lenientJsonParse(raw) as Partial<AIProgram> | null;
            setPhase({ kind: 'streaming', raw, partial });
          } else if (evtName === 'usage') {
            setTokens({ in: parsed.tokensIn, out: parsed.tokensOut });
          } else if (evtName === 'complete') {
            // Server already validated/stored the parsed program. We
            // fetch the generation record AFTER the stream closes
            // (below) to get the parsed JSON. Just record the
            // success/failure outcome here; if it failed, render
            // the error inline now since we're not going to fetch.
            if (!parsed.parsedOk) {
              setPhase({
                kind: 'failed',
                raw,
                message: parsed.errorMessage ?? 'Failed to parse model output.',
              });
            }
          }
        }
      }
    } catch (e) {
      if ((e as Error).name === 'AbortError') {
        setPhase({ kind: 'failed', raw, message: 'Cancelled.' });
      } else {
        setPhase({
          kind: 'failed',
          raw,
          message: (e as Error).message,
        });
      }
      return;
    }

    // After stream closes, fetch the generation row to get the parsed
    // program (we don't try to re-parse client-side — server already did).
    const id = generationIdRef.current;
    if (id) {
      const r = await fetch(`/api/ai/generations/${id}`);
      if (r.ok) {
        const gen = await r.json();
        if (gen.status === 'completed' && gen.parsedProgram) {
          setPhase({
            kind: 'parsed',
            raw,
            program: JSON.parse(gen.parsedProgram) as AIProgram,
          });
          return;
        }
        if (gen.status === 'failed') {
          setPhase({
            kind: 'failed',
            raw,
            message: gen.errorMessage ?? 'Failed.',
          });
          return;
        }
      }
    }
  };

  // Capture the generationId in a ref so the async fetch after the
  // stream has access to it (the closure above sees the initial null).
  const generationIdRef = useRef<string | null>(null);
  useEffect(() => {
    generationIdRef.current = generationId;
  }, [generationId]);

  const handleCancel = () => {
    abortRef.current?.abort();
  };

  return (
    <div className="space-y-6">
      <div className="text-xs text-zinc-500 uppercase tracking-wider">
        Provider: <span className="text-zinc-300">{providerLabel}</span>
        {' · '}Model: <span className="text-zinc-300">{modelLabel}</span>
      </div>

      <section className="bg-zinc-900 border border-zinc-800 rounded p-4 space-y-4">
        <Field label="Template">
          <select
            value={templateId}
            onChange={(e) => setTemplateId(e.target.value)}
            className={inputClass}
            disabled={phase.kind === 'streaming'}
          >
            {templates.map((t) => (
              <option key={t.id} value={t.id}>
                {t.isBuiltIn ? '★ ' : ''}
                {t.name}
              </option>
            ))}
          </select>
          {selectedTemplate?.description && (
            <p className="text-xs text-zinc-500 mt-1">
              {selectedTemplate.description}
            </p>
          )}
        </Field>

        <Field label="Your specifics">
          <textarea
            value={userInput}
            onChange={(e) => setUserInput(e.target.value)}
            placeholder="e.g. 8 weeks, 4 days per week, heavy leg emphasis. I have a meet in 6 weeks. Bench Press is at 245x5, Squat 365x3, Deadlift 425x3."
            rows={6}
            className={inputClass}
            disabled={phase.kind === 'streaming'}
          />
        </Field>

        <label className="flex items-start gap-2 text-xs text-zinc-300">
          <input
            type="checkbox"
            checked={includeHistory}
            onChange={(e) => setIncludeHistory(e.target.checked)}
            disabled={phase.kind === 'streaming' || workoutCount === 0}
            className="mt-0.5"
          />
          <span>
            Include my workout history as context{' '}
            <span className="text-zinc-500">
              ({workoutCount === 0
                ? 'no workouts logged yet — disabled'
                : `last 90 days · summarizes per-exercise frequency, recent weights, stagnations`}
              )
            </span>
          </span>
        </label>

        <div className="flex items-center gap-2">
          {phase.kind === 'streaming' ? (
            <button
              type="button"
              onClick={handleCancel}
              className="inline-flex items-center gap-2 px-4 py-2 rounded border border-red-900 text-red-400 text-xs uppercase tracking-wider hover:bg-red-900/30"
            >
              <Square className="w-3.5 h-3.5" />
              Cancel
            </button>
          ) : (
            <button
              type="button"
              onClick={handleGenerate}
              disabled={!userInput.trim()}
              className="inline-flex items-center gap-2 px-5 py-2 rounded bg-white text-black font-bold text-xs uppercase tracking-wider hover:bg-gray-100 disabled:bg-zinc-700 disabled:text-zinc-500"
            >
              <Sparkles className="w-4 h-4" />
              Generate
            </button>
          )}
        </div>
      </section>

      {(phase.kind === 'streaming' || phase.kind === 'failed' || phase.kind === 'parsed') && (
        <section className="space-y-3">
          <div className="flex items-center justify-between">
            <h2 className="text-sm font-semibold text-white uppercase tracking-wider">
              {phase.kind === 'streaming' ? 'Generating...' : 'Response'}
            </h2>
            {(tokens.in != null || tokens.out != null) && (
              <span className="text-[11px] text-zinc-500 uppercase tracking-wider">
                {tokens.in ?? '?'} in · {tokens.out ?? '?'} out
              </span>
            )}
          </div>

          {phase.kind === 'streaming' && (
            <>
              {phase.partial ? (
                <PartialPreview partial={phase.partial} />
              ) : (
                <div className="text-xs text-zinc-500 italic">
                  Waiting for the first parseable JSON...
                  <Loader2 className="inline w-3 h-3 animate-spin ml-2" />
                </div>
              )}
              <details className="text-xs text-zinc-500">
                <summary className="cursor-pointer">Raw stream</summary>
                <div className="bg-zinc-950 border border-zinc-800 rounded p-3 font-mono text-[11px] text-zinc-400 max-h-80 overflow-auto whitespace-pre-wrap mt-2">
                  {phase.raw || '(waiting for first token...)'}
                  <Loader2 className="inline w-3 h-3 animate-spin ml-2" />
                </div>
              </details>
            </>
          )}

          {phase.kind === 'failed' && (
            <>
              <div className="bg-red-950/40 border border-red-900 rounded p-3 text-sm text-red-300">
                {phase.message}
              </div>
              {phase.raw && (
                <details className="text-xs text-zinc-500">
                  <summary className="cursor-pointer">Raw response</summary>
                  <pre className="bg-zinc-950 border border-zinc-800 rounded p-3 mt-2 whitespace-pre-wrap">
                    {phase.raw}
                  </pre>
                </details>
              )}
            </>
          )}

          {phase.kind === 'parsed' && generationId && (
            <ProgramPreview
              program={phase.program}
              generationId={generationId}
              exercises={exercises}
              onApplied={(programId) => {
                router.push(`/main/programs/${programId}`);
              }}
            />
          )}
        </section>
      )}
    </div>
  );
}

function ProgramPreview({
  program: initial,
  generationId,
  exercises,
  onApplied,
}: {
  program: AIProgram;
  generationId: string;
  exercises: LibraryExercise[];
  onApplied: (programId: string) => void;
}) {
  const [program, setProgram] = useState<AIProgram>(initial);
  const [applying, setApplying] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [startDate, setStartDate] = useState(
    new Date().toISOString().slice(0, 10),
  );
  const [activate, setActivate] = useState(true);

  const exerciseLookup = useMemo(
    () => new Map(exercises.map((e) => [e.id, e])),
    [exercises],
  );

  const unresolvedCount = useMemo(() => {
    let n = 0;
    for (const w of program.weeks)
      for (const d of w.days)
        for (const ex of d.exercises) if (!ex.exerciseId) n++;
    return n;
  }, [program]);

  const setExerciseId = (
    weekIdx: number,
    dayIdx: number,
    exIdx: number,
    newId: string | null,
  ) => {
    setProgram((p) => {
      const next = structuredClone(p);
      next.weeks[weekIdx].days[dayIdx].exercises[exIdx].exerciseId = newId;
      return next;
    });
  };

  const removeExercise = (weekIdx: number, dayIdx: number, exIdx: number) => {
    setProgram((p) => {
      const next = structuredClone(p);
      next.weeks[weekIdx].days[dayIdx].exercises.splice(exIdx, 1);
      // Renumber order
      next.weeks[weekIdx].days[dayIdx].exercises.forEach(
        (ex: AIExercise, i: number) => {
          ex.order = i;
        },
      );
      return next;
    });
  };

  const handleApply = async () => {
    if (unresolvedCount > 0) {
      setError(
        `Resolve all ${unresolvedCount} unknown exercise(s) before applying.`,
      );
      return;
    }
    setError(null);
    setApplying(true);
    try {
      const res = await fetch('/api/ai/apply', {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          generationId,
          program,
          startDate,
          isActive: activate,
        }),
      });
      const body = await res.json();
      if (!res.ok) throw new Error(body.error ?? `HTTP ${res.status}`);
      onApplied(body.programId);
    } catch (e) {
      setError((e as Error).message);
    } finally {
      setApplying(false);
    }
  };

  return (
    <div className="bg-zinc-900 border border-zinc-800 rounded p-4 space-y-4">
      <div>
        <h3 className="text-lg font-bold text-white">{program.name}</h3>
        <p className="text-xs text-zinc-500 mt-1">
          {program.type} · {program.durationWeeks} week
          {program.durationWeeks === 1 ? '' : 's'} · {program.weeks.length}{' '}
          week{program.weeks.length === 1 ? '' : 's'} planned
        </p>
        {program.description && (
          <p className="text-sm text-zinc-300 mt-2">{program.description}</p>
        )}
      </div>

      {unresolvedCount > 0 && (
        <div className="rounded bg-amber-950/30 border border-amber-900 px-3 py-2 text-xs text-amber-200">
          {unresolvedCount} exercise(s) the AI couldn&apos;t map to your
          library. Pick a replacement or remove them before applying.
        </div>
      )}

      <div className="space-y-3">
        {program.weeks.map((w, wIdx) => (
          <details
            key={w.weekNumber}
            open={wIdx === 0}
            className="bg-zinc-950 border border-zinc-800 rounded"
          >
            <summary className="cursor-pointer px-3 py-2 text-sm text-white">
              Week {w.weekNumber}
              {w.phase && (
                <span className="text-zinc-500"> · {w.phase}</span>
              )}
              <span className="text-zinc-600 text-xs">
                {' '}
                ({w.days.length} day{w.days.length === 1 ? '' : 's'})
              </span>
            </summary>
            <div className="p-3 space-y-2">
              {w.days.map((d, dIdx) => (
                <div
                  key={d.dayOfWeek}
                  className="bg-zinc-900 border border-zinc-800 rounded p-3"
                >
                  <p className="text-xs font-semibold text-zinc-300 uppercase tracking-wider">
                    {DAY_LABELS[d.dayOfWeek]}
                    {d.name && (
                      <span className="text-zinc-500 normal-case font-normal">
                        {' '}
                        · {d.name}
                      </span>
                    )}
                  </p>
                  <ul className="mt-2 space-y-2">
                    {d.exercises.map((ex, eIdx) => {
                      const isUnknown = !ex.exerciseId;
                      const lib = ex.exerciseId
                        ? exerciseLookup.get(ex.exerciseId)
                        : null;
                      return (
                        <li
                          key={eIdx}
                          className={`text-sm ${isUnknown ? 'bg-amber-950/30 border border-amber-900' : 'bg-zinc-950 border border-zinc-800'} rounded p-2`}
                        >
                          <div className="flex items-start justify-between gap-2">
                            <div className="min-w-0 flex-1">
                              <div className="text-white">
                                {lib?.name ?? ex.exerciseName}
                                {isUnknown && (
                                  <span className="ml-2 text-[10px] uppercase tracking-wider text-amber-400">
                                    not in library
                                  </span>
                                )}
                              </div>
                              {(ex.sets || ex.repsMin || ex.repsMax || ex.rpe || ex.restSeconds) && (
                                <div className="text-xs text-zinc-500 mt-0.5">
                                  {ex.sets ? `${ex.sets}×` : ''}
                                  {ex.repsMin === ex.repsMax || !ex.repsMax
                                    ? (ex.repsMin ?? '?')
                                    : `${ex.repsMin}-${ex.repsMax}`}
                                  {ex.rpe ? ` @ RPE ${ex.rpe}` : ''}
                                  {ex.restSeconds ? ` · rest ${ex.restSeconds}s` : ''}
                                </div>
                              )}
                              {ex.notes && (
                                <div className="text-xs text-zinc-400 mt-1 italic">
                                  {ex.notes}
                                </div>
                              )}
                            </div>
                            <button
                              type="button"
                              onClick={() => removeExercise(wIdx, dIdx, eIdx)}
                              className="text-xs text-red-400 hover:text-red-300 px-1"
                              title="Remove from program"
                            >
                              ✕
                            </button>
                          </div>
                          {isUnknown && (
                            <div className="mt-2">
                              <select
                                value=""
                                onChange={(e) =>
                                  setExerciseId(wIdx, dIdx, eIdx, e.target.value || null)
                                }
                                className="w-full text-xs px-2 py-1 rounded border border-amber-900 bg-zinc-900 text-white"
                              >
                                <option value="">
                                  Map to existing exercise...
                                </option>
                                {exercises.map((opt) => (
                                  <option key={opt.id} value={opt.id}>
                                    {opt.name} ({opt.type})
                                  </option>
                                ))}
                              </select>
                            </div>
                          )}
                        </li>
                      );
                    })}
                  </ul>
                </div>
              ))}
            </div>
          </details>
        ))}
      </div>

      <div className="border-t border-zinc-800 pt-4 space-y-3">
        <div className="grid grid-cols-2 gap-3">
          <Field label="Start date">
            <input
              type="date"
              value={startDate}
              onChange={(e) => setStartDate(e.target.value)}
              className={inputClass}
            />
          </Field>
          <label className="flex items-end gap-2">
            <input
              type="checkbox"
              checked={activate}
              onChange={(e) => setActivate(e.target.checked)}
              className="mb-2"
            />
            <span className="text-xs text-zinc-300 mb-2">
              Activate this program after applying
            </span>
          </label>
        </div>

        {error && (
          <div className="rounded bg-red-900/50 px-3 py-2 border border-red-800 text-xs text-red-400">
            {error}
          </div>
        )}

        <button
          type="button"
          onClick={handleApply}
          disabled={applying || unresolvedCount > 0}
          className="px-5 py-2 rounded bg-emerald-700 text-white font-bold text-xs uppercase tracking-wider hover:bg-emerald-600 disabled:bg-zinc-700 disabled:text-zinc-500"
        >
          {applying ? (
            <>
              <Loader2 className="inline w-4 h-4 animate-spin mr-2" />
              Applying...
            </>
          ) : (
            'Apply this program'
          )}
        </button>
      </div>
    </div>
  );
}

const inputClass =
  'w-full px-3 py-2 text-sm rounded border border-zinc-700 bg-zinc-800 text-white placeholder:text-zinc-500 focus:outline-none focus:ring-2 focus:ring-white/30';

function Field({ label, children }: { label: string; children: React.ReactNode }) {
  return (
    <label className="block">
      <span className="text-[11px] font-semibold text-zinc-400 uppercase tracking-wider block mb-1">
        {label}
      </span>
      {children}
    </label>
  );
}

function PartialPreview({ partial }: { partial: Partial<AIProgram> }) {
  const weeks = (partial.weeks as AIWeek[] | undefined) ?? [];
  return (
    <div className="bg-zinc-950 border border-zinc-800 rounded p-3 space-y-2">
      <div className="flex items-center gap-2 text-xs">
        <Loader2 className="w-3 h-3 animate-spin text-zinc-500" />
        <span className="text-zinc-400">
          Building program...{' '}
          {partial.name && (
            <span className="text-white font-semibold">{partial.name}</span>
          )}
          {partial.type && (
            <span className="text-zinc-500"> · {partial.type}</span>
          )}
          {typeof partial.durationWeeks === 'number' && (
            <span className="text-zinc-500"> · {partial.durationWeeks} wk</span>
          )}
        </span>
      </div>
      {weeks.length > 0 && (
        <ul className="text-xs text-zinc-300 space-y-1">
          {weeks.map((w, i) => (
            <li key={i}>
              <span className="text-zinc-500">Week {w?.weekNumber ?? '?'}:</span>{' '}
              {Array.isArray(w?.days)
                ? `${w.days.length} day${w.days.length === 1 ? '' : 's'} (${
                    w.days.reduce(
                      (n: number, d: AIDay) =>
                        n + (Array.isArray(d?.exercises) ? d.exercises.length : 0),
                      0,
                    )
                  } exercises)`
                : '...'}
              {w?.phase && (
                <span className="text-zinc-500"> · {w.phase}</span>
              )}
            </li>
          ))}
        </ul>
      )}
    </div>
  );
}