import React, { useEffect, useMemo, useState } from "react";
import { createRoot } from "react-dom/client";
import {
  Activity,
  Braces,
  Code2,
  Database,
  FileCode2,
  GitBranch,
  Layers3,
  RefreshCw,
  Search,
  Server,
  Sparkles,
  Upload,
  Zap
} from "lucide-react";
import "./styles.css";

const DEFAULT_CODE = `from collections import defaultdict


def normalize_user(user):
    email = user["email"].strip().lower()
    return {
        "id": user["id"],
        "email": email,
        "active": bool(user.get("active", True)),
    }


def group_active_users(users):
    groups = defaultdict(list)
    for user in users:
        normalized = normalize_user(user)
        if normalized["active"]:
            domain = normalized["email"].split("@")[-1]
            groups[domain].append(normalized)
    return dict(groups)


def find_user_by_email(users, email):
    target = email.strip().lower()
    for user in users:
        normalized = normalize_user(user)
        if normalized["email"] == target:
            return normalized
    return None
`;

const DEFAULT_QUERY =
  "Where are active users grouped by email domain after normalizing their records?";

const SOURCE_LABELS = {
  raw_code: "Raw code",
  cross_file_propagated: "Propagated",
  semantic_synthetic: "Semantic"
};

const SOURCE_ICONS = {
  raw_code: FileCode2,
  cross_file_propagated: GitBranch,
  semantic_synthetic: Sparkles
};

let extractorPromise = null;
let embeddingBackend = "hashing";

function getExtractor() {
  if (!extractorPromise) {
    extractorPromise = import("@xenova/transformers").then(({ env, pipeline }) => {
      env.allowLocalModels = false;
      env.useBrowserCache = true;
      return pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2");
    });
  }
  return extractorPromise;
}

function formatNumber(value) {
  if (typeof value !== "number") return value;
  return value.toLocaleString(undefined, { maximumFractionDigits: 4 });
}

function sourceName(source) {
  return SOURCE_LABELS[source] || source;
}

function tokenize(text) {
  const words = (text.toLowerCase().match(/[a-z_][a-z0-9_]*|\d+(?:\.\d+)?/g) || []);
  const bigrams = words.slice(1).map((word, index) => `${words[index]}_${word}`);
  return [...words, ...bigrams];
}

function hashToken(token, dimensions) {
  let hash = 2166136261;
  for (let index = 0; index < token.length; index += 1) {
    hash ^= token.charCodeAt(index);
    hash = Math.imul(hash, 16777619);
  }
  return Math.abs(hash) % dimensions;
}

function embedTextHash(text, dimensions = 384) {
  const vector = new Float32Array(dimensions);
  for (const token of tokenize(text)) {
    const bucket = hashToken(token, dimensions);
    vector[bucket] += token.length % 2 === 0 ? 1 : -1;
  }
  let norm = 0;
  for (const value of vector) norm += value * value;
  norm = Math.sqrt(norm);
  if (norm > 0) {
    for (let index = 0; index < vector.length; index += 1) vector[index] /= norm;
  }
  return vector;
}

async function embedText(text) {
  try {
    const extractor = await getExtractor();
    const output = await extractor(text || " ", {
      pooling: "mean",
      normalize: true
    });
    embeddingBackend = "Xenova/all-MiniLM-L6-v2";
    return Float32Array.from(output.data);
  } catch (error) {
    console.warn("Falling back to hashing embeddings:", error);
    embeddingBackend = "hashing-fallback";
    return embedTextHash(text);
  }
}

function dot(left, right) {
  let score = 0;
  for (let index = 0; index < left.length; index += 1) score += left[index] * right[index];
  return score;
}

function extractFunctions(source, filePath) {
  const lines = source.replace(/\r\n/g, "\n").split("\n");
  const starts = [];
  const defPattern = /^(\s*)(async\s+def|def)\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(/;
  lines.forEach((line, index) => {
    const match = line.match(defPattern);
    if (match) {
      starts.push({ line: index, indent: match[1].length, name: match[3] });
    }
  });
  return starts.map((start, index) => {
    const next = starts.find((candidate, candidateIndex) => (
      candidateIndex > index && candidate.indent <= start.indent
    ));
    const endLine = next ? next.line : lines.length;
    const rawCode = lines.slice(start.line, endLine).join("\n").trim();
    const calls = Array.from(rawCode.matchAll(/\b([A-Za-z_][A-Za-z0-9_]*)\s*\(/g))
      .map((match) => match[1])
      .filter((name) => !["if", "for", "while", "return", "def"].includes(name));
    const uniqueCalls = [...new Set(calls.filter((name) => name !== start.name))];
    const words = start.name.replace(/_/g, " ");
    const ownSummary = `Function ${start.name} handles ${words}. Calls: ${uniqueCalls.slice(0, 8).join(", ") || "none"}.`;
    return {
      key: `${filePath}::${start.name}`,
      file_path: filePath,
      qualified_name: start.name,
      symbol_type: "function",
      start_line: start.line + 1,
      end_line: endLine,
      raw_code: rawCode,
      own_sketch: `${lines[start.line]}\n${ownSummary}`,
      propagated_text: `${ownSummary}\n\n${rawCode}`,
      direct_callees: uniqueCalls
    };
  });
}

function semanticQuestions(record) {
  const nameWords = record.qualified_name.replace(/_/g, " ");
  const questions = [
    `Where is ${nameWords} implemented?`,
    `Which function handles ${nameWords}?`,
    `How does ${nameWords} work?`
  ];
  if (/email|domain/i.test(record.raw_code)) {
    questions.push("Where are users grouped or matched by email domain?");
  }
  if (/active/i.test(record.raw_code)) {
    questions.push("Which function filters or processes active records?");
  }
  if (/normalize|strip|lower/i.test(record.raw_code)) {
    questions.push("Where is input normalized or cleaned?");
  }
  return questions;
}

async function buildBrowserIndex(source, filePath) {
  const functions = extractFunctions(source, filePath);
  const byName = Object.fromEntries(functions.map((record) => [record.qualified_name, record]));
  const vectors = [];
  for (const record of functions) {
    const calleeContext = record.direct_callees
      .map((name) => byName[name]?.own_sketch)
      .filter(Boolean)
      .join("\n");
    record.propagated_text = [record.propagated_text, calleeContext].filter(Boolean).join("\n\n");
    for (const [sourceNameValue, text] of [
      ["raw_code", record.raw_code],
      ["cross_file_propagated", record.propagated_text]
    ]) {
      vectors.push({
        record,
        source: sourceNameValue,
        text,
        scope: null,
        evidence: [],
        vector: await embedText(text)
      });
    }
    for (const question of semanticQuestions(record)) {
      vectors.push({
        record,
        source: "semantic_synthetic",
        text: question,
        scope: "behavior",
        evidence: [record.qualified_name],
        vector: await embedText(question)
      });
    }
  }
  const dimensions = vectors[0]?.vector.length || 384;
  return {
    dimensions,
    functions,
    vectors,
    stats: {
      schema_version: "coderag-browser-neural-vector-db-v1",
      database_path: "browser-memory",
      dimensions,
      embedding_backend: embeddingBackend,
      function_count: functions.length,
      vector_count: vectors.length,
      sources: vectors.reduce((counts, row) => {
        counts[row.source] = (counts[row.source] || 0) + 1;
        return counts;
      }, {}),
      active_index_mode: "browser_static",
      active_file_path: filePath
    }
  };
}

async function searchBrowserIndex(index, query, topK, perSource, selectedSources) {
  const queryVector = await embedText(query);
  const candidateRows = new Map();
  for (const sourceValue of selectedSources) {
    const ranked = index.vectors
      .filter((row) => row.source === sourceValue)
      .map((row) => ({ ...row, score: dot(row.vector, queryVector) }))
      .sort((left, right) => right.score - left.score)
      .slice(0, perSource);
    for (const row of ranked) {
      const existing = candidateRows.get(row.record.key);
      if (!existing || row.score > existing.best_score) {
        candidateRows.set(row.record.key, {
          ...row.record,
          source_scores: { ...(existing?.source_scores || {}), [sourceValue]: Number(row.score.toFixed(6)) },
          best_source: sourceValue,
          best_score: Number(row.score.toFixed(6)),
          matched_text: row.text,
          matched_scope: row.scope,
          matched_evidence: row.evidence
        });
      } else {
        existing.source_scores[sourceValue] = Number(row.score.toFixed(6));
      }
    }
  }
  const hits = [...candidateRows.values()]
    .sort((left, right) => right.best_score - left.best_score)
    .slice(0, topK);
  return {
    query,
    hits,
    candidate_count: candidateRows.size,
    per_source: perSource,
    sources: selectedSources,
    stats: index.stats
  };
}

function App() {
  const [code, setCode] = useState(DEFAULT_CODE);
  const [filename, setFilename] = useState("user_index.py");
  const [query, setQuery] = useState(DEFAULT_QUERY);
  const [topK, setTopK] = useState(6);
  const [perSource, setPerSource] = useState(5);
  const [sources, setSources] = useState([
    "raw_code",
    "cross_file_propagated",
    "semantic_synthetic"
  ]);
  const [stats, setStats] = useState(null);
  const [result, setResult] = useState(null);
  const [activeHit, setActiveHit] = useState(null);
  const [status, setStatus] = useState("idle");
  const [error, setError] = useState("");
  const [clientIndex, setClientIndex] = useState(null);

  useEffect(() => {
    refreshStats();
  }, []);

  useEffect(() => {
    setActiveHit(result?.hits?.[0] || null);
  }, [result]);

  const sourceCounts = useMemo(() => stats?.sources || {}, [stats]);
  const hasIndexedCode = stats?.active_index_mode === "pasted_code";

  async function refreshStats() {
    try {
      const response = await fetch("/api/coderag/stats");
      const payload = await response.json();
      if (!response.ok) throw new Error(payload.error || "Stats failed");
      setStats(payload);
    } catch (err) {
      setStatus("indexing");
      const fallback = await buildBrowserIndex(code, filename || "user_index.py");
      setClientIndex(fallback);
      setStats(fallback.stats);
      setError("");
      setStatus("idle");
    }
  }

  async function indexCode() {
    setStatus("indexing");
    setError("");
    setResult(null);
    try {
      const response = await fetch("/api/coderag/index-code", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ code, filename })
      });
      const payload = await response.json();
      if (!response.ok) throw new Error(payload.error || "Indexing failed");
      setStats(payload);
      setStatus("idle");
    } catch (err) {
      const fallback = await buildBrowserIndex(code, filename || "user_index.py");
      setClientIndex(fallback);
      setStats(fallback.stats);
      setError("");
      setStatus("idle");
    }
  }

  async function rebuildSampleDb() {
    setStatus("rebuilding");
    setError("");
    try {
      const response = await fetch("/api/coderag/rebuild", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({})
      });
      const payload = await response.json();
      if (!response.ok) throw new Error(payload.error || "Rebuild failed");
      setStats(payload);
      setResult(null);
      setStatus("idle");
    } catch (err) {
      const fallback = await buildBrowserIndex(DEFAULT_CODE, "sample_users.py");
      setCode(DEFAULT_CODE);
      setFilename("sample_users.py");
      setClientIndex(fallback);
      setStats(fallback.stats);
      setResult(null);
      setError("");
      setStatus("idle");
    }
  }

  async function search() {
    setStatus("searching");
    setError("");
    try {
      const response = await fetch("/api/coderag/search", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          query,
          top_k: topK,
          per_source: perSource,
          sources
        })
      });
      const payload = await response.json();
      if (!response.ok) throw new Error(payload.error || "Search failed");
      setResult(payload);
      setStats(payload.stats);
      setStatus("idle");
    } catch (err) {
      const fallback = clientIndex || await buildBrowserIndex(code, filename || "user_index.py");
      setClientIndex(fallback);
      const payload = await searchBrowserIndex(fallback, query, topK, perSource, sources);
      setResult(payload);
      setStats(payload.stats);
      setError("");
      setStatus("idle");
    }
  }

  function toggleSource(source) {
    setSources((current) => {
      if (current.includes(source) && current.length > 1) {
        return current.filter((item) => item !== source);
      }
      if (!current.includes(source)) return [...current, source];
      return current;
    });
  }

  return (
    <main className="app-shell">
      <section className="topbar">
        <div>
          <div className="eyebrow">
            <Database size={16} />
            Paste code, build vectors, ask questions
          </div>
          <h1>CodeRAG Retrieval Console</h1>
        </div>
        <button
          className="icon-button"
          onClick={rebuildSampleDb}
          disabled={status === "rebuilding"}
          title="Restore benchmark vector database"
        >
          <RefreshCw size={18} />
        </button>
      </section>

      <section className="status-strip">
        <Metric icon={Server} label="Functions" value={stats?.function_count ?? "-"} />
        <Metric icon={Layers3} label="Vector rows" value={stats?.vector_count ?? "-"} />
        <Metric icon={Braces} label="Dimensions" value={stats?.dimensions ?? "-"} />
        <Metric icon={Activity} label="Candidates" value={result?.candidate_count ?? "-"} />
      </section>

      <section className="workbench">
        <aside className="code-pane">
          <div className="panel-title">
            <Code2 size={18} />
            Code workspace
          </div>
          <label className="filename-field">
            File name
            <input
              value={filename}
              onChange={(event) => setFilename(event.target.value)}
              spellCheck="false"
            />
          </label>
          <textarea
            className="code-input"
            value={code}
            onChange={(event) => setCode(event.target.value)}
            spellCheck="false"
          />
          <button
            className="primary-button"
            onClick={indexCode}
            disabled={status === "indexing" || !code.trim()}
          >
            <Upload size={18} />
            Build code index
          </button>
          {hasIndexedCode && (
            <div className="success-box">
              Active index: {stats.function_count} functions from {stats.active_file_path}
            </div>
          )}
        </aside>

        <aside className="query-pane">
          <div className="panel-title">
            <Search size={18} />
            Question
          </div>
          <textarea
            className="question-input"
            value={query}
            onChange={(event) => setQuery(event.target.value)}
            spellCheck="false"
          />

          <div className="control-group">
            <label>
              Top matches
              <input
                type="number"
                min="1"
                max="20"
                value={topK}
                onChange={(event) => setTopK(Number(event.target.value))}
              />
            </label>
            <label>
              Per source
              <input
                type="number"
                min="1"
                max="10"
                value={perSource}
                onChange={(event) => setPerSource(Number(event.target.value))}
              />
            </label>
          </div>

          <div className="source-grid">
            {Object.keys(SOURCE_LABELS).map((source) => {
              const Icon = SOURCE_ICONS[source];
              const active = sources.includes(source);
              return (
                <button
                  key={source}
                  className={`source-toggle ${active ? "active" : ""}`}
                  onClick={() => toggleSource(source)}
                  title={`${sourceName(source)} vector source`}
                >
                  <Icon size={17} />
                  <span>{sourceName(source)}</span>
                  <small>{sourceCounts[source] ?? 0}</small>
                </button>
              );
            })}
          </div>

          <button
            className="primary-button"
            onClick={search}
            disabled={status === "searching" || !query.trim()}
          >
            <Zap size={18} />
            Retrieve top matches
          </button>

          {error && <div className="error-box">{error}</div>}
          {status === "indexing" && <div className="notice">Parsing functions, creating V7 sketches, generating semantic vectors, and writing SQLite rows.</div>}
          {status === "searching" && <div className="notice">Scoring vectors and merging candidates from selected sources.</div>}
          {status === "rebuilding" && <div className="notice">Restoring the benchmark repository vector database.</div>}
        </aside>

        <section className="results-pane">
          <div className="panel-title">
            <Layers3 size={18} />
            Top matches
          </div>
          {!result ? (
            <div className="empty-state">
              Build the code index, ask a question, and CodeRAG will retrieve the most relevant functions.
            </div>
          ) : (
            <div className="hit-list">
              {result.hits.map((hit, index) => (
                <HitRow
                  key={hit.key}
                  hit={hit}
                  index={index}
                  active={activeHit?.key === hit.key}
                  onClick={() => setActiveHit(hit)}
                />
              ))}
            </div>
          )}
        </section>
      </section>

      <section className="detail-pane wide-detail">
        <div className="panel-title">
          <FileCode2 size={18} />
          Retrieved evidence
        </div>
        {!activeHit ? (
          <div className="empty-state">Select a match to inspect its code, V7 sketch, source scores, and semantic provenance.</div>
        ) : (
          <HitDetail hit={activeHit} />
        )}
      </section>
    </main>
  );
}

function Metric({ icon: Icon, label, value }) {
  return (
    <div className="metric">
      <Icon size={18} />
      <div>
        <span>{label}</span>
        <strong>{formatNumber(value)}</strong>
      </div>
    </div>
  );
}

function HitRow({ hit, index, active, onClick }) {
  return (
    <button className={`hit-row ${active ? "active" : ""}`} onClick={onClick}>
      <span className="rank">{index + 1}</span>
      <span className="hit-main">
        <strong>{hit.qualified_name}</strong>
        <small>{hit.file_path}</small>
      </span>
      <span className={`source-pill ${hit.best_source}`}>{sourceName(hit.best_source)}</span>
      <span className="score">{hit.best_score.toFixed(4)}</span>
    </button>
  );
}

function HitDetail({ hit }) {
  return (
    <div className="detail-grid">
      <div className="detail-stack">
        <div className="detail-header">
          <div>
            <h2>{hit.qualified_name}</h2>
            <p>
              {hit.file_path}
              {hit.start_line ? `:${hit.start_line}-${hit.end_line}` : ""}
            </p>
          </div>
          <span className={`source-pill ${hit.best_source}`}>{sourceName(hit.best_source)}</span>
        </div>

        <section className="score-card">
          <h3>Source scores</h3>
          <div className="score-bars">
            {Object.entries(hit.source_scores).map(([source, score]) => (
              <div className="bar-row" key={source}>
                <span>{sourceName(source)}</span>
                <div className="bar-track">
                  <div className={`bar-fill ${source}`} style={{ width: `${Math.max(2, Math.min(100, score * 100))}%` }} />
                </div>
                <strong>{score.toFixed(4)}</strong>
              </div>
            ))}
          </div>
        </section>

        <section className="score-card">
          <h3>Matched vector text</h3>
          <p className="matched-text">{hit.matched_text || "No text stored for this vector."}</p>
          {hit.matched_scope && <span className="scope-chip">{hit.matched_scope}</span>}
          {hit.matched_evidence?.length > 0 && (
            <div className="evidence-list">
              {hit.matched_evidence.map((item) => (
                <code key={item}>{item}</code>
              ))}
            </div>
          )}
        </section>
      </div>

      <CodeBlock title="Raw code" value={hit.raw_code || "Code body unavailable."} />
      <CodeBlock title="V7 sketch and propagated context" value={hit.propagated_text || hit.own_sketch || "Sketch unavailable."} />
    </div>
  );
}

function CodeBlock({ title, value }) {
  return (
    <section className="code-card">
      <h3>{title}</h3>
      <pre>{value}</pre>
    </section>
  );
}

createRoot(document.getElementById("root")).render(<App />);