frontend/src/components/Word2VecTools.tsx

import { useState } from "react";
import { api, getErrorMessage } from "../api";
import type { W2VQueryResult, W2VSimilarWord, CompareResponse } from "../types";
import { scoreColor } from "../utils/colors";
import ScoreBar from "./ScoreBar";
import StatusMessage from "./StatusMessage";
import DocumentViewer from "./DocumentViewer";

export default function Word2VecTools() {
  const [error, setError] = useState("");

  // Similar words
  const [simWord, setSimWord] = useState("");
  const [simTopK, setSimTopK] = useState(10);
  const [simResults, setSimResults] = useState<W2VSimilarWord[]>([]);
  const [simLoading, setSimLoading] = useState(false);

  // Compare
  const [compTextA, setCompTextA] = useState("");
  const [compTextB, setCompTextB] = useState("");
  const [compResult, setCompResult] = useState<CompareResponse | null>(null);
  const [compLoading, setCompLoading] = useState(false);

  // Search
  const [queryText, setQueryText] = useState("");
  const [queryTopK, setQueryTopK] = useState(5);
  const [queryResults, setQueryResults] = useState<W2VQueryResult[]>([]);
  const [queryLoading, setQueryLoading] = useState(false);

  async function handleSimilarWords() {
    setSimLoading(true); setError("");
    try {
      const res = await api.w2vSimilarWords({ word: simWord, top_k: simTopK });
      setSimResults(res.similar);
    } catch (err) {
      setError(getErrorMessage(err));
    } finally {
      setSimLoading(false);
    }
  }

  async function handleCompare() {
    setCompLoading(true); setError("");
    try {
      const res = await api.w2vCompare({ text_a: compTextA, text_b: compTextB });
      setCompResult(res);
    } catch (err) {
      setError(getErrorMessage(err));
    } finally {
      setCompLoading(false);
    }
  }

  async function handleQuery() {
    setQueryLoading(true); setError("");
    try {
      const res = await api.w2vQuery({ text: queryText, top_k: queryTopK });
      setQueryResults(res.results);
    } catch (err) {
      setError(getErrorMessage(err));
    } finally {
      setQueryLoading(false);
    }
  }

  return (
    <div>
      {error && <StatusMessage type="err" message={error} />}

      <div style={{ display: "grid", gridTemplateColumns: "1fr 1fr", gap: 16 }}>
        {/* Similar Words */}
        <div className="panel">
          <h3 style={{ marginTop: 0 }}>Similar Words</h3>
          <p className="panel-desc">
            Find words that appear in similar contexts using Word2Vec static embeddings.
          </p>
          <div className="form-row">
            <div className="form-group">
              <label>Word</label>
              <input value={simWord} onChange={e => setSimWord(e.target.value)}
                onKeyDown={e => e.key === "Enter" && handleSimilarWords()}
                placeholder="e.g. pizza" />
            </div>
            <div className="form-group form-group-sm">
              <label>Top K</label>
              <input type="number" value={simTopK} onChange={e => setSimTopK(+e.target.value)}
                min={1} max={50} style={{ width: 60 }} />
            </div>
            <div className="form-group form-group-sm">
              <label>&nbsp;</label>
              <button className="btn btn-primary" onClick={handleSimilarWords}
                disabled={simLoading || !simWord.trim()}>
                {simLoading ? "..." : "Find"}
              </button>
            </div>
          </div>

          {simResults.length > 0 && (
            <table className="data-table" style={{ marginTop: 8 }}>
              <thead>
                <tr><th>Word</th><th>Similarity</th></tr>
              </thead>
              <tbody>
                {simResults.map((r, i) => (
                  <tr key={i}>
                    <td style={{ fontWeight: 600 }}>{r.word}</td>
                    <td><ScoreBar score={r.score} /></td>
                  </tr>
                ))}
              </tbody>
            </table>
          )}
        </div>

        {/* Compare Texts */}
        <div className="panel">
          <h3 style={{ marginTop: 0 }}>Compare Texts</h3>
          <p className="panel-desc">
            Sentence similarity via averaged word vectors.
          </p>
          <div className="form-group" style={{ marginBottom: 8 }}>
            <label>Text A</label>
            <input value={compTextA} onChange={e => setCompTextA(e.target.value)}
              placeholder="pizza gives me homework" />
          </div>
          <div className="form-group" style={{ marginBottom: 8 }}>
            <label>Text B</label>
            <input value={compTextB} onChange={e => setCompTextB(e.target.value)}
              placeholder="school gives me homework" />
          </div>
          <button className="btn btn-primary" onClick={handleCompare}
            disabled={compLoading || !compTextA.trim() || !compTextB.trim()}>
            {compLoading ? "..." : "Compare"}
          </button>

          {compResult && (
            <div className="similarity-gauge" style={{ marginTop: 12 }}>
              <div className="similarity-value"
                style={{ color: scoreColor(compResult.similarity) }}>
                {compResult.similarity.toFixed(4)}
              </div>
              <div className="similarity-label">Word2Vec Cosine Similarity</div>
            </div>
          )}
        </div>
      </div>

      {/* Semantic Search — full width */}
      <div className="panel">
        <h3 style={{ marginTop: 0 }}>Semantic Search</h3>
        <p className="panel-desc">
          Search your corpus using averaged Word2Vec vectors.
        </p>
        <div className="form-row">
          <div className="form-group" style={{ flex: 1 }}>
            <label>Query</label>
            <input value={queryText} onChange={e => setQueryText(e.target.value)}
              onKeyDown={e => e.key === "Enter" && handleQuery()}
              placeholder="a place where children learn" />
          </div>
          <div className="form-group form-group-sm">
            <label>Top K</label>
            <input type="number" value={queryTopK} onChange={e => setQueryTopK(+e.target.value)}
              min={1} max={20} style={{ width: 60 }} />
          </div>
          <div className="form-group form-group-sm">
            <label>&nbsp;</label>
            <button className="btn btn-primary" onClick={handleQuery}
              disabled={queryLoading || !queryText.trim()}>
              {queryLoading ? "Searching..." : "Search"}
            </button>
          </div>
        </div>

        {queryResults.length > 0 && (
          <div style={{ marginTop: 8 }}>
            {queryResults.map((r, i) => (
              <DocumentViewer key={i} docId={r.doc_id}>
                <div className="result-card" style={{ cursor: "pointer" }}>
                  <div className="result-header">
                    <span>#{r.rank} <span className="tag">{r.doc_id}</span></span>
                    <ScoreBar score={r.score} />
                  </div>
                  <div className="result-text">{r.text}</div>
                </div>
              </DocumentViewer>
            ))}
          </div>
        )}
      </div>
    </div>
  );
}