import { useEffect, useMemo, useState } from 'react';
import type { ComponentType, ReactNode } from 'react';
import { useParams } from 'react-router-dom';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { Skeleton } from '@/components/ui/skeleton';
import { listRuns, loadEpisode, rewardTotal } from '@/lib/demo-kit/episode';
import { usePlayer } from '@/lib/demo-kit/player';
import { loadDemoModule } from '@/lib/demo-kit/registry';
import type { AnyDemoModule } from '@/lib/demo-kit/registry';
import type { DemoEpisode, DemoStep, RunRef, StoryBeat } from '@/lib/demo-kit/types';
import { useRunParam, useSpeedParam, useStepParam, useTabParam } from '@/lib/url-state';
import * as st from '@/content/styles';
import { cn } from '@/lib/utils';
import { BeatSection } from './BeatSection';
import { BlindCompare } from './BlindCompare';
import { CodeReceipt } from './CodeReceipt';
import { DemoErrorBoundary } from './DemoErrorBoundary';
import { EnvAnatomy } from './EnvAnatomy';
import { LimitsCallout } from './LimitsCallout';
import { MetricMover } from './MetricMover';
import { ModelCallPanel } from './ModelCallPanel';
import { PlayYourself } from './PlayYourself';
import { ProvenanceCard } from './ProvenanceCard';
import { ReasoningDrawer } from './ReasoningDrawer';
import { ReasoningPanel } from './ReasoningPanel';
import { RewardBreakdown } from './RewardBreakdown';
import { RewardEditor } from './RewardEditor';
import type { RewardArm } from './RewardEditor';
import { SegmentedControl } from './SegmentedControl';
import { SlotRegion } from './SlotRegion';
import { StatStrip } from './StatStrip';
import type { Stat } from './StatStrip';
import { StepTimeline } from './StepTimeline';
import { RecordedBadge, TracePlayer } from './TracePlayer';
import { VerifyBadge } from './VerifyBadge';
import { formatOrDash, useIsDesktop } from './format';

const REPO_BLOB = 'https://github.com/karti-ai/PIG-Demo/blob/main/';

/** The tab the step-detail strip opens on. Kept out of the URL when it is this. */
const DEFAULT_DETAIL_TAB = 'reasoning';

/** Reserved slug for the shell's own hand-written demo. Dev builds only. */
const MOCK_SLUG = '__mock';

export interface DemoBundle {
  demo: AnyDemoModule;
  runs: RunRef[];
  episodes: Record<string, DemoEpisode>;
}

type LoadState =
  | { status: 'loading' }
  | { status: 'ready'; bundle: DemoBundle }
  | { status: 'error'; message: string };

/**
 * A demo's module plus every recorded run it has.
 *
 * `loadDemoModule` and `loadEpisode` both cache their promises, so the route
 * loader having already fetched the module makes this resolve without a second
 * request. Runs are loaded with `allSettled` on purpose: one unreadable trace
 * drops that arm rather than blanking the page.
 */
async function loadBundle(slug: string): Promise<DemoBundle> {
  if (slug === MOCK_SLUG) {
    // Dynamic, so the mock lands in its own chunk and production never fetches
    // it. A static import would ship several hundred lines of fake trace to
    // every visitor of every real demo.
    const mock = await import('./mock');
    return { demo: mock.mockDemo, runs: mock.mockRuns, episodes: mock.mockEpisodes };
  }

  const demo = await loadDemoModule(slug);
  const runs = await listRuns(slug).catch(() => [] as RunRef[]);
  const settled = await Promise.allSettled(runs.map((run) => loadEpisode(run)));

  const episodes: Record<string, DemoEpisode> = {};
  settled.forEach((outcome, index) => {
    const run = runs[index];
    if (!run) return;
    if (outcome.status === 'fulfilled') episodes[run.id] = outcome.value;
    else console.error(`[pig-demo] dropped run "${run.id}":`, outcome.reason);
  });

  return { demo, runs: runs.filter((run) => episodes[run.id] !== undefined), episodes };
}

export interface DemoShellProps {
  /** Overrides the route param. Useful for previews and tests. */
  slug?: string;
  /** Skips loading entirely when the caller already has the bundle. */
  bundle?: DemoBundle;
}

/**
 * The route component every demo is rendered through.
 *
 * It owns four things and no more: loading, the narrative beats, the URL state,
 * and the page's single polite live region. Everything visual is delegated to
 * the surfaces in this directory, and the demo module is never reached into —
 * the shell only ever calls `adapt` and renders `Surface`.
 */
export function DemoShell({ slug: slugProp, bundle }: DemoShellProps) {
  const params = useParams();
  const slug = slugProp ?? params['slug'] ?? '';
  const [state, setState] = useState<LoadState>(
    bundle ? { status: 'ready', bundle } : { status: 'loading' },
  );

  useEffect(() => {
    if (bundle) {
      setState({ status: 'ready', bundle });
      return;
    }
    let live = true;
    setState({ status: 'loading' });
    loadBundle(slug)
      .then((loaded) => {
        if (live) setState({ status: 'ready', bundle: loaded });
      })
      .catch((error: unknown) => {
        if (!live) return;
        setState({
          status: 'error',
          message: error instanceof Error ? error.message : String(error),
        });
      });
    return () => {
      live = false;
    };
  }, [slug, bundle]);

  if (state.status === 'loading') return <ShellSkeleton />;
  if (state.status === 'error') {
    return (
      <main className={cn(st.shell, 'py-16')}>
        <div role="alert" className="card max-w-xl p-6">
          <h1 className={st.h2}>That demo is not here</h1>
          <p className={cn(st.prose, 'mt-3')}>{state.message}</p>
          <a href="/gallery" className={cn(st.btnSecondary, 'mt-6')}>
            Back to the demos
          </a>
        </div>
      </main>
    );
  }

  return (
    // A second boundary inside the route's own: this one is keyed to the demo
    // so a crash names it, and resetting re-renders the surfaces rather than
    // re-navigating.
    <DemoErrorBoundary demoTitle={state.bundle.demo.meta.title}>
      <DemoBody bundle={state.bundle} />
    </DemoErrorBoundary>
  );
}

function DemoBody({ bundle }: { bundle: DemoBundle }) {
  const { demo, runs, episodes } = bundle;
  const isDesktop = useIsDesktop();

  const [runParam, setRunParam] = useRunParam();
  const [stepParam, setStepParam] = useStepParam();
  const [tabParam, setTabParam] = useTabParam(DEFAULT_DETAIL_TAB);
  const [speedParam, setSpeedParam] = useSpeedParam();

  const run = useMemo(
    () => runs.find((candidate) => candidate.id === runParam) ?? runs[0],
    [runs, runParam],
  );
  const episode = run ? episodes[run.id] : undefined;

  const steps = useMemo<DemoStep<unknown>[]>(
    () => (episode ? demo.adapt(episode) : []),
    [demo, episode],
  );

  const player = usePlayer(steps, {
    initialIndex: stepParam,
    initialSpeed: speedParam,
    onIndexChange: setStepParam,
  });

  // The URL is the other writer of this state — Back, a pasted permalink, the
  // run switcher. The player is the source of truth while it is running, so it
  // only follows the URL when the two have actually diverged.
  const { seek } = player;
  useEffect(() => {
    if (stepParam !== player.index) seek(stepParam);
    // Intentionally keyed on the URL only: including `player.index` here would
    // re-run the effect on the player's own advance and fight it.
  }, [stepParam, seek]);

  const hasBeat = (surface: StoryBeat['surface']) =>
    demo.narrative.beats.some((beat) => beat.surface === surface);
  const timelineInSplit = !hasBeat('scrubber');
  const extras = demo.tabs ?? [];
  const extrasInCustomBeat = hasBeat('custom');

  const arms = useMemo<RewardArm[]>(
    () =>
      runs.map((candidate) => {
        const armEpisode = episodes[candidate.id];
        const arm: RewardArm = {
          id: candidate.id,
          label: candidate.label,
          values: armEpisode?.rewards ?? {},
        };
        if (candidate.intervention) arm.note = candidate.intervention;
        return arm;
      }),
    [runs, episodes],
  );

  const blindPair = useMemo(() => {
    for (let i = 0; i < runs.length; i += 1) {
      for (let j = i + 1; j < runs.length; j += 1) {
        const left = runs[i];
        const right = runs[j];
        if (!left || !right || left.seed !== right.seed) continue;
        const leftEpisode = episodes[left.id];
        const rightEpisode = episodes[right.id];
        if (!leftEpisode || !rightEpisode) continue;
        return { left, right, leftEpisode, rightEpisode };
      }
    }
    // Two runs on different seeds are two different puzzles; showing them side
    // by side would be a comparison of luck.
    return null;
  }, [runs, episodes]);

  if (!run || !episode || steps.length === 0) {
    return (
      <main className={cn(st.shell, 'py-16')}>
        <h1 className={st.h2}>{demo.meta.title}</h1>
        <p className={cn(st.prose, 'mt-3 max-w-prose')}>
          No recorded run is available for this demo yet. The environment and its grader are in
          the repository; the traces come from the eval command on the provenance card.
        </p>
        <EnvAnatomy anatomy={demo.anatomy} rewardLine={demo.meta.rewardLine} className="mt-8" />
      </main>
    );
  }

  const Surface = demo.Surface as ComponentType<{ state: unknown; compact?: boolean }>;
  const current = steps[player.index];
  const lastStep = steps[steps.length - 1];

  const heroStats: Stat[] = [
    {
      label: 'Outcome',
      value: episode.outcome,
      tone: episode.outcome === 'solved' ? 'positive' : 'warning',
      ...(episode.truncated ? { hint: 'Truncated before a terminal state' } : {}),
    },
    {
      label: 'Total reward',
      value: formatOrDash(rewardTotal(episode.rewards, demo.reward.components)),
      tone: 'brand',
      hint: 'Shipped weights',
    },
    { label: 'Steps', value: steps.length, hint: 'Model calls in this run' },
    { label: 'Seed', value: episode.seed, hint: 'The same seed reproduces this board' },
  ];

  const timeline = (
    <StepTimeline
      steps={steps}
      current={player.index}
      onSelect={(next) => {
        player.pause();
        player.seek(next);
      }}
      Surface={Surface}
      onTogglePlay={player.toggle}
    />
  );

  const detailTabs: { id: string; label: string; content: ReactNode }[] = [
    {
      id: 'reasoning',
      label: 'Reasoning',
      content: isDesktop ? (
        <ReasoningPanel
          reasoning={current?.reasoning ?? null}
          durationMs={current?.call?.durationMs ?? null}
          playing={player.isPlaying}
          speed={player.speed}
          stepIndex={player.index}
        />
      ) : (
        // Under `lg` there is no column for this, and putting it below the
        // board means watching the run with the thinking off-screen. The sheet
        // is mounted only here, so vaul never locks body scroll on desktop.
        <ReasoningDrawer
          reasoning={current?.reasoning ?? null}
          durationMs={current?.call?.durationMs ?? null}
          playing={player.isPlaying}
          speed={player.speed}
          stepIndex={player.index}
        />
      ),
    },
    {
      id: 'call',
      label: 'Model call',
      content: (
        <div className="space-y-3">
          <ModelCallPanel call={current?.call ?? null} />
          {current?.reply ? (
            <div className="card p-3">
              <h3 className="text-sm font-semibold">Reply</h3>
              <p className="mt-1 whitespace-pre-wrap font-mono text-xs leading-relaxed">
                {current.reply}
              </p>
            </div>
          ) : null}
        </div>
      ),
    },
    ...(extrasInCustomBeat
      ? []
      : extras.map((tab) => ({ id: tab.id, label: tab.label, content: <tab.Component /> }))),
  ];
  const activeTab = detailTabs.some((tab) => tab.id === tabParam) ? tabParam : DEFAULT_DETAIL_TAB;

  const renderSurface = (beat: StoryBeat): ReactNode => {
    switch (beat.surface) {
      case 'hero':
        return (
          <div className="grid grid-cols-1 gap-4 lg:grid-cols-[auto_minmax(0,1fr)] lg:items-start">
            <div className="card w-fit p-4">
              {lastStep ? <Surface state={lastStep.state} /> : null}
            </div>
            <div className="space-y-3">
              <StatStrip stats={heroStats} />
              <RecordedBadge
                model={run.model}
                capturedAt={run.capturedAt}
                {...(run.intervention ? { intervention: run.intervention } : {})}
                className="ml-0 w-fit"
              />
              <p className={cn(st.prose, 'max-w-prose')}>{demo.narrative.thesis}</p>
            </div>
          </div>
        );

      case 'anatomy':
        return <EnvAnatomy anatomy={demo.anatomy} rewardLine={demo.meta.rewardLine} />;

      case 'split-play':
        return (
          <div className="space-y-3">
            {demo.interactive ? (
              <>
                <PlayYourself demo={demo} seed={run.seed} />
                <div className="space-y-1 pt-4">
                  <h3 className={st.h3}>What the model did</h3>
                  <p className={cn(st.prose, 'max-w-prose text-sm')}>
                    Same hidden answer, same rules, same budget — replayed at the
                    speed it actually happened.
                  </p>
                </div>
              </>
            ) : null}

            {runs.length > 1 ? (
              <RunSwitcher
                runs={runs}
                activeId={run.id}
                onSelect={(id) => {
                  player.pause();
                  // The run param setter also zeroes `step`: step 6 of a
                  // nine-turn rollout is not step 6 of a three-turn one.
                  setRunParam(id);
                }}
              />
            ) : null}

            <TracePlayer
              playing={player.isPlaying}
              onPlayingChange={(next) => (next ? player.play() : player.pause())}
              speed={player.speed}
              onSpeedChange={(next) => {
                player.setSpeed(next);
                setSpeedParam(next);
                // `instant` is a destination, not a rate. The player only
                // consumes it while running, so choosing it from a paused
                // transport has to start the run — otherwise the button
                // visibly does nothing, which reads as broken.
                if (next === 'instant') player.play();
              }}
              onRestart={player.restart}
              step={player.index}
              stepCount={steps.length}
              onStepChange={(next) => {
                player.pause();
                player.seek(next);
              }}
              progress={player.progress}
              timingIsReal={player.timingIsReal}
              model={run.model}
              capturedAt={run.capturedAt}
              {...(run.intervention ? { intervention: run.intervention } : {})}
            />

            <div className="grid grid-cols-1 gap-3 lg:grid-cols-[auto_minmax(0,1fr)] lg:items-start">
              <div className="card w-fit p-4">
                {current ? <Surface state={current.state} /> : null}
              </div>

              <div className="min-w-0">
                <Tabs value={activeTab} onValueChange={setTabParam}>
                  <TabsList aria-label="Details for this step" className="w-full overflow-x-auto">
                    {detailTabs.map((tab) => (
                      <TabsTrigger key={tab.id} value={tab.id}>
                        {tab.label}
                      </TabsTrigger>
                    ))}
                  </TabsList>
                  {detailTabs.map((tab) => (
                    <TabsContent key={tab.id} value={tab.id}>
                      {tab.content}
                    </TabsContent>
                  ))}
                </Tabs>
              </div>
            </div>

            {timelineInSplit ? timeline : null}
            <SlotRegion id="below-board" />
          </div>
        );

      case 'scrubber':
        return (
          <div className="space-y-3">
            {timeline}
            {player.timingIsReal ? null : (
              <p className="text-xs text-muted">
                Some steps in this run carried no recorded latency, so their dwell on the timeline
                is the player's fallback rather than a measurement.
              </p>
            )}
            <SlotRegion id="below-timeline" />
          </div>
        );

      case 'reward-editor':
        return (
          <div className="space-y-4">
            <RewardBreakdown
              spec={demo.reward}
              values={episode.rewards}
              {...(episode.metrics ? { metrics: episode.metrics } : {})}
            />
            <VerifyBadge demo={demo} episode={episode} />
            {arms.length > 1 ? <RewardEditor spec={demo.reward} arms={arms} /> : null}
          </div>
        );

      case 'metric': {
        const currentTotal = rewardTotal(episode.rewards, demo.reward.components);
        const points = arms
          .map((arm) => ({ x: arm.label, y: rewardTotal(arm.values, demo.reward.components) }))
          .filter((point): point is { x: string; y: number } => point.y !== null);
        const baselineArm = arms[0];
        const baselineValue = baselineArm
          ? rewardTotal(baselineArm.values, demo.reward.components)
          : null;
        return (
          <div className="space-y-4">
            <MetricMover
              label={`Total reward — ${run.label}`}
              value={currentTotal ?? 0}
              {...(baselineArm && baselineValue !== null && arms.length > 1
                ? { baseline: { value: baselineValue, label: baselineArm.label } }
                : {})}
              series={points}
              caption="Every point is a recorded run scored by the same grader. Nothing here is a projection."
            />
            {blindPair ? (
              <BlindCompare
                seed={blindPair.left.seed}
                Surface={Surface}
                a={{
                  runId: blindPair.left.id,
                  label: blindPair.left.label,
                  model: blindPair.left.model,
                  ...(blindPair.left.intervention
                    ? { intervention: blindPair.left.intervention }
                    : {}),
                  steps: demo.adapt(blindPair.leftEpisode),
                  total: rewardTotal(blindPair.leftEpisode.rewards, demo.reward.components),
                }}
                b={{
                  runId: blindPair.right.id,
                  label: blindPair.right.label,
                  model: blindPair.right.model,
                  ...(blindPair.right.intervention
                    ? { intervention: blindPair.right.intervention }
                    : {}),
                  steps: demo.adapt(blindPair.rightEpisode),
                  total: rewardTotal(blindPair.rightEpisode.rewards, demo.reward.components),
                }}
              />
            ) : null}
          </div>
        );
      }

      case 'receipt':
        return (
          <div className="grid grid-cols-1 gap-4 lg:grid-cols-2 lg:items-start">
            <ProvenanceCard provenance={demo.provenance} run={run} />
            <CodeReceipt
              code={demo.reward.source.code}
              path={demo.reward.source.path}
              {...(demo.reward.source.marker ? { marker: demo.reward.source.marker } : {})}
              href={`${REPO_BLOB}${demo.reward.source.path}`}
            />
            <SlotRegion id="after-receipts" className="lg:col-span-2" />
          </div>
        );

      case 'limits':
        return <LimitsCallout limits={demo.narrative.limits} />;

      case 'custom':
        return extras.length > 0 ? (
          <div className="space-y-4">
            {extras.map((tab) => (
              <tab.Component key={tab.id} />
            ))}
          </div>
        ) : (
          <SlotRegion id="before-limits" />
        );

      default:
        return null;
    }
  };

  return (
    <main className={cn(st.shell, 'pb-16')}>
      {/*
        The page's ONE live region. Every step change lands here and nowhere
        else: with reduced motion the board animation is gone, so this sentence
        is the only thing that tells a screen-reader user what just happened.
      */}
      <div aria-live="polite" aria-atomic="true" className="sr-only">
        {current?.announce ?? ''}
      </div>

      <header className="pt-8">
        <p className={cn(st.eyebrow, 'text-accent-fg')}>For {demo.meta.persona}</p>
        <h1 className={cn(st.h2, 'mt-1')}>{demo.meta.title}</h1>
        <p className={cn(st.lede, 'mt-2 max-w-prose')}>{demo.meta.tagline}</p>
        <p className="mt-4 max-w-prose border-l-2 border-brand pl-3 text-sm italic leading-relaxed text-fg">
          {demo.narrative.anxiety}
        </p>
      </header>

      <div className="divide-y divide-border">
        {demo.narrative.beats.map((beat, index) => (
          <BeatSection key={beat.id} beat={beat} number={index + 1}>
            {renderSurface(beat)}
          </BeatSection>
        ))}
      </div>
    </main>
  );
}

function RunSwitcher({
  runs,
  activeId,
  onSelect,
}: {
  runs: RunRef[];
  activeId: string;
  onSelect: (id: string) => void;
}) {
  // Two axes, not one list. With four arms over eight seeds a flat control is
  // thirty buttons carrying four distinct labels, which reads as a bug. Arms
  // are grouped by `label` because that is what an arm IS in the manifest —
  // the shell has no other notion of one, and inventing a field for it would
  // put demo-specific structure into the contract.
  const arms = useMemo(() => {
    const byLabel = new Map<string, RunRef[]>();
    for (const run of runs) {
      const list = byLabel.get(run.label);
      if (list) list.push(run);
      else byLabel.set(run.label, [run]);
    }
    return [...byLabel.entries()].map(([label, group]) => ({ label, runs: group }));
  }, [runs]);

  const active = runs.find((r) => r.id === activeId) ?? runs[0];
  if (!active) return null;
  const activeArm = arms.find((a) => a.label === active.label) ?? arms[0];
  if (!activeArm) return null;

  const pickArm = (label: string) => {
    const arm = arms.find((a) => a.label === label);
    if (!arm) return;
    // Hold the seed across an arm change where the arm has it. Comparing two
    // agents means comparing them on the SAME hidden word; silently jumping to
    // a different seed would make the comparison meaningless while looking fine.
    const sameSeed = arm.runs.find((r) => r.seed === active.seed);
    onSelect((sameSeed ?? arm.runs[0])!.id);
  };

  return (
    <div className="flex flex-wrap items-center gap-x-4 gap-y-2">
      <SegmentedControl
        label="Agent"
        options={arms.map((arm) => {
          const intervention = arm.runs.find((r) => r.intervention)?.intervention;
          return {
            value: arm.label,
            label: arm.label,
            ...(intervention ? { title: intervention } : {}),
          };
        })}
        value={activeArm.label}
        onChange={pickArm}
        className="border border-border p-1"
        optionClassName="tap px-3 text-sm"
      />
      {activeArm.runs.length > 1 ? (
        <SegmentedControl
          label="Hidden word"
          options={activeArm.runs.map((run) => ({
            value: run.id,
            label: `#${run.seed}`,
          }))}
          value={active.id}
          onChange={onSelect}
          className="border border-border p-1"
          optionClassName="tap px-2.5 text-xs nums"
        />
      ) : null}
    </div>
  );
}


/**
 * The loading state. Shaped like the page it becomes, and with no spinner: a
 * spinner here would imply a live model call, which is the one thing the whole
 * page is at pains to say is not happening.
 */
function ShellSkeleton() {
  return (
    <div className={cn(st.shell, 'py-10')} aria-busy="true">
      <p className="sr-only">Loading the recorded run.</p>
      <Skeleton className="h-8 w-64" />
      <Skeleton className="mt-3 h-4 w-full max-w-md" />
      <div className="mt-10 grid grid-cols-1 gap-3 sm:grid-cols-2 lg:grid-cols-4">
        {[0, 1, 2, 3].map((index) => (
          <Skeleton key={index} className="h-28" />
        ))}
      </div>
      <Skeleton className="mt-6 h-64" />
    </div>
  );
}
