frontend/src/components/Word2VecPanel.tsx

import { useState, useEffect } from "react";
import { api, getErrorMessage } from "../api";
import type { W2VInitResponse } from "../types";
import StatusMessage from "./StatusMessage";
import LogViewer from "./LogViewer";
import MetricCard from "./MetricCard";

interface Props {
  onReady: (ready: boolean, info?: { vocab_size: number; sentences: number; vector_size: number }) => void;
}

export default function Word2VecPanel({ onReady }: Props) {
  const [statusChecked, setStatusChecked] = useState(false);
  const [trainResult, setTrainResult] = useState<W2VInitResponse | null>(null);

  const [vectorSize, setVectorSize] = useState(100);
  const [windowSize, setWindowSize] = useState(5);
  const [w2vEpochs, setW2vEpochs] = useState(50);
  const [showAdvanced, setShowAdvanced] = useState(false);
  const [initLoading, setInitLoading] = useState(false);
  const [error, setError] = useState("");

  useEffect(() => {
    api.w2vStatus().then(res => {
      if (res.ready) {
        onReady(true, { vocab_size: res.vocab_size!, sentences: res.sentences!, vector_size: res.vector_size! });
      }
      setStatusChecked(true);
    }).catch(() => setStatusChecked(true));
  }, []);

  async function handleTrainFromEngine() {
    setInitLoading(true); setError(""); setTrainResult(null);
    try {
      const res = await api.w2vInitFromEngine({
        vector_size: vectorSize,
        window: windowSize,
        epochs: w2vEpochs,
      });
      setTrainResult(res);
    } catch (err) {
      setError(getErrorMessage(err));
    } finally {
      setInitLoading(false);
    }
  }

  if (!statusChecked) {
    return <div className="panel"><p>Checking Word2Vec status...</p></div>;
  }

  // Training complete — show results + continue button
  if (trainResult) {
    return (
      <div>
        <div className="panel">
          <h2>Training Complete</h2>
          <div className="metric-grid">
            <MetricCard value={trainResult.vocab_size} label="Vocabulary" />
            <MetricCard value={trainResult.sentences} label="Sentences" />
            <MetricCard value={trainResult.vector_size} label="Dimensions" />
            <MetricCard value={`${trainResult.seconds}s`} label="Train Time" />
          </div>
          <StatusMessage type="ok" message="Word2Vec model trained and saved. It will persist across restarts." />
          <button className="btn btn-primary" style={{ marginTop: 12 }}
            onClick={() => onReady(true, { vocab_size: trainResult.vocab_size, sentences: trainResult.sentences, vector_size: trainResult.vector_size })}>
            Continue to Analysis
          </button>
        </div>

        <LogViewer active={false} />
      </div>
    );
  }

  // Training form
  return (
    <div>
      <div className="panel">
        <h2>Word2Vec Baseline (gensim)</h2>
        <p className="panel-desc">
          Static embeddings — one vector per word, no context awareness.
          Train on all documents loaded in the engine to use as a baseline comparison.
        </p>

        <button className="advanced-toggle" onClick={() => setShowAdvanced(!showAdvanced)}>
          {showAdvanced ? "\u25be" : "\u25b8"} Advanced Settings
        </button>

        {showAdvanced && (
          <div className="advanced-section">
            <div className="form-row">
              <div className="form-group" style={{ maxWidth: 120 }}>
                <label>Vector Size</label>
                <input type="number" value={vectorSize} onChange={e => setVectorSize(+e.target.value)} min={50} max={300} />
              </div>
              <div className="form-group" style={{ maxWidth: 120 }}>
                <label>Window</label>
                <input type="number" value={windowSize} onChange={e => setWindowSize(+e.target.value)} min={2} max={15} />
              </div>
              <div className="form-group" style={{ maxWidth: 120 }}>
                <label>Epochs</label>
                <input type="number" value={w2vEpochs} onChange={e => setW2vEpochs(+e.target.value)} min={5} max={200} />
              </div>
            </div>
          </div>
        )}

        <button className="btn btn-primary" onClick={handleTrainFromEngine}
          disabled={initLoading} style={{ marginTop: 8 }}>
          {initLoading ? <><span className="spinner" /> Training on all engine documents...</> : "Train Word2Vec"}
        </button>

        <LogViewer active={initLoading} />
      </div>

      {error && <StatusMessage type="err" message={error} />}
    </div>
  );
}