frontend/src/App.tsx

import { useState, useEffect, Fragment } from "react";
import type { CorpusStats } from "./types";
import { api, checkConnection } from "./api";
import TrainingPanel from "./components/TrainingPanel";
import EngineSetup from "./components/EngineSetup";
import SemanticSearch from "./components/SemanticSearch";
import TextCompare from "./components/TextCompare";
import KeywordAnalysis from "./components/KeywordAnalysis";
import KeywordMatcher from "./components/KeywordMatcher";
import BatchAnalysis from "./components/BatchAnalysis";
import SimilarWords from "./components/SimilarWords";
import ContextAnalysis from "./components/ContextAnalysis";
import Word2VecPanel from "./components/Word2VecPanel";
import Word2VecTools from "./components/Word2VecTools";
import DatasetPanel from "./components/DatasetPanel";
import MetricCard from "./components/MetricCard";
import "./styles.css";

type NavGroup = "data" | "training" | "analysis";
type TrainingTab = "model" | "w2v";
type AnalysisTab = "context" | "words" | "search" | "compare" | "keyword" | "match" | "batch";

const STEPS: { id: NavGroup; label: string; needsIndex?: boolean }[] = [
  { id: "data", label: "Data & Setup" },
  { id: "training", label: "Training" },
  { id: "analysis", label: "Analysis", needsIndex: true },
];

const TRAINING_TABS: { id: TrainingTab; label: string }[] = [
  { id: "model", label: "Fine-tune Model" },
  { id: "w2v", label: "Word2Vec Baseline" },
];

const ANALYSIS_TABS: { id: AnalysisTab; label: string }[] = [
  { id: "context", label: "Context" },
  { id: "words", label: "Similar Words" },
  { id: "search", label: "Search" },
  { id: "compare", label: "Compare" },
  { id: "keyword", label: "Keywords" },
  { id: "match", label: "Matcher" },
  { id: "batch", label: "Batch" },
];

export default function App() {
  const [group, setGroup] = useState<NavGroup>("data");
  const [trainingTab, setTrainingTab] = useState<TrainingTab>("model");
  const [analysisTab, setAnalysisTab] = useState<AnalysisTab>("context");
  const [stats, setStats] = useState<CorpusStats | null>(null);
  const [showManualSetup, setShowManualSetup] = useState(false);
  const [serverError, setServerError] = useState<string | null>(null);
  const [w2vReady, setW2vReady] = useState(false);
  const [w2vInfo, setW2vInfo] = useState<{ vocab_size: number; sentences: number; vector_size: number } | null>(null);
  const [resetLoading, setResetLoading] = useState(false);
  const ready = stats !== null && stats.index_built;

  useEffect(() => {
    checkConnection().then((err) => {
      setServerError(err);
      if (!err) {
        api.getStats().then(setStats).catch(() => {});
        api.w2vStatus().then(res => {
          if (res.ready) {
            setW2vReady(true);
            setW2vInfo({ vocab_size: res.vocab_size!, sentences: res.sentences!, vector_size: res.vector_size! });
          }
        }).catch(() => {});
      }
    });
    const interval = setInterval(() => {
      checkConnection().then(setServerError);
    }, 15000);
    return () => clearInterval(interval);
  }, []);

  function handleW2vReady(ready: boolean, info?: { vocab_size: number; sentences: number; vector_size: number }) {
    setW2vReady(ready);
    setW2vInfo(ready && info ? info : null);
  }

  async function handleReset() {
    setResetLoading(true);
    try {
      await api.w2vReset();
      setW2vReady(false);
      setW2vInfo(null);
    } catch {
      // ignore
    } finally {
      setResetLoading(false);
    }
  }

  function handleStepClick(id: NavGroup, needsIndex?: boolean) {
    if (needsIndex && !ready) return;
    setGroup(id);
  }

  // ── W2V trained: stats bar + analysis tabs, no stepper ──
  if (w2vReady && w2vInfo) {
    return (
      <div className="app">
        <header className="app-header">
          <h1>Contextual Similarity Engine</h1>
          {stats && (
            <div className="header-stats">
              <span className="badge">{stats.model_name}</span>
              <span className="badge">{stats.total_documents} docs</span>
              <span className="badge">{stats.total_chunks} chunks</span>
            </div>
          )}
        </header>

        {serverError && (
          <div className="server-error-banner">
            <strong>Server unavailable:</strong> {serverError}
          </div>
        )}

        {/* W2V stats bar */}
        <div className="content">
          <div className="panel">
            <div style={{ display: "flex", alignItems: "center", justifyContent: "space-between", flexWrap: "wrap", gap: 12 }}>
              <h2 style={{ margin: 0 }}>Word2Vec Baseline</h2>
              <button className="btn btn-secondary" onClick={handleReset} disabled={resetLoading}
                style={{ fontSize: "0.85em" }}>
                {resetLoading ? "Resetting..." : "Reset & Retrain"}
              </button>
            </div>
            <div className="metric-grid" style={{ marginTop: 12 }}>
              <MetricCard value={w2vInfo.vocab_size} label="Vocabulary" />
              <MetricCard value={w2vInfo.sentences} label="Sentences" />
              <MetricCard value={w2vInfo.vector_size} label="Dimensions" />
            </div>
          </div>

          {/* W2V-specific tools: Similar Words, Compare, Semantic Search */}
          <Word2VecTools />
        </div>

        {/* Transformer Analysis sub-tabs */}
        <nav className="subtabs">
          {ANALYSIS_TABS.map((t) => (
            <button
              key={t.id}
              className={`subtab ${analysisTab === t.id ? "subtab-active" : ""}`}
              onClick={() => setAnalysisTab(t.id)}
            >
              {t.label}
            </button>
          ))}
        </nav>

        {/* Analysis content */}
        <main className="content">
          {analysisTab === "context" && <ContextAnalysis />}
          {analysisTab === "words" && <SimilarWords />}
          {analysisTab === "search" && <SemanticSearch />}
          {analysisTab === "compare" && <TextCompare />}
          {analysisTab === "keyword" && <KeywordAnalysis />}
          {analysisTab === "match" && <KeywordMatcher />}
          {analysisTab === "batch" && <BatchAnalysis />}
        </main>
      </div>
    );
  }

  // ── Normal stepper flow ──
  return (
    <div className="app">
      <header className="app-header">
        <h1>Contextual Similarity Engine</h1>
        {stats && (
          <div className="header-stats">
            <span className="badge">{stats.model_name}</span>
            <span className="badge">{stats.total_documents} docs</span>
            <span className="badge">{stats.total_chunks} chunks</span>
            <span className={`badge ${stats.index_built ? "badge-ok" : "badge-warn"}`}>
              {stats.index_built ? "Index ready" : "Index not built"}
            </span>
          </div>
        )}
      </header>

      {serverError && (
        <div className="server-error-banner">
          <strong>Server unavailable:</strong> {serverError}
        </div>
      )}

      {/* Progress Stepper */}
      <nav className="stepper">
        {STEPS.map((step, i) => {
          const disabled = step.needsIndex && !ready;
          const active = group === step.id;
          const done = step.id === "data" && ready;
          return (
            <Fragment key={step.id}>
              {i > 0 && (
                <div className={`stepper-line ${!disabled ? "stepper-line-active" : ""}`} />
              )}
              <div className="stepper-item">
                <button
                  className={`stepper-circle ${active ? "stepper-active" : ""} ${done && !active ? "stepper-done" : ""}`}
                  onClick={() => handleStepClick(step.id, step.needsIndex)}
                  disabled={disabled}
                >
                  {done && !active ? "\u2713" : i + 1}
                </button>
                <span className={`stepper-label ${active ? "stepper-label-active" : ""}`}>
                  {step.label}
                </span>
              </div>
            </Fragment>
          );
        })}
      </nav>

      {/* Sub-tabs */}
      {group === "training" && (
        <nav className="subtabs">
          {TRAINING_TABS.map((t) => (
            <button
              key={t.id}
              className={`subtab ${trainingTab === t.id ? "subtab-active" : ""}`}
              onClick={() => setTrainingTab(t.id)}
            >
              {t.label}
            </button>
          ))}
        </nav>
      )}

      {group === "analysis" && (
        <nav className="subtabs">
          {ANALYSIS_TABS.map((t) => (
            <button
              key={t.id}
              className={`subtab ${analysisTab === t.id ? "subtab-active" : ""}`}
              onClick={() => setAnalysisTab(t.id)}
            >
              {t.label}
            </button>
          ))}
        </nav>
      )}

      {/* Content */}
      <main className="content">
        {group === "data" && (
          <>
            <DatasetPanel onStatsUpdate={setStats} />
            <button
              className="collapsible-toggle"
              onClick={() => setShowManualSetup(!showManualSetup)}
            >
              <span className="collapsible-arrow">{showManualSetup ? "\u25be" : "\u25b8"}</span>
              Or add documents manually
            </button>
            {showManualSetup && <EngineSetup onStatsUpdate={setStats} />}
          </>
        )}

        {group === "training" && trainingTab === "model" && <TrainingPanel />}
        {group === "training" && trainingTab === "w2v" && <Word2VecPanel onReady={handleW2vReady} />}

        {group === "analysis" && analysisTab === "context" && <ContextAnalysis />}
        {group === "analysis" && analysisTab === "words" && <SimilarWords />}
        {group === "analysis" && analysisTab === "search" && <SemanticSearch />}
        {group === "analysis" && analysisTab === "compare" && <TextCompare />}
        {group === "analysis" && analysisTab === "keyword" && <KeywordAnalysis />}
        {group === "analysis" && analysisTab === "match" && <KeywordMatcher />}
        {group === "analysis" && analysisTab === "batch" && <BatchAnalysis />}
      </main>
    </div>
  );
}