frontend/src/components/EngineSetup.tsx

import { useState } from "react";
import { api, getErrorMessage } from "../api";
import type { CorpusStats } from "../types";
import StatusMessage from "./StatusMessage";
import Select from "./Select";

interface Props {
  onStatsUpdate: (stats: CorpusStats) => void;
}

const MODELS = [
  { value: "all-MiniLM-L6-v2", label: "all-MiniLM-L6-v2 (fast, 384-dim)" },
  { value: "all-mpnet-base-v2", label: "all-mpnet-base-v2 (best quality, 768-dim)" },
  { value: "BAAI/bge-large-en-v1.5", label: "BAAI/bge-large-en-v1.5 (high accuracy, 1024-dim)" },
];

export default function EngineSetup({ onStatsUpdate }: Props) {
  const [model, setModel] = useState("all-MiniLM-L6-v2");
  const [chunkSize, setChunkSize] = useState(512);
  const [chunkOverlap, setChunkOverlap] = useState(128);
  const [batchSize, setBatchSize] = useState(64);

  const [docId, setDocId] = useState("");
  const [docText, setDocText] = useState("");

  const [showAdvanced, setShowAdvanced] = useState(false);
  const [status, setStatus] = useState<{ type: "ok" | "err" | "loading"; msg: string } | null>(null);
  const [initialized, setInitialized] = useState(false);
  const [docsAdded, setDocsAdded] = useState<string[]>([]);

  async function handleInit() {
    setStatus({ type: "loading", msg: "Loading model..." });
    try {
      const res = await api.init({
        model_name: model,
        chunk_size: chunkSize,
        chunk_overlap: chunkOverlap,
        batch_size: batchSize,
      });
      setInitialized(true);
      setDocsAdded([]);
      setStatus({ type: "ok", msg: `Model "${res.model}" loaded in ${res.load_time_seconds}s` });
    } catch (e: unknown) {
      setStatus({ type: "err", msg: getErrorMessage(e) });
    }
  }

  async function handleAddDoc() {
    if (!docId.trim() || !docText.trim()) return;
    setStatus({ type: "loading", msg: `Adding document "${docId}"...` });
    try {
      const res = await api.addDocument({ doc_id: docId, text: docText });
      setDocsAdded((prev) => [...prev, res.doc_id]);
      setStatus({ type: "ok", msg: `Added "${res.doc_id}": ${res.num_chunks} chunks` });
      setDocId("");
      setDocText("");
    } catch (e: unknown) {
      setStatus({ type: "err", msg: getErrorMessage(e) });
    }
  }

  async function handleBuildIndex() {
    setStatus({ type: "loading", msg: "Building FAISS index..." });
    try {
      const res = await api.buildIndex();
      setStatus({
        type: "ok",
        msg: `Index built: ${res.total_chunks} vectors (dim=${res.embedding_dim}) in ${res.build_time_seconds}s`,
      });
      const stats = await api.getStats();
      onStatsUpdate(stats);
    } catch (e: unknown) {
      setStatus({ type: "err", msg: getErrorMessage(e) });
    }
  }

  return (
    <div>
      {/* Step 1: Initialize engine */}
      <div className="panel">
        <h2>1. Initialize Engine</h2>
        <div className="form-row">
          <div className="form-group">
            <label>Model</label>
            <Select options={MODELS} value={model} onChange={setModel} />
          </div>
        </div>

        <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 form-group-md">
                <label>Chunk Size</label>
                <input type="number" value={chunkSize} onChange={(e) => setChunkSize(+e.target.value)} />
              </div>
              <div className="form-group form-group-md">
                <label>Overlap</label>
                <input type="number" value={chunkOverlap} onChange={(e) => setChunkOverlap(+e.target.value)} />
              </div>
              <div className="form-group form-group-md">
                <label>Batch Size</label>
                <input type="number" value={batchSize} onChange={(e) => setBatchSize(+e.target.value)} />
              </div>
            </div>
          </div>
        )}

        <button className="btn btn-primary" onClick={handleInit} style={{ marginTop: 8 }}>
          Initialize
        </button>
      </div>

      {/* Step 2: Add documents */}
      <div className="panel">
        <h2>2. Add Documents</h2>
        {docsAdded.length > 0 && (
          <div style={{ marginBottom: 12 }}>
            {docsAdded.map((id) => (
              <span key={id} className="tag">{id}</span>
            ))}
          </div>
        )}
        <div className="form-row">
          <div className="form-group form-group-lg">
            <label>Document ID</label>
            <input
              value={docId}
              onChange={(e) => setDocId(e.target.value)}
              placeholder="e.g. chapter_1"
              disabled={!initialized}
            />
          </div>
        </div>
        <div className="form-group mb-2">
          <label>Document Text</label>
          <textarea
            value={docText}
            onChange={(e) => setDocText(e.target.value)}
            placeholder="Paste your document text here..."
            rows={8}
            disabled={!initialized}
          />
        </div>
        <button className="btn btn-primary" onClick={handleAddDoc} disabled={!initialized || !docId || !docText}>
          Add Document
        </button>
      </div>

      {/* Step 3: Build index */}
      <div className="panel">
        <h2>3. Build Index</h2>
        <p className="panel-desc">
          Embeds all chunks and builds a FAISS index for fast similarity search.
          This must be done after adding all documents.
        </p>
        <button
          className="btn btn-primary"
          onClick={handleBuildIndex}
          disabled={!initialized || docsAdded.length === 0}
        >
          Build Index
        </button>
      </div>

      {status && <StatusMessage type={status.type} message={status.msg} />}
    </div>
  );
}