MelodyMuse/src/pages/HomePage.tsx

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { Youtube } from "lucide-react";
import { InputPanel } from "../components/InputPanel";
import type { InputValues } from "../components/InputPanel";
import {
  ResultsPanel,
  type RegeneratingMap,
  type SectionKey,
} from "../components/ResultsPanel";
import { StickyZipBar } from "../components/StickyZipBar";
import { HistoryPanel } from "../components/HistoryPanel";
import { Footer } from "../components/Footer";
import { EqualizerIcon } from "../components/EqualizerIcon";
import { useToast } from "../lib/toast";
import { generateSong, getServerStatus } from "../lib/llm";
import { addToHistory, type HistoryEntry } from "../lib/history";
import { buildSongZip, downloadBlob, sanitizeFilename } from "../lib/zip";
import {
  SECTION_LABELS,
  type SongAssets,
  type VideoPrompt,
} from "../lib/types";

// localStorage key for the in-progress input form. Lets the user survive
// an accidental refresh without losing what they were typing.
const DRAFT_KEY = "melodymuse-draft";
const DRAFT_DEBOUNCE_MS = 400;

const DEFAULT_INPUT: InputValues = {
  idea: "",
  language: "English",
  customLanguage: "",
  style_hint: "",
  mood: "",
  vocals: "vocals",
};

function resolveLanguage(v: InputValues): string {
  return v.language === "Other"
    ? v.customLanguage.trim() || "English"
    : v.language;
}

function isAnyBusy(loading: boolean, regenerating: RegeneratingMap): boolean {
  if (loading) return true;
  for (const v of Object.values(regenerating)) if (v) return true;
  return false;
}

function loadDraft(): InputValues | null {
  if (typeof window === "undefined") return null;
  try {
    const raw = window.localStorage.getItem(DRAFT_KEY);
    if (!raw) return null;
    const parsed = JSON.parse(raw) as Partial<InputValues> | null;
    if (!parsed || typeof parsed !== "object") return null;
    return { ...DEFAULT_INPUT, ...parsed };
  } catch {
    return null;
  }
}

export function HomePage() {
  const toast = useToast();

  // Hydrate the input from any previously-saved draft so the user can recover
  // an accidental refresh.
  const [input, setInput] = useState<InputValues>(
    () => loadDraft() ?? DEFAULT_INPUT,
  );
  const [assets, setAssets] = useState<SongAssets | null>(null);
  // The "last generated" snapshot. The cards compare their current editable
  // values against these to decide whether to show the Revert button.
  const [originalAssets, setOriginalAssets] = useState<SongAssets | null>(null);
  const [loading, setLoading] = useState(false);
  const [regenerating, setRegenerating] = useState<RegeneratingMap>({});
  const [selectedTitleIndex, setSelectedTitleIndex] = useState(0);
  const [serverOk, setServerOk] = useState<boolean | null>(null);

  // AbortController for the *current* network request. Refs are used so the
  // cancel handler always sees the latest value without re-binding.
  const abortRef = useRef<AbortController | null>(null);

  const anyRegenerating = isAnyBusy(loading, regenerating);

  // Debounced auto-save: persist the input form ~400ms after the last edit so
  // a refresh doesn't wipe what the user was typing.
  useEffect(() => {
    const id = window.setTimeout(() => {
      try {
        window.localStorage.setItem(DRAFT_KEY, JSON.stringify(input));
      } catch {
        // Quota or disabled — fine, history is the more durable store.
      }
    }, DRAFT_DEBOUNCE_MS);
    return () => window.clearTimeout(id);
  }, [input]);

  // Check the server on mount. If unreachable, surface a quiet header hint.
  useEffect(() => {
    const ctrl = new AbortController();
    (async () => {
      try {
        const s = await getServerStatus(ctrl.signal);
        setServerOk(Boolean(s.ok));
      } catch {
        setServerOk(false);
      }
    })();
    return () => ctrl.abort();
  }, []);

  // Clean up any in-flight request on unmount.
  useEffect(() => () => abortRef.current?.abort(), []);

  const selectedTitle = useMemo(() => {
    if (!assets || assets.titles.length === 0) return undefined;
    return assets.titles[
      Math.min(selectedTitleIndex, assets.titles.length - 1)
    ];
  }, [assets, selectedTitleIndex]);

  const cancel = useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
  }, []);

  // Esc cancels any in-flight generation.
  useEffect(() => {
    if (!anyRegenerating) return;
    const onKey = (e: KeyboardEvent) => {
      if (e.key === "Escape") {
        e.preventDefault();
        cancel();
      }
    };
    window.addEventListener("keydown", onKey);
    return () => window.removeEventListener("keydown", onKey);
  }, [anyRegenerating, cancel]);

  // Build the request payload from the current form values. Kept as a
  // function so all three call sites (full / regenerate all / per-section)
  // stay in sync.
  const buildRequest = useCallback(
    (
      section: GenerateCall["section"],
      extra: Partial<GenerateCall> = {},
    ): GenerateCall => ({
      input: input.idea.trim(),
      language: resolveLanguage(input),
      ...(input.mood.trim() ? { mood: input.mood.trim() } : {}),
      ...(input.style_hint.trim()
        ? { style_hint: input.style_hint.trim() }
        : {}),
      vocals: input.vocals,
      section,
      ...extra,
    }),
    [input],
  );

  const runGeneration = useCallback(
    async (
      build: () => GenerateCall,
      onSuccess: (result: Partial<SongAssets>) => void,
      busySetter: (b: boolean) => void,
      successMsg: string,
    ) => {
      // Cancel any in-flight request before starting a new one.
      abortRef.current?.abort();
      const ctrl = new AbortController();
      abortRef.current = ctrl;

      busySetter(true);
      try {
        const req = build();
        const result = await generateSong(req, { signal: ctrl.signal });
        onSuccess(result);
        toast.success(successMsg);
      } catch (err) {
        if (err instanceof DOMException && err.name === "AbortError") {
          toast.info("Generation cancelled");
        } else {
          toast.error(
            err instanceof Error
              ? err.message
              : "Generation failed — please try again",
          );
        }
      } finally {
        busySetter(false);
        if (abortRef.current === ctrl) abortRef.current = null;
      }
    },
    [toast],
  );

  const handleGenerate = useCallback(async () => {
    if (!input.idea.trim()) return;
    // Clear the old assets so the skeleton state shows during generation
    // instead of the previous song's cards.
    setAssets(null);
    setOriginalAssets(null);
    await runGeneration(
      () => buildRequest("all"),
      (result) => {
        const next = result as SongAssets;
        setAssets(next);
        setOriginalAssets(next);
        setSelectedTitleIndex(0);
        addToHistory(input, next);
      },
      setLoading,
      "Song assets generated",
    );
  }, [input, buildRequest, runGeneration]);

  const handleRegenerateAll = useCallback(async () => {
    if (!input.idea.trim()) return;
    setAssets(null);
    setOriginalAssets(null);
    await runGeneration(
      () => buildRequest("all"),
      (result) => {
        const next = result as SongAssets;
        setAssets(next);
        setOriginalAssets(next);
        setSelectedTitleIndex(0);
        addToHistory(input, next);
      },
      (b) => {
        setLoading(b);
        setRegenerating((m) => ({ ...m, all: b }));
      },
      "Regenerated all sections",
    );
  }, [input, buildRequest, runGeneration]);

  const handleRegenerateSection = useCallback(
    async (section: SectionKey) => {
      if (!input.idea.trim() || !assets) return;
      await runGeneration(
        () => buildRequest(section, { context: assets }),
        (result) => {
          // Merge partial into current assets, and update the original
          // snapshot for the keys that came back. Any keys the user has
          // hand-edited that the model *didn't* return stay as-is.
          const merged: SongAssets = { ...assets, ...result };
          setAssets(merged);
          setOriginalAssets((prev) => (prev ? { ...prev, ...result } : merged));
        },
        (b) => setRegenerating((m) => ({ ...m, [section]: b })),
        `Regenerated ${SECTION_LABELS[section]}`,
      );
    },
    [input, assets, buildRequest, runGeneration],
  );

  const handleDownload = useCallback(async () => {
    if (!assets || !selectedTitle) return;
    try {
      const blob = await buildSongZip(assets, selectedTitle);
      const safe = sanitizeFilename(selectedTitle);
      downloadBlob(blob, `${safe || "song"}.zip`);
    } catch {
      toast.error("Could not generate ZIP — please try again");
    }
  }, [assets, selectedTitle, toast]);

  // Asset editing handlers
  const handleAssetChange = useCallback(
    <K extends keyof SongAssets>(key: K, value: SongAssets[K]) => {
      setAssets((curr) => (curr ? { ...curr, [key]: value } : curr));
    },
    [],
  );

  const handleVideoPromptChange = useCallback(
    (index: number, value: string) => {
      setAssets((curr) => {
        if (!curr) return curr;
        const next: VideoPrompt[] = curr.video_prompts.map((p, i) =>
          i === index ? { ...p, prompt: value } : p,
        );
        return { ...curr, video_prompts: next };
      });
    },
    [],
  );

  // Restore a previous generation from history.
  const handleLoadHistory = useCallback(
    (entry: HistoryEntry) => {
      if (anyRegenerating) cancel();
      setInput(entry.input);
      setAssets(entry.assets);
      setOriginalAssets(entry.assets);
      setSelectedTitleIndex(0);
      toast.info(`Loaded "${truncate(entry.input.idea, 60)}"`);
    },
    [anyRegenerating, cancel, toast],
  );

  return (
    <div className="min-h-screen pb-32">
      <header className="sticky top-0 z-20 bg-bg/80 backdrop-blur border-b border-border">
        <div className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 h-14 flex items-center justify-between">
          <div className="flex items-center gap-2.5">
            <EqualizerIcon size={18} className="shrink-0" />
            <span className="text-sm font-medium text-fg-muted">
              MelodyMuse
            </span>
          </div>
          <div className="flex items-center gap-3">
            <a
              href="https://www.youtube.com/@AIWentNonsense"
              target="_blank"
              rel="noopener noreferrer"
              className="inline-flex items-center gap-1.5 text-xs text-fg-muted hover:text-fg transition-colors"
              title="AI Went Nonsense on YouTube"
            >
              <Youtube className="w-4 h-4" />
              <span className="hidden sm:inline">YouTube</span>
            </a>
            {serverOk === false && (
              <span
                className="text-xs text-rose-300 hidden sm:inline"
                title="The MelodyMuse server is unreachable. Generation will fail."
              >
                Server offline
              </span>
            )}
          </div>
        </div>
      </header>

      <main className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-6 lg:py-10">
        <div className="grid grid-cols-1 lg:grid-cols-[35%_minmax(0,1fr)] gap-8">
          <div className="lg:sticky lg:top-20 lg:self-start space-y-4">
            <InputPanel
              values={input}
              onChange={setInput}
              onSubmit={handleGenerate}
              loading={loading}
              disabled={anyRegenerating}
              cancelButton={
                <button
                  type="button"
                  onClick={cancel}
                  className="btn-secondary w-full"
                >
                  Cancel generation
                </button>
              }
            />
            <HistoryPanel onLoad={handleLoadHistory} />
          </div>

          <div>
            <ResultsPanel
              assets={assets}
              originalAssets={originalAssets}
              loading={loading}
              selectedTitleIndex={selectedTitleIndex}
              onSelectTitle={setSelectedTitleIndex}
              onChange={handleAssetChange}
              onChangeVideoPrompt={handleVideoPromptChange}
              onRegenerateAll={handleRegenerateAll}
              onRegenerateSection={handleRegenerateSection}
              onCancel={cancel}
              anyRegenerating={anyRegenerating}
              regenerating={regenerating}
            />
          </div>
        </div>
      </main>

      {assets && (
        <StickyZipBar
          title={selectedTitle}
          busy={loading}
          onDownload={handleDownload}
        />
      )}

      <Footer />
    </div>
  );
}

function truncate(s: string, n: number): string {
  const t = s.trim();
  if (t.length <= n) return t || "previous generation";
  return t.slice(0, n - 1) + "…";
}

type GenerateCall = Parameters<typeof generateSong>[0];