'use client';

import { useEffect, useMemo, useRef, useState } from 'react';
import { useRouter } from 'next/navigation';
import { Loader2, Sparkles } from 'lucide-react';
import { lenientJsonParse } from '@/lib/ai/lenientJson';
import { estimateCost, formatCost } from '@/lib/ai/pricing';
import type { AiWorkoutDraft } from '@/lib/ai/workoutDraft';

interface LibraryExercise {
  id: string;
  name: string;
  type: string;
}

// AI output shape — mirrors lib/ai/workoutSchema.ts (AIWorkout).
interface AIWorkoutExercise {
  exerciseId: string | null;
  exerciseName: string;
  order: number;
  sets?: number | null;
  reps?: number | null;
  suggestedWeight?: number | null;
  suggestedWeightUnit?: 'lbs' | 'kg' | null;
  rpe?: number | null;
  gear?: number | null;
  durationSeconds?: number | null;
  notes?: string | null;
}
interface AIWorkout {
  name: string;
  notes?: string | null;
  exercises: AIWorkoutExercise[];
}

// The ephemeral draft we hand to the New Workout form via sessionStorage.
export const AI_WORKOUT_DRAFT_KEY = 'ai-workout-draft';

type Phase =
  | { kind: 'idle' }
  | { kind: 'streaming'; raw: string; lastPartial: Partial<AIWorkout> | null }
  | { kind: 'failed'; raw: string; message: string };

export default function GenerateWorkoutClient({
  exercises,
  providerLabel,
  modelLabel,
  workoutCount,
}: {
  exercises: LibraryExercise[];
  providerLabel: string;
  modelLabel: string;
  workoutCount: number;
}) {
  const router = useRouter();
  const [userInput, setUserInput] = useState('');
  const [includeHistory, setIncludeHistory] = useState(workoutCount >= 1);
  const [phase, setPhase] = useState<Phase>({ kind: 'idle' });
  // The editable suggestion once parsed. Lifted to the parent so the
  // Refine action can send the user's current edits back as the prior
  // workout. null until the first successful parse.
  const [workout, setWorkout] = useState<AIWorkout | null>(null);
  // Refine instruction lives here (not in WorkoutPreview) because the
  // preview unmounts while streaming — keeping it in the parent means a
  // failed refine doesn't lose what the user typed; we clear it only on
  // a successful regeneration.
  const [refineInput, setRefineInput] = useState('');
  const [tokens, setTokens] = useState<{ in?: number; out?: number; durationMs?: number }>({});
  const closeStreamRef = useRef<(() => void) | null>(null);

  const streaming = phase.kind === 'streaming';

  /**
   * Run a generation. `priorWorkout` present → REVISION mode: `input`
   * is the change instruction and the model re-emits the full workout.
   */
  const runGeneration = async (input: string, priorWorkout?: AIWorkout) => {
    if (!input.trim()) return;
    setPhase({ kind: 'streaming', raw: '', lastPartial: null });
    setTokens({});

    let id: string;
    try {
      const res = await fetch('/api/ai/generate-workout', {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          userInput: input,
          includeHistory,
          priorWorkout: priorWorkout ?? null,
        }),
      });
      const body = await res.json().catch(() => ({}));
      if (!res.ok) {
        setPhase({ kind: 'failed', raw: '', message: body.error ?? `HTTP ${res.status}` });
        return;
      }
      id = body.id;
    } catch (e) {
      setPhase({ kind: 'failed', raw: '', message: (e as Error).message });
      return;
    }

    attachStream(id);
  };

  const attachStream = (id: string) => {
    const es = new EventSource(`/api/ai/generations/${id}/stream`);
    closeStreamRef.current = () => es.close();
    let raw = '';
    let lastPartial: Partial<AIWorkout> | null = null;

    es.addEventListener('text', (ev) => {
      const data = JSON.parse((ev as MessageEvent).data);
      raw += data.delta;
      const next = lenientJsonParse(raw) as Partial<AIWorkout> | null;
      if (next) lastPartial = next; // sticky — kills flicker between parses
      setPhase({ kind: 'streaming', raw, lastPartial });
    });
    es.addEventListener('usage', (ev) => {
      const data = JSON.parse((ev as MessageEvent).data);
      setTokens((t) => ({ ...t, in: data.tokensIn, out: data.tokensOut }));
    });
    es.addEventListener('complete', async (ev) => {
      const data = JSON.parse((ev as MessageEvent).data);
      es.close();
      closeStreamRef.current = null;
      setTokens((t) => ({
        ...t,
        in: data.tokensIn ?? t.in,
        out: data.tokensOut ?? t.out,
        durationMs: data.durationMs,
      }));
      if (data.parsedOk) {
        const r = await fetch(`/api/ai/generations/${id}`);
        if (r.ok) {
          const gen = await r.json();
          if (gen.parsedProgram) {
            setWorkout(JSON.parse(gen.parsedProgram) as AIWorkout);
            setRefineInput(''); // consumed — clear only on success
            setPhase({ kind: 'idle' });
            return;
          }
        }
      }
      setPhase({
        kind: 'failed',
        raw,
        message: data.errorMessage ?? 'Failed to parse model output.',
      });
    });
    es.onerror = () => {
      if (es.readyState === EventSource.CLOSED) {
        closeStreamRef.current = null;
        setPhase((p) =>
          p.kind === 'streaming'
            ? {
                kind: 'failed',
                raw: p.raw,
                message:
                  'Stream disconnected. The generation may still be running — check AI · History.',
              }
            : p,
        );
      }
    };
  };

  // Warn before unload while streaming (the runner keeps going server-side).
  useEffect(() => {
    if (!streaming) return;
    const onBeforeUnload = (e: BeforeUnloadEvent) => {
      e.preventDefault();
      e.returnValue = '';
    };
    window.addEventListener('beforeunload', onBeforeUnload);
    return () => window.removeEventListener('beforeunload', onBeforeUnload);
  }, [streaming]);

  // Detach on unmount; the server keeps generating regardless.
  useEffect(() => () => closeStreamRef.current?.(), []);

  const costStr = useMemo(() => {
    if (tokens.in == null || tokens.out == null) return null;
    return formatCost(
      estimateCost({
        provider: providerLabel,
        model: modelLabel,
        tokensIn: tokens.in,
        tokensOut: tokens.out,
      }),
    );
  }, [providerLabel, modelLabel, tokens.in, tokens.out]);

  const showResult = streaming || phase.kind === 'failed' || workout != null;

  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="Describe today's workout">
          <textarea
            value={userInput}
            onChange={(e) => setUserInput(e.target.value)}
            placeholder="e.g. Upper body, focus on shoulders. Overhead press, 4 working sets. Abs: landmine rotation, oblique cable, landmine hollow-body hold. Pull-ups, biceps and triceps."
            rows={6}
            className={inputClass}
            disabled={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={streaming || workoutCount === 0}
            className="mt-0.5"
          />
          <span>
            Use my history to suggest weights{' '}
            <span className="text-zinc-500">
              ({workoutCount === 0
                ? 'no workouts logged yet — disabled'
                : 'last 90 days · recent working weights per exercise'}
              )
            </span>
          </span>
        </label>

        <button
          type="button"
          onClick={() => runGeneration(userInput)}
          disabled={!userInput.trim() || streaming}
          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" />
          {workout ? 'Regenerate' : 'Generate workout'}
        </button>
      </section>

      {showResult && (
        <section className="space-y-3">
          <div className="flex items-center justify-between">
            <h2 className="text-sm font-semibold text-white uppercase tracking-wider">
              {streaming ? 'Generating…' : 'Suggested workout'}
            </h2>
            <span className="text-[11px] text-zinc-500 uppercase tracking-wider">
              {tokens.in != null && (
                <>
                  {tokens.in} in · {tokens.out ?? '?'} out
                </>
              )}
              {costStr && <> · {costStr}</>}
              {tokens.durationMs != null && (
                <> · {(tokens.durationMs / 1000).toFixed(1)}s</>
              )}
            </span>
          </div>

          {streaming && (
            <>
              {phase.lastPartial ? (
                <PartialPreview partial={phase.lastPartial} />
              ) : (
                <div className="text-xs text-zinc-500 italic flex items-center gap-2">
                  <Loader2 className="w-3 h-3 animate-spin" />
                  Waiting for the first parseable JSON…
                </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>
              )}
            </>
          )}

          {!streaming && workout && (
            <WorkoutPreview
              workout={workout}
              setWorkout={(updater) => setWorkout((w) => (w ? updater(w) : w))}
              exercises={exercises}
              refineInput={refineInput}
              setRefineInput={setRefineInput}
              onRefine={() => runGeneration(refineInput, workout)}
              onUse={(draft) => {
                sessionStorage.setItem(AI_WORKOUT_DRAFT_KEY, JSON.stringify(draft));
                router.push('/main/workouts/new?from=ai');
              }}
            />
          )}
        </section>
      )}
    </div>
  );
}

function WorkoutPreview({
  workout,
  setWorkout,
  exercises,
  refineInput,
  setRefineInput,
  onRefine,
  onUse,
}: {
  workout: AIWorkout;
  setWorkout: (updater: (w: AIWorkout) => AIWorkout) => void;
  exercises: LibraryExercise[];
  refineInput: string;
  setRefineInput: (v: string) => void;
  onRefine: () => void;
  onUse: (draft: AiWorkoutDraft) => void;
}) {
  const [error, setError] = useState<string | null>(null);

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

  const unresolvedCount = useMemo(
    () =>
      workout.exercises.filter(
        (ex) => !ex.exerciseId || !exerciseLookup.has(ex.exerciseId),
      ).length,
    [workout, exerciseLookup],
  );

  const updateExercise = (
    idx: number,
    patch: Partial<AIWorkoutExercise>,
  ) => {
    setWorkout((w) => {
      const next = structuredClone(w);
      next.exercises[idx] = { ...next.exercises[idx], ...patch };
      return next;
    });
  };

  const removeExercise = (idx: number) => {
    setWorkout((w) => {
      const next = structuredClone(w);
      next.exercises.splice(idx, 1);
      next.exercises.forEach((ex, i) => (ex.order = i));
      return next;
    });
  };

  const numOrNull = (v: string) => {
    if (v.trim() === '') return null;
    const n = Number(v);
    return Number.isFinite(n) ? n : null;
  };

  const handleUse = () => {
    if (unresolvedCount > 0) {
      setError(`Map or remove the ${unresolvedCount} unknown exercise(s) first.`);
      return;
    }
    setError(null);
    onUse({
      name: workout.name,
      notes: workout.notes ?? undefined,
      exercises: workout.exercises.map((ex) => ({
        exerciseId: ex.exerciseId!,
        sets: ex.sets && ex.sets > 0 ? ex.sets : 3,
        reps: ex.reps ?? undefined,
        suggestedWeight: ex.suggestedWeight ?? undefined,
        suggestedWeightUnit: ex.suggestedWeightUnit ?? undefined,
        rpe: ex.rpe ?? undefined,
        gear: ex.gear ?? undefined,
        durationSeconds: ex.durationSeconds ?? undefined,
        notes: ex.notes ?? undefined,
      })),
    });
  };

  return (
    <div className="bg-zinc-900 border border-zinc-800 rounded p-4 space-y-4">
      <div>
        <input
          value={workout.name}
          onChange={(e) => setWorkout((w) => ({ ...w, name: e.target.value }))}
          className="text-lg font-bold text-white bg-transparent border-b border-transparent hover:border-zinc-700 focus:border-zinc-500 focus:outline-none w-full"
        />
        {workout.notes && (
          <p className="text-sm text-zinc-400 mt-1 italic">{workout.notes}</p>
        )}
        <p className="text-xs text-zinc-500 mt-1">
          {workout.exercises.length} exercise
          {workout.exercises.length === 1 ? '' : 's'}
        </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 using this workout.
        </div>
      )}

      <ul className="space-y-2">
        {workout.exercises.map((ex, idx) => {
          const isUnknown = !ex.exerciseId || !exerciseLookup.has(ex.exerciseId);
          const lib = ex.exerciseId ? exerciseLookup.get(ex.exerciseId) : null;
          return (
            <li
              key={idx}
              className={`rounded p-3 ${
                isUnknown
                  ? 'bg-amber-950/30 border border-amber-900'
                  : 'bg-zinc-950 border border-zinc-800'
              }`}
            >
              <div className="flex items-start justify-between gap-2">
                <div className="text-white text-sm">
                  {lib?.name ?? ex.exerciseName}
                  {isUnknown && (
                    <span className="ml-2 text-[10px] uppercase tracking-wider text-amber-400">
                      not in library
                    </span>
                  )}
                </div>
                <button
                  type="button"
                  onClick={() => removeExercise(idx)}
                  className="text-xs text-red-400 hover:text-red-300 px-1"
                  title="Remove from workout"
                >
                  ✕
                </button>
              </div>

              {isUnknown && (
                <select
                  value={ex.exerciseId ?? ''}
                  onChange={(e) =>
                    updateExercise(idx, { exerciseId: e.target.value || null })
                  }
                  className="mt-2 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 className="mt-2 grid grid-cols-3 gap-2">
                <NumField
                  label="Sets"
                  value={ex.sets}
                  onChange={(v) => updateExercise(idx, { sets: v })}
                />
                <NumField
                  label="Reps"
                  value={ex.reps}
                  onChange={(v) => updateExercise(idx, { reps: v })}
                />
                <NumField
                  label={`Weight${ex.suggestedWeightUnit ? ` (${ex.suggestedWeightUnit})` : ''}`}
                  value={ex.suggestedWeight}
                  step="any"
                  onChange={(v) => updateExercise(idx, { suggestedWeight: v })}
                />
              </div>
              {ex.notes && (
                <div className="text-xs text-zinc-400 mt-2 italic">{ex.notes}</div>
              )}
            </li>
          );
        })}
      </ul>

      <div className="border-t border-zinc-800 pt-4 space-y-3">
        <Field label="Refine (send a change back to the AI)">
          <div className="flex gap-2">
            <input
              value={refineInput}
              onChange={(e) => setRefineInput(e.target.value)}
              onKeyDown={(e) => {
                // Cleared by the parent only on a successful regeneration, so
                // a failed refine keeps what the user typed.
                if (e.key === 'Enter' && refineInput.trim()) onRefine();
              }}
              placeholder="e.g. make overhead press 5 sets; swap the oblique exercise"
              className={inputClass}
            />
            <button
              type="button"
              onClick={() => {
                if (refineInput.trim()) onRefine();
              }}
              disabled={!refineInput.trim()}
              className="shrink-0 px-4 py-2 rounded bg-zinc-700 text-white font-bold text-xs uppercase tracking-wider hover:bg-zinc-600 disabled:bg-zinc-800 disabled:text-zinc-600"
            >
              Refine
            </button>
          </div>
        </Field>

        {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={handleUse}
          disabled={unresolvedCount > 0 || workout.exercises.length === 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"
        >
          Use this workout
        </button>
        <p className="text-[11px] text-zinc-500">
          Opens a pre-filled workout — nothing is saved until you save it there.
        </p>
      </div>
    </div>
  );

  function NumField({
    label,
    value,
    step,
    onChange,
  }: {
    label: string;
    value?: number | null;
    step?: string;
    onChange: (v: number | null) => void;
  }) {
    return (
      <label className="block">
        <span className="text-[10px] font-semibold text-zinc-500 uppercase tracking-wider block mb-0.5">
          {label}
        </span>
        <input
          type="number"
          inputMode="decimal"
          step={step}
          value={value ?? ''}
          onChange={(e) => onChange(numOrNull(e.target.value))}
          className="w-full px-2 py-1 text-sm rounded border border-zinc-700 bg-zinc-800 text-white focus:outline-none focus:ring-1 focus:ring-white/30"
        />
      </label>
    );
  }
}

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<AIWorkout> }) {
  const exercises = (partial.exercises as AIWorkoutExercise[] | 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 workout…{' '}
          {partial.name && (
            <span className="text-white font-semibold">{partial.name}</span>
          )}
        </span>
      </div>
      {exercises.length > 0 && (
        <ul className="text-xs text-zinc-300 space-y-1">
          {exercises.map((ex, i) => (
            <li key={i}>
              <span className="text-zinc-500">{(ex?.order ?? i) + 1}.</span>{' '}
              {ex?.exerciseName ?? '…'}
              {ex?.sets ? (
                <span className="text-zinc-500">
                  {' '}
                  · {ex.sets}×{ex.reps ?? '?'}
                  {ex.suggestedWeight != null
                    ? ` @ ${ex.suggestedWeight}${ex.suggestedWeightUnit ?? ''}`
                    : ''}
                </span>
              ) : null}
            </li>
          ))}
        </ul>
      )}
    </div>
  );
}
