import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  api, assetFileURL, cancelJob, connectWS, deleteAsset, deleteJob, getMe, getSettings,
  listAssets, listJobs, listOperators, logout, retryJob, runJob, uploadFile,
} from "./api";
import Viewer from "./Viewer";
import SplatViewer from "./SplatViewer";
import Settings from "./Settings";
import Login from "./Login";
import Dashboard from "./Dashboard";

const MODEL_EXTS = ["glb", "gltf", "obj", "fbx", "ply", "stl"];
const fmtSize = (n) =>
  n > 1e9 ? (n / 1e9).toFixed(1) + " GB" : n > 1e6 ? (n / 1e6).toFixed(1) + " MB"
  : n > 1e3 ? (n / 1e3).toFixed(0) + " KB" : n + " B";
const KIND_ICON = { video: "🎬", image: "🖼", model: "🧊", splat: "✨", frames: "🎞", colmap_dataset: "📐", folder: "📁", other: "📄" };
const isViewable = (a) => a && (a.kind === "splat" || MODEL_EXTS.includes(a.name.split(".").pop().toLowerCase()));

function ParamForm({ schema, values, onChange }) {
  const props = schema?.properties || {};
  if (!Object.keys(props).length) return <p className="dim">No parameters.</p>;
  return (
    <div className="params">
      {Object.entries(props).map(([key, def]) => {
        const val = values[key] ?? def.default ?? "";
        const set = (v) => onChange({ ...values, [key]: v });
        return (
          <label key={key} title={def.description || ""}>
            <span>{key}</span>
            {def.enum ? (
              <select value={val} onChange={(e) => set(maybeNum(def, e.target.value))}>
                {def.enum.map((o) => <option key={o} value={o}>{o}</option>)}
              </select>
            ) : def.type === "boolean" ? (
              <input type="checkbox" checked={!!val} onChange={(e) => set(e.target.checked)} />
            ) : def.type === "number" || def.type === "integer" ? (
              <input type="number" value={val} step={def.type === "integer" ? 1 : "any"}
                min={def.minimum} max={def.maximum}
                onChange={(e) => set(e.target.value === "" ? "" : (def.type === "integer" ? parseInt(e.target.value) : parseFloat(e.target.value)))} />
            ) : (
              <input type="text" value={val} onChange={(e) => set(e.target.value)} />
            )}
          </label>
        );
      })}
    </div>
  );
}
const maybeNum = (def, v) => (def.type === "integer" ? parseInt(v) : def.type === "number" ? parseFloat(v) : v);

function AssetPreview({ asset }) {
  const [folder, setFolder] = useState(null);
  const ext = asset ? asset.name.split(".").pop().toLowerCase() : "";
  useEffect(() => {
    setFolder(null);
    if (asset && (asset.kind === "frames" || asset.kind === "folder" || asset.kind === "colmap_dataset")) {
      api(assetFileURL(asset.id)).then(setFolder).catch(() => {});
    }
  }, [asset?.id]);

  if (!asset) return <div className="preview empty">Select an asset — or drop / paste files anywhere</div>;
  if (asset.kind === "splat") return <SplatViewer url={assetFileURL(asset.id)} />;
  if (MODEL_EXTS.includes(ext)) return <Viewer url={assetFileURL(asset.id)} ext={ext} />;
  if (asset.kind === "image") return <div className="preview"><img src={assetFileURL(asset.id)} alt={asset.name} /></div>;
  if (asset.kind === "video") return <div className="preview"><video src={assetFileURL(asset.id)} controls /></div>;
  if (folder?.files) {
    const imgs = folder.files.filter((f) => /\.(jpg|jpeg|png)$/i.test(f));
    return (
      <div className="preview grid">
        {imgs.slice(0, 60).map((f) => <img key={f} src={assetFileURL(asset.id, f)} loading="lazy" alt={f} />)}
        {imgs.length > 60 && <span className="dim">+{imgs.length - 60} more</span>}
        {!imgs.length && <span className="dim">{folder.files.length} files (no image preview)</span>}
      </div>
    );
  }
  if (ext === "json") return <JsonPreview url={assetFileURL(asset.id)} />;
  return <div className="preview empty">No preview for .{ext} — download to use it</div>;
}

function JsonPreview({ url }) {
  const [text, setText] = useState("…");
  useEffect(() => { fetch(url).then((r) => r.text()).then((t) => setText(t.slice(0, 20000))); }, [url]);
  return <pre className="preview json">{text}</pre>;
}

function CompareGrid({ assets }) {
  return (
    <div className="compare-grid">
      {assets.map((a) => (
        <div key={a.id} className="compare-cell">
          <div className="compare-cap">{a.name} · {fmtSize(a.size)}{a.meta?.tool ? ` · ${a.meta.tool}` : ""}</div>
          {a.kind === "splat"
            ? <SplatViewer url={assetFileURL(a.id)} />
            : <Viewer url={assetFileURL(a.id)} ext={a.name.split(".").pop().toLowerCase()} />}
        </div>
      ))}
    </div>
  );
}

export default function App() {
  const [user, setUser] = useState(undefined);  // undefined=loading, null=logged out
  const [view, setView] = useState("workspace"); // workspace | dashboard
  const [assets, setAssets] = useState([]);
  const [operators, setOperators] = useState([]);
  const [jobs, setJobs] = useState([]);
  const [settings, setSettings] = useState({ _env_set: [] });
  const [selected, setSelected] = useState(null);
  const [opId, setOpId] = useState(null);
  const [params, setParams] = useState({});
  const [openJob, setOpenJob] = useState(null);
  const [showSettings, setShowSettings] = useState(false);
  const [compareMode, setCompareMode] = useState(false);
  const [compareSel, setCompareSel] = useState([]);
  const fileInput = useRef(null);
  const wsRef = useRef(null);

  const refreshAssets = useCallback(() => listAssets().then(setAssets).catch(() => {}), []);
  const refreshJobs = useCallback(() => listJobs().then(setJobs).catch(() => {}), []);
  const refreshSettings = useCallback(() => getSettings().then(setSettings).catch(() => {}), []);

  // auth bootstrap
  useEffect(() => {
    getMe().then(setUser).catch(() => setUser(null));
    const onUnauth = () => setUser(null);
    window.addEventListener("mb-unauth", onUnauth);
    return () => window.removeEventListener("mb-unauth", onUnauth);
  }, []);

  // load workspace data + connect the live socket once, after login
  useEffect(() => {
    if (!user) return;
    refreshAssets(); refreshJobs();
    if (user.is_owner) refreshSettings();
    listOperators().then(setOperators).catch(() => {});
    if (!wsRef.current) {
      wsRef.current = connectWS((msg) => {
        if (msg.type === "job") {
          setJobs((js) => {
            const i = js.findIndex((j) => j.id === msg.job.id);
            if (i === -1) return [msg.job, ...js];
            const copy = [...js]; copy[i] = msg.job; return copy;
          });
        }
        if (msg.type === "assets_changed") refreshAssets();
      });
    }
  }, [user?.username]);

  const doLogout = async () => { try { await logout(); } catch { /* ignore */ } window.location.reload(); };

  useEffect(() => {
    const drop = async (e) => { e.preventDefault(); for (const f of e.dataTransfer?.files || []) await uploadFile(f); refreshAssets(); };
    const paste = async (e) => { for (const it of e.clipboardData?.items || []) { const f = it.getAsFile?.(); if (f) await uploadFile(f); } refreshAssets(); };
    const prevent = (e) => e.preventDefault();
    window.addEventListener("drop", drop); window.addEventListener("dragover", prevent); window.addEventListener("paste", paste);
    return () => { window.removeEventListener("drop", drop); window.removeEventListener("dragover", prevent); window.removeEventListener("paste", paste); };
  }, [refreshAssets]);

  const selectedAsset = assets.find((a) => a.id === selected) || null;
  const envSet = settings._env_set || [];
  // no-input operators (accepts: []) are always available; others filter by asset kind
  const usableOps = useMemo(() =>
    operators.filter((op) => !op.accepts.length || !selectedAsset || op.accepts.includes(selectedAsset.kind)),
    [operators, selectedAsset]);
  const opGroups = useMemo(() => {
    const groups = {};
    for (const op of usableOps) (groups[op.category || "other"] ||= []).push(op);
    return groups;
  }, [usableOps]);
  const activeOp = operators.find((o) => o.id === opId);
  const needsInput = (activeOp?.accepts?.length ?? 0) > 0;
  const gatedEnv = (activeOp?.requires_env || []).filter((e) => !envSet.includes(e));
  const gated = gatedEnv.length > 0;

  useEffect(() => {
    if (activeOp) {
      const defaults = {};
      for (const [k, d] of Object.entries(activeOp.params_schema?.properties || {}))
        if (d.default !== undefined) defaults[k] = d.default;
      setParams(defaults);
    }
  }, [opId]);
  useEffect(() => {
    if (usableOps.length && !usableOps.find((o) => o.id === opId)) setOpId(usableOps[0].id);
  }, [usableOps]);

  const launch = async () => {
    if (!activeOp || gated || (needsInput && !selectedAsset)) return;
    await runJob(activeOp.id, needsInput ? selectedAsset.id : null, params);
    refreshJobs();
  };

  const clickAsset = (a) => {
    if (compareMode && isViewable(a)) {
      setCompareSel((s) => s.includes(a.id) ? s.filter((x) => x !== a.id) : [...s, a.id]);
    } else {
      setSelected(a.id);
    }
  };
  const compareAssets = compareSel.map((id) => assets.find((a) => a.id === id)).filter(Boolean);

  if (user === undefined) return <div className="login-screen"><p className="dim">…</p></div>;
  if (user === null) return <Login onLogin={setUser} />;

  return (
    <div className="layout">
      <header>
        <h1>MODEL<span>BEAST</span></h1>
        <span className="dim">M3 Ultra · 256GB · local-first 3D factory</span>
        <div className="spacer" />
        <button className={view === "dashboard" ? "active" : ""}
          onClick={() => setView(view === "dashboard" ? "workspace" : "dashboard")}>📊 Dashboard</button>
        <button className={compareMode ? "active" : ""} onClick={() => { setCompareMode(!compareMode); setCompareSel([]); setView("workspace"); }}>
          ⊞ Compare{compareSel.length ? ` (${compareSel.length})` : ""}
        </button>
        {user.is_owner && <button onClick={() => setShowSettings(true)}>⚙ Settings</button>}
        <button onClick={() => fileInput.current.click()}>+ Add files</button>
        <span className="user-chip" title={user.is_owner ? "owner" : "guest — local models only"}>
          {user.username}{user.is_owner ? " · owner" : ""}
        </span>
        <button onClick={doLogout} title="log out">⏻</button>
        <input ref={fileInput} type="file" multiple hidden
          onChange={async (e) => { for (const f of e.target.files) await uploadFile(f); e.target.value = ""; refreshAssets(); }} />
      </header>

      <aside className="assets">
        <h2>Assets</h2>
        {assets.map((a) => {
          const checked = compareSel.includes(a.id);
          return (
            <div key={a.id}
              className={"asset" + (a.id === selected && !compareMode ? " sel" : "") + (checked ? " checked" : "")}
              onClick={() => clickAsset(a)}>
              {compareMode && <input type="checkbox" checked={checked} readOnly disabled={!isViewable(a)} />}
              <span className="icon">{KIND_ICON[a.kind] || "📄"}</span>
              <div className="meta">
                <div className="name">{a.name}</div>
                <div className="dim">{a.kind} · {fmtSize(a.size)}</div>
              </div>
              <a className="dl" href={assetFileURL(a.id)} download onClick={(e) => e.stopPropagation()} title="download">↓</a>
              <button className="x" title="delete" onClick={(e) => { e.stopPropagation(); deleteAsset(a.id).then(refreshAssets); if (selected === a.id) setSelected(null); }}>×</button>
            </div>
          );
        })}
        {!assets.length && <p className="dim">Drop videos, images, or 3D files anywhere.</p>}
      </aside>

      <main>
        {view === "dashboard"
          ? <Dashboard jobs={jobs} />
          : compareMode && compareAssets.length
            ? <CompareGrid assets={compareAssets} />
            : <AssetPreview asset={selectedAsset} />}
      </main>

      <aside className="workbench">
        <h2>Run</h2>
        <select value={opId || ""} onChange={(e) => setOpId(e.target.value)}>
          {Object.entries(opGroups).map(([cat, ops]) => (
            <optgroup key={cat} label={cat}>
              {ops.map((op) => <option key={op.id} value={op.id}>{op.name}</option>)}
            </optgroup>
          ))}
        </select>
        {activeOp && <p className="dim">{activeOp.description}</p>}
        {gated && <p className="warn">🔒 Needs {gatedEnv.join(", ")} — add it in Settings.</p>}
        {activeOp && !needsInput && <p className="dim">No input asset needed — just set the parameters and run.</p>}
        {activeOp && <ParamForm schema={activeOp.params_schema} values={params} onChange={setParams} />}
        <button className="go" disabled={!activeOp || gated || (needsInput && !selectedAsset)} onClick={launch}>
          ▶ Run {activeOp?.name || ""}
        </button>
        {!user.is_owner && <p className="dim guest-note">local models only — everything you run is generated on this machine</p>}

        <h2>Jobs</h2>
        <div className="jobs">
          {jobs.map((j) => (
            <div key={j.id} className={"job " + j.status}>
              <div className="job-row" onClick={() => setOpenJob(openJob === j.id ? null : j.id)}>
                <span className="status">{j.status === "running" ? "⏳" : j.status === "done" ? "✅" : j.status === "error" ? "❌" : j.status === "cancelled" ? "⛔" : "🕐"}</span>
                <span className="jop">{j.operator}</span>
                <span className="dim jid">{j.username ? `${j.username}` : j.id}</span>
                <span className="job-actions" onClick={(e) => e.stopPropagation()}>
                  {j.status === "running" && <button onClick={() => cancelJob(j.id)}>cancel</button>}
                  {(j.status === "error" || j.status === "done" || j.status === "cancelled") && <button onClick={() => retryJob(j.id).then(refreshJobs)}>retry</button>}
                  {j.status !== "running" && <button onClick={() => deleteJob(j.id).then(refreshJobs)}>×</button>}
                </span>
              </div>
              {openJob === j.id && (
                <pre className="log">{j.error ? `ERROR: ${j.error}\n\n` : ""}{j.log || "(no output yet)"}</pre>
              )}
            </div>
          ))}
          {!jobs.length && <p className="dim">No jobs yet.</p>}
        </div>
      </aside>

      {showSettings && user.is_owner && <Settings onClose={() => setShowSettings(false)} onSaved={setSettings} />}
    </div>
  );
}