"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";

import { LyricsEditor } from "@/components/lyrics-editor";
import { Queue } from "@/components/queue";
import { EMPTY_REFERENCE, ReferenceDrop, type ReferenceState } from "@/components/reference-drop";
import { RenderPanel, type RenderSettings } from "@/components/render-panel";
import { ScorePanel } from "@/components/score-panel";
import { StyleField } from "@/components/style-field";
import { Transport } from "@/components/transport";
import { IDLE_LIFT, useLyricLift } from "@/lib/lyrics-lift";
import { DEFAULT_ADVANCED, STARTER_LYRICS, STARTER_STYLE } from "@/lib/presets";
import type { ComfyCapabilities, Job } from "@/lib/types";

/**
 * The SSE stream emits a job the moment the store creates it, which can beat
 * the POST response that also carries it. Merging by id keeps one row per job
 * whichever arrives first, and keeps the newer snapshot of the two.
 */
function upsertJob(jobs: Job[], incoming: Job): Job[] {
  const existing = jobs.find((job) => job.id === incoming.id);
  const merged = existing ? { ...existing, ...incoming } : incoming;
  return [merged, ...jobs.filter((job) => job.id !== incoming.id)].sort(
    (a, b) => b.createdAt - a.createdAt,
  );
}

const INITIAL_SETTINGS: RenderSettings = {
  scoreSource: "model",
  mode: "full",
  maxDuration: 240,
  seed: 0,
  randomizeSeed: true,
  tiledDecode: false,
  realAudioMatch: false,
  realAudioStrength: 1,
  instrumentalArrangement: false,
  instrumentalStrength: 1,
  outputFormat: "flac",
  checkpoint: "",
  audioEncoder: "",
  advanced: DEFAULT_ADVANCED,
};

export function Studio() {
  const [style, setStyle] = useState(STARTER_STYLE);
  const [lyrics, setLyrics] = useState(STARTER_LYRICS);
  const [reference, setReference] = useState<ReferenceState>(EMPTY_REFERENCE);
  const [settings, setSettings] = useState<RenderSettings>(INITIAL_SETTINGS);
  const [capabilities, setCapabilities] = useState<ComfyCapabilities | null>(null);
  const [jobs, setJobs] = useState<Job[]>([]);
  const [selectedId, setSelectedId] = useState<string | null>(null);
  const [manualScore, setManualScore] = useState("");
  const { lift, start: startLift, reset: resetLift, setLift } = useLyricLift();
  const [scoreOpen, setScoreOpen] = useState(false);
  const [playToken, setPlayToken] = useState(0);
  const [busy, setBusy] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const autoSelected = useRef<string | null>(null);

  // Safari will not decode FLAC, so do not hand it a FLAC by default.
  useEffect(() => {
    const probe = document.createElement("audio");
    if (!probe.canPlayType("audio/flac") && !probe.canPlayType("audio/x-flac")) {
      setSettings((current) => ({ ...current, outputFormat: "mp3" }));
    }
  }, []);

  // Keep looking until ComfyUI answers, so starting it later just works.
  useEffect(() => {
    let cancelled = false;
    let timer: ReturnType<typeof setTimeout> | null = null;

    const poll = () => {
      void fetch("/api/capabilities")
        .then((res) => res.json())
        .then((data: ComfyCapabilities) => {
          if (cancelled) return;
          setCapabilities(data);
          setSettings((current) => ({
            ...current,
            checkpoint: current.checkpoint || pickWeight(data.checkpoints, /yue2/i),
            audioEncoder: current.audioEncoder || pickWeight(data.audioEncoders, /sheetsage/i),
          }));
          if (!data.connected) timer = setTimeout(poll, 8000);
        })
        .catch(() => {
          if (!cancelled) timer = setTimeout(poll, 8000);
        });
    };

    poll();
    return () => {
      cancelled = true;
      if (timer) clearTimeout(timer);
    };
  }, []);

  useEffect(() => {
    const source = new EventSource("/api/events");

    source.addEventListener("snapshot", (event) => {
      setJobs(JSON.parse((event as MessageEvent).data) as Job[]);
    });

    source.addEventListener("job", (event) => {
      setJobs((current) => upsertJob(current, JSON.parse((event as MessageEvent).data) as Job));
    });

    return () => source.close();
  }, []);

  // The newest finished take takes over the transport as soon as it lands.
  useEffect(() => {
    const done = jobs.find((job) => job.stage === "done");
    if (done && autoSelected.current !== done.id) {
      autoSelected.current = done.id;
      setSelectedId(done.id);
    }
  }, [jobs]);

  const selected = useMemo(
    () => jobs.find((job) => job.id === selectedId) ?? null,
    [jobs, selectedId],
  );

  const running = useMemo(
    () => jobs.find((job) => !["done", "failed", "cancelled"].includes(job.stage)) ?? null,
    [jobs],
  );

  // An edited score wins over whatever the selected take came back with.
  const activeScore = manualScore || selected?.abc || "";

  // Dropping a reference switches the score source, because that is what the
  // drop was for. It used to force the plan to melody as well, which quietly
  // undid a deliberate "melody and chords" choice, and that is the setting
  // that decides how much of the reference survives. The plan is left alone.
  const handleReference = useCallback(
    (next: ReferenceState) => {
      setReference(next);
      setSettings((current) => ({
        ...current,
        scoreSource:
          next.status === "ready"
            ? "reference"
            : current.scoreSource === "reference"
              ? "model"
              : current.scoreSource,
      }));

      if (next.status === "ready" && next.filename) {
        void startLift(next.filename, next.start, next.duration);
      } else if (next.status !== "uploading") {
        resetLift();
      }
    },
    [startLift, resetLift],
  );

  // Filling the field is the point, but only when there is nothing to lose.
  // Untouched starter text is not the operator's work; anything they typed is,
  // and that only ever gets replaced by someone pressing the button.
  const lyricsAreDisposable = !lyrics.trim() || lyrics === STARTER_LYRICS;

  useEffect(() => {
    if (lift.status !== "done" || !lift.lyrics) return;
    if (!lyricsAreDisposable) return;
    setLyrics(lift.lyrics);
    setLift((current) => ({ ...current, lyrics: null }));
  }, [lift.status, lift.lyrics, lyricsAreDisposable, setLift]);

  const applyLift = useCallback(() => {
    if (!lift.lyrics) return;
    setLyrics(lift.lyrics);
    setLift(IDLE_LIFT);
  }, [lift.lyrics, setLift]);

  const retryLift = useCallback(() => {
    if (reference.status !== "ready" || !reference.filename) return;
    void startLift(reference.filename, reference.start, reference.duration);
  }, [reference.status, reference.filename, reference.start, reference.duration, startLift]);

  const render = useCallback(async () => {
    setBusy(true);
    setError(null);
    try {
      const res = await fetch("/api/generate", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          style,
          lyrics,
          mode: settings.mode,
          scoreSource: settings.scoreSource,
          abc: settings.scoreSource === "manual" ? activeScore : "",
          referenceFilename: reference.filename,
          referenceStart: reference.start,
          referenceDuration: reference.duration,
          maxDuration: settings.maxDuration,
          seed: settings.seed,
          randomizeSeed: settings.randomizeSeed,
          checkpoint: settings.checkpoint,
          audioEncoder: settings.audioEncoder,
          tiledDecode: settings.tiledDecode,
          realAudioMatch: settings.realAudioMatch,
          realAudioStrength: settings.realAudioStrength,
          instrumentalArrangement: settings.instrumentalArrangement,
          instrumentalStrength: settings.instrumentalStrength,
          outputFormat: settings.outputFormat,
          advanced: settings.advanced,
        }),
      });
      const body = await res.json();
      if (!res.ok) throw new Error(body.error ?? "The render could not start");
      setJobs((current) => upsertJob(current, body as Job));
    } catch (caught) {
      setError(caught instanceof Error ? caught.message : "The render could not start");
    } finally {
      setBusy(false);
    }
  }, [style, lyrics, settings, reference.filename, reference.start, reference.duration, activeScore]);

  const cancel = useCallback(async (job: Job) => {
    await fetch(`/api/jobs/${job.id}`, { method: "DELETE" }).catch(() => undefined);
  }, []);

  return (
    <div className="flex min-h-[100dvh] flex-col">
      <Header capabilities={capabilities} />

      <main className="mx-auto grid w-full max-w-[1560px] flex-1 gap-5 px-5 py-6 sm:px-7 lg:grid-cols-[minmax(0,1.55fr)_minmax(330px,0.85fr)]">
        <div className="flex min-w-0 flex-col gap-5">
          <StyleField value={style} onChange={setStyle} />
          <ReferenceDrop state={reference} onChange={handleReference} />
          <LyricsEditor
            value={lyrics}
            onChange={setLyrics}
            lift={lift}
            onApplyLift={applyLift}
            onDismissLift={() => setLift(IDLE_LIFT)}
            onRetryLift={retryLift}
          />
        </div>

        {/*
          Deliberately not a sticky, independently scrolling column. It used to
          be, sized against a hardcoded header height, and the moment the error
          banner expanded the header grew past that number: the column's bottom
          ran off screen and Takes could only be reached through a thin inner
          scrollbar. One page scroll cannot trap anything.
        */}
        <div className="flex min-w-0 flex-col gap-5">
          <RenderPanel
            settings={settings}
            onChange={setSettings}
            capabilities={capabilities}
            hasReference={reference.status === "ready"}
            hasScore={Boolean(activeScore)}
            onOpenScore={() => setScoreOpen(true)}
            busy={busy}
            error={error}
            onRender={render}
          />
          <Queue
            jobs={jobs}
            selectedId={selectedId}
            onSelect={(job) => {
              setSelectedId(job.id);
              setManualScore("");
              setPlayToken((token) => token + 1);
            }}
            onCancel={cancel}
          />
        </div>
      </main>

      <Transport
        job={selected}
        running={running}
        playToken={playToken}
        onOpenScore={() => setScoreOpen(true)}
      />

      <ScorePanel
        open={scoreOpen}
        score={activeScore}
        onClose={() => setScoreOpen(false)}
        onApply={(next) => {
          setManualScore(next);
          setSettings((current) => ({ ...current, scoreSource: "manual" }));
        }}
      />
    </div>
  );
}

/** Prefer the weight that looks like the one we want over whatever sorts first. */
function pickWeight(options: string[], wanted: RegExp): string {
  return options.find((name) => wanted.test(name)) ?? options[0] ?? "";
}

function Header({ capabilities }: { capabilities: ComfyCapabilities | null }) {
  const connected = capabilities?.connected ?? false;
  const missing = capabilities?.missingNodes ?? [];

  return (
    <>
      {/*
        Only the bar is sticky. The notices below it are diagnostic, and an
        expanded engine log is tall enough that pinning it would eat a third of
        a laptop viewport and push the right column out of reach.
      */}
      <header className="sticky top-0 z-20 border-b border-line bg-void/90 backdrop-blur">
        <div className="mx-auto flex h-16 max-w-[1560px] items-center justify-between gap-6 px-5 sm:px-7">
          <div className="flex items-baseline gap-3">
            <span className="text-[15px] font-medium tracking-tight text-ink">YuE2 Studio</span>
            <span className="hidden text-[12px] text-ink-faint sm:inline">
              Full songs from a style, a lyric sheet and a reference
            </span>
          </div>

          <div className="flex items-center gap-2.5">
            {/* Real semantic state, the only status indicator on the page. */}
            <span
              aria-hidden
              className={`h-1.5 w-1.5 rounded-full ${connected ? "bg-citron" : "bg-alarm"}`}
            />
            <span className="font-mono text-[11.5px] text-ink-dim">
              {capabilities === null
                ? "checking"
                : capabilities.mock
                  ? "mock backend"
                  : connected
                    ? capabilities.url.replace(/^https?:\/\//, "")
                    : "ComfyUI unreachable"}
            </span>
            {capabilities?.buildStamp ? (
              <span
                className="hidden font-mono text-[11px] text-ink-faint sm:inline"
                title="The image this container was built from"
              >
                {capabilities.buildStamp}
              </span>
            ) : null}
            </div>
          </div>
      </header>

      {capabilities && !capabilities.connected ? (
        <div className="border-b border-alarm/30 bg-alarm/10 px-5 py-2 sm:px-7">
          <p className="text-center text-[12px] text-alarm">
            Cannot reach ComfyUI at {capabilities.url}. It may still be loading the model, and this
            page picks it up on its own once it answers.
          </p>
          {capabilities.engineLog ? (
            <details className="mx-auto mt-1.5 max-w-[80ch]">
              <summary className="cursor-pointer text-center text-[11.5px] text-alarm/80">
                Engine log
              </summary>
              <pre className="console-scroll mt-1.5 max-h-40 overflow-auto rounded-console border border-alarm/30 bg-void px-3 py-2 font-mono text-[11px] leading-relaxed whitespace-pre-wrap text-ink-dim">
                {capabilities.engineLog}
              </pre>
            </details>
          ) : null}
        </div>
      ) : null}

      {missing.length > 0 ? (
        <p className="border-b border-line bg-raised px-5 py-2 text-center text-[12px] text-ink-dim sm:px-7">
          This ComfyUI build is missing {missing.slice(0, 3).join(", ")}. Update to v0.35.0 or newer.
        </p>
      ) : null}
    </>
  );
}