import { useEffect, useRef, useState } from "react";
import "./index.css";

import { AuthForms } from "./components/AuthForms";
import { FileDropzone } from "./components/FileDropzone";
import { LiveControls, type LiveStatus } from "./components/LiveControls";
import { ModelLoadProgress } from "./components/ModelLoadProgress";
import { ModelPicker } from "./components/ModelPicker";
import { TranscriptView } from "./components/TranscriptView";
import { TranscriptionOptions } from "./components/TranscriptionOptions";
import { decodeAudioFile } from "./lib/audio";
import { LiveTranscriber, type LiveTranscribeResult } from "./lib/live";
import { startMicCapture, type MicCapture } from "./lib/mic";
import { getDownloadedModels } from "./lib/modelCache";
import { registerServiceWorker } from "./pwa/register";
import {
  AUTO_LANGUAGE,
  DEFAULT_DEVICE_PREFERENCE,
  DEFAULT_MODEL_ID,
  DEFAULT_TASK,
  MODELS,
  getModel,
} from "./lib/models";
import {
  applyLoadProgress,
  initialLoadProgress,
  isWorkerToUiMessage,
  type DevicePreference,
  type InferenceDevice,
  type TranscriptChunk,
  type UiToWorkerMessage,
  type WhisperTask,
} from "./worker/protocol";

interface TranscriptResult {
  text: string;
  chunks: TranscriptChunk[];
}

const MODEL_IDS = MODELS.map((model) => model.id);

export function App() {
  const workerRef = useRef<Worker | null>(null);

  // Live-mode session (M7, D-016). The mic and transcriber live outside React
  // state; only display text goes through setState. Segment passes are
  // correlated with worker responses via requestId.
  const liveRef = useRef<{ mic: MicCapture; transcriber: LiveTranscriber } | null>(
    null,
  );
  const pendingLive = useRef(
    new Map<
      string,
      {
        resolve: (result: LiveTranscribeResult) => void;
        reject: (error: Error) => void;
      }
    >(),
  );
  const nextLiveRequest = useRef(0);

  const [busy, setBusy] = useState(false);
  const [decoding, setDecoding] = useState(false);
  const [loadStarted, setLoadStarted] = useState(false);
  const [loadProgress, setLoadProgress] = useState(initialLoadProgress);
  const [partialText, setPartialText] = useState("");
  const [result, setResult] = useState<TranscriptResult | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [fileName, setFileName] = useState<string | null>(null);

  const [liveStatus, setLiveStatus] = useState<LiveStatus>("idle");
  const [liveCommitted, setLiveCommitted] = useState("");
  const [liveTentative, setLiveTentative] = useState("");

  const [modelId, setModelId] = useState(DEFAULT_MODEL_ID);
  const [devicePreference, setDevicePreference] = useState<DevicePreference>(
    DEFAULT_DEVICE_PREFERENCE,
  );
  const [activeDevice, setActiveDevice] = useState<InferenceDevice | null>(null);
  const [language, setLanguage] = useState(AUTO_LANGUAGE);
  const [task, setTask] = useState<WhisperTask>(DEFAULT_TASK);
  const [downloaded, setDownloaded] = useState<ReadonlySet<string>>(new Set());

  // PWA state (M3, D-007): is a service worker controlling this page (app
  // shell available offline), and is an updated version waiting?
  const [swControlled, setSwControlled] = useState(false);
  const [applyUpdate, setApplyUpdate] = useState<(() => void) | null>(null);

  function refreshDownloaded() {
    // Downloaded badges (D-005): query Cache Storage for cached weights.
    getDownloadedModels(MODEL_IDS).then(setDownloaded, () => {});
  }

  // Terminate the worker when the app unmounts (also keeps HMR tidy).
  useEffect(() => {
    refreshDownloaded();
    // Registration is idempotent, so the StrictMode double-invoke is fine.
    registerServiceWorker({
      onControlled: () => setSwControlled(true),
      onUpdateReady: (apply) => setApplyUpdate(() => apply),
    });
    return () => {
      void liveRef.current?.mic.stop().catch(() => {});
      liveRef.current = null;
      workerRef.current?.terminate();
      workerRef.current = null;
    };
  }, []);

  function takeLiveRequest(requestId: string | undefined) {
    if (requestId === undefined) return undefined;
    const pending = pendingLive.current.get(requestId);
    if (pending) pendingLive.current.delete(requestId);
    return pending;
  }

  function ensureWorker(): Worker {
    if (workerRef.current) return workerRef.current;
    // Inference always runs in a dedicated Web Worker (D-003), never on the
    // main thread (Architecture.md §5). The worker script is bundled by the
    // server at startup and served at a stable URL (see src/index.ts) because
    // Bun's bundler does not follow `new Worker(new URL(...))` graphs yet.
    const worker = new Worker("inference.worker.js", { type: "module" });
    worker.addEventListener("message", (event: MessageEvent<unknown>) => {
      const message = event.data;
      if (!isWorkerToUiMessage(message)) return;
      switch (message.type) {
        case "load-progress":
          setLoadProgress((state) => applyLoadProgress(state, message));
          if (message.status === "ready") {
            // Surface the backend the worker actually selected (D-004) and
            // refresh the downloaded badges now that weights are cached.
            setActiveDevice(message.device);
            refreshDownloaded();
          }
          break;
        case "partial":
          // Live segment passes (requestId set) resolve wholesale via
          // "result"; their token stream is not rendered incrementally.
          if (message.requestId === undefined) {
            setPartialText((text) => text + message.text);
          }
          break;
        case "result": {
          const pending = takeLiveRequest(message.requestId);
          if (pending) {
            pending.resolve({ text: message.text, chunks: message.chunks });
          } else if (message.requestId === undefined) {
            setResult({ text: message.text, chunks: message.chunks });
            setBusy(false);
          }
          break;
        }
        case "error": {
          const pending = takeLiveRequest(message.requestId);
          if (pending) {
            pending.reject(new Error(message.message));
          } else {
            setError(message.message);
            setBusy(false);
          }
          break;
        }
      }
    });
    worker.addEventListener("error", (event) => {
      const failure = new Error(
        event.message || "The inference worker crashed unexpectedly.",
      );
      // Unblock any live passes so a live session can drain and stop.
      for (const pending of pendingLive.current.values()) pending.reject(failure);
      pendingLive.current.clear();
      setError(failure.message);
      setBusy(false);
    });
    workerRef.current = worker;
    return worker;
  }

  function postToWorker(worker: Worker, message: UiToWorkerMessage, transfer?: Transferable[]) {
    worker.postMessage(message, transfer ?? []);
  }

  /**
   * Post a load request and reset per-load UI state so the progress panel
   * restarts cleanly on model/device switches. Idempotent in the worker for
   * an unchanged model+device; a change disposes the old pipeline first.
   */
  function startLoad(nextModelId: string, preference: DevicePreference): Worker {
    const worker = ensureWorker();
    setLoadProgress(initialLoadProgress);
    setLoadStarted(true);
    setError(null);
    postToWorker(worker, { type: "load", modelId: nextModelId, device: preference });
    return worker;
  }

  function handleModelSelect(nextModelId: string) {
    if (busy || liveStatus !== "idle" || nextModelId === modelId) return;
    setModelId(nextModelId);
    setActiveDevice(null);
    setPartialText("");
    setResult(null);
    startLoad(nextModelId, devicePreference);
  }

  function handleDevicePreference(preference: DevicePreference) {
    if (busy || liveStatus !== "idle" || preference === devicePreference) return;
    setDevicePreference(preference);
    setActiveDevice(null);
    // Only reload eagerly if a load was already in motion; otherwise the next
    // file drop picks up the new preference.
    if (loadStarted) startLoad(modelId, preference);
  }

  async function handleFile(file: File) {
    if (busy || liveStatus !== "idle") return;
    setBusy(true);
    setError(null);
    setPartialText("");
    setResult(null);
    setFileName(file.name);

    // Kick off the (idempotent) model load immediately so the download runs
    // while we decode the audio locally.
    const worker = ensureWorker();
    if (!loadStarted) {
      startLoad(modelId, devicePreference);
    } else {
      postToWorker(worker, { type: "load", modelId, device: devicePreference });
    }

    const multilingual = getModel(modelId)?.multilingual ?? false;

    setDecoding(true);
    try {
      // Decoded entirely on-device via Web Audio (D-012); the audio never
      // touches the network (Architecture.md §5).
      const audio = await decodeAudioFile(file);
      // Transfer the buffer to the worker — zero copy. Language/task only
      // apply to multilingual models; `.en` variants take neither.
      postToWorker(
        worker,
        multilingual
          ? { type: "transcribe", audio, language, task }
          : { type: "transcribe", audio },
        [audio.buffer as ArrayBuffer],
      );
    } catch (decodeError) {
      setError(decodeError instanceof Error ? decodeError.message : String(decodeError));
      setBusy(false);
    } finally {
      setDecoding(false);
    }
  }

  // -------------------------------------------------------------------------
  // Live microphone transcription (M7, D-016)
  // -------------------------------------------------------------------------

  /** One inference pass over a live segment, correlated via requestId. */
  function transcribeLiveSegment(
    worker: Worker,
    audio: Float32Array,
    options: { language?: string; task?: WhisperTask },
  ): Promise<LiveTranscribeResult> {
    const requestId = `live-${++nextLiveRequest.current}`;
    return new Promise((resolve, reject) => {
      pendingLive.current.set(requestId, { resolve, reject });
      // The controller hands over ownership, so the buffer can transfer.
      postToWorker(worker, { type: "transcribe", audio, requestId, ...options }, [
        audio.buffer as ArrayBuffer,
      ]);
    });
  }

  async function startLive() {
    if (busy || liveStatus !== "idle") return;
    setError(null);
    setResult(null);
    setPartialText("");
    setFileName(null);
    setLiveCommitted("");
    setLiveTentative("");
    setLiveStatus("starting");

    // Idempotent model load, like the file flow — a pending download runs
    // while the mic opens; segment passes queue in the worker until ready.
    const worker = ensureWorker();
    if (!loadStarted) {
      startLoad(modelId, devicePreference);
    } else {
      postToWorker(worker, { type: "load", modelId, device: devicePreference });
    }

    // Language/task are fixed for the session; the options are disabled while
    // live. `.en` models take neither (same rule as the file flow).
    const options = (getModel(modelId)?.multilingual ?? false) ? { language, task } : {};
    const transcriber = new LiveTranscriber({
      tentativeIntervalMs: 2500,
      minTentativeSeconds: 2.0,
      segmenter: { commitSilenceSeconds: 2.5, maxSegmentSeconds: 29 },
      transcribe: (audio) => transcribeLiveSegment(worker, audio, options),
      onUpdate: (update) => {
        setLiveCommitted(update.committedText);
        setLiveTentative(update.tentativeText);
        if (update.error) setError(update.error);
      },
    });

    try {
      const mic = await startMicCapture((chunk) => transcriber.push(chunk));
      liveRef.current = { mic, transcriber };
      setLiveStatus("listening");
    } catch (micError) {
      setLiveStatus("idle");
      setError(
        micError instanceof DOMException && micError.name === "NotAllowedError"
          ? "Microphone access was denied. Allow it in your browser's site settings and try again."
          : micError instanceof Error
            ? micError.message
            : String(micError),
      );
    }
  }

  async function stopLive() {
    const session = liveRef.current;
    if (!session || liveStatus !== "listening") return;
    setLiveStatus("stopping");
    // Stop feeding audio first, then drain the remaining segment passes.
    await session.mic.stop();
    try {
      const final = await session.transcriber.stop();
      setResult(final);
      setLiveCommitted("");
      setLiveTentative("");
    } finally {
      liveRef.current = null;
      setLiveStatus("idle");
    }
  }

  const model = getModel(modelId);
  const modelReady = loadProgress.ready;
  const showModelProgress = loadStarted && !modelReady && !error;
  const transcribing = busy && !decoding && modelReady;
  const liveActive = liveStatus !== "idle";

  return (
    <main className="app">
      <header className="app-header">
        <div>
          <h1>AU Whisper Web</h1>
          <p className="tagline">
            Speech-to-text that runs entirely in your browser. Your audio never
            leaves this device.
          </p>
        </div>
        {/* Session UI (M4). Anonymous-first: logged out, this is only an
            unobtrusive "Sign in" button — everything else is unchanged.
            Renders nothing until the /me probe confirms the server has auth
            enabled (runtime flag, D-015). */}
        <AuthForms />
      </header>

      <ModelPicker
        models={MODELS}
        selectedId={modelId}
        onSelect={handleModelSelect}
        devicePreference={devicePreference}
        onDevicePreferenceChange={handleDevicePreference}
        activeDevice={activeDevice}
        downloaded={downloaded}
        offlineReady={swControlled}
        disabled={busy || liveActive}
      />

      {model?.multilingual && (
        <TranscriptionOptions
          language={language}
          onLanguageChange={setLanguage}
          task={task}
          onTaskChange={setTask}
          disabled={busy || liveActive}
        />
      )}

      <FileDropzone onFile={handleFile} disabled={busy || liveActive} />

      <LiveControls
        status={liveStatus}
        onStart={() => void startLive()}
        onStop={() => void stopLive()}
        disabled={busy}
      />

      {fileName && (
        <p className="file-name">
          File: <strong>{fileName}</strong>
          {decoding && " — decoding audio…"}
        </p>
      )}

      {error && (
        <p className="error" role="alert">
          {error}
        </p>
      )}

      {showModelProgress && (
        <ModelLoadProgress modelId={modelId} progress={loadProgress} />
      )}

      <TranscriptView
        partialText={liveActive ? liveCommitted : partialText}
        tentativeText={liveActive ? liveTentative : ""}
        result={result}
        transcribing={liveActive || transcribing}
        activeLabel={
          liveStatus === "stopping"
            ? "Finishing up…"
            : liveActive
              ? "Listening…"
              : undefined
        }
      />

      {applyUpdate && (
        <div className="update-toast" role="status">
          <span>A new version of Whisper Web is available.</span>
          <button type="button" onClick={applyUpdate}>
            Reload
          </button>
          <button
            type="button"
            className="update-toast-dismiss"
            onClick={() => setApplyUpdate(null)}
          >
            Later
          </button>
        </div>
      )}

      <footer className="app-footer">
        <p>
          Model: <code>{modelId}</code> · Engine:{" "}
          {activeDevice === "webgpu"
            ? "WebGPU (on-device)"
            : activeDevice === "webnn"
              ? "WebNN (on-device)"
              : activeDevice === "wasm"
                ? "WASM (on-device)"
                : "on-device — backend selected on first load"}
        </p>
      </footer>
    </main>
  );
}

export default App;