frontend/src/components/AstTreeView.tsx

import { useState, useEffect, useMemo, useCallback } from "react";
import { ChevronRight, ChevronDown, Search, X } from "lucide-react";
import type { AstNode } from "@/lib/api";

// ── Node type → CSS class mapping ─────────────────────────────────────

function getNodeClass(type: string): string {
  const t = type.toLowerCase();
  if (t.includes("keyword") || t.includes("statement") || t.includes("clause")) return "ast-keyword";
  if (t.includes("literal") || t.includes("quoted") || t.includes("string")) return "ast-literal";
  if (t.includes("numeric") || t.includes("integer") || t.includes("float")) return "ast-numeric";
  if (t.includes("comment")) return "ast-comment";
  if (t.includes("function") || t.includes("func")) return "ast-function";
  if (t.includes("operator") || t.includes("comparison") || t.includes("arithmetic")) return "ast-operator";
  if (t.includes("type") || t.includes("datatype")) return "ast-type";
  if (t.includes("whitespace") || t.includes("newline") || t.includes("indent")) return "ast-whitespace";
  return "ast-default";
}

function getNodeBadgeStyle(type: string): React.CSSProperties {
  const t = type.toLowerCase();
  if (t.includes("keyword") || t.includes("statement") || t.includes("clause"))
    return { background: "oklch(0.72 0.18 270 / 0.15)", color: "oklch(0.72 0.18 270)", border: "1px solid oklch(0.72 0.18 270 / 0.3)" };
  if (t.includes("literal") || t.includes("quoted") || t.includes("string"))
    return { background: "oklch(0.75 0.16 145 / 0.15)", color: "oklch(0.75 0.16 145)", border: "1px solid oklch(0.75 0.16 145 / 0.3)" };
  if (t.includes("numeric") || t.includes("integer") || t.includes("float"))
    return { background: "oklch(0.78 0.18 55 / 0.15)", color: "oklch(0.78 0.18 55)", border: "1px solid oklch(0.78 0.18 55 / 0.3)" };
  if (t.includes("comment"))
    return { background: "oklch(0.50 0.010 264 / 0.15)", color: "oklch(0.55 0.010 264)", border: "1px solid oklch(0.50 0.010 264 / 0.3)" };
  if (t.includes("function") || t.includes("func"))
    return { background: "oklch(0.78 0.16 210 / 0.15)", color: "oklch(0.78 0.16 210)", border: "1px solid oklch(0.78 0.16 210 / 0.3)" };
  if (t.includes("operator") || t.includes("comparison"))
    return { background: "oklch(0.82 0.010 264 / 0.10)", color: "oklch(0.82 0.010 264)", border: "1px solid oklch(0.82 0.010 264 / 0.3)" };
  if (t.includes("whitespace") || t.includes("newline"))
    return { background: "oklch(0.20 0.008 264 / 0.5)", color: "oklch(0.38 0.008 264)", border: "1px solid oklch(0.25 0.008 264 / 0.3)" };
  return { background: "oklch(0.20 0.010 264 / 0.5)", color: "oklch(0.72 0.010 264)", border: "1px solid oklch(0.28 0.010 264 / 0.3)" };
}

// ── Search helpers ─────────────────────────────────────────────────────

function collectMatchIds(node: AstNode, query: string, matchIds: Set<string>, ancestorIds: Set<string>, path: string[]): boolean {
  const q = query.toLowerCase();
  const matches =
    node.type.toLowerCase().includes(q) ||
    node.name.toLowerCase().includes(q) ||
    (node.raw?.toLowerCase().includes(q) ?? false);

  let childMatched = false;
  for (const child of node.children) {
    if (collectMatchIds(child, query, matchIds, ancestorIds, [...path, node.id])) {
      childMatched = true;
    }
  }

  if (matches) {
    matchIds.add(node.id);
    path.forEach((id) => ancestorIds.add(id));
  }

  return matches || childMatched;
}

function highlightText(text: string, query: string): React.ReactNode {
  if (!query) return text;
  const idx = text.toLowerCase().indexOf(query.toLowerCase());
  if (idx === -1) return text;
  return (
    <>
      {text.slice(0, idx)}
      <mark style={{ background: "oklch(0.68 0.16 210 / 0.35)", color: "inherit", borderRadius: "2px" }}>
        {text.slice(idx, idx + query.length)}
      </mark>
      {text.slice(idx + query.length)}
    </>
  );
}

// ── Single tree node ───────────────────────────────────────────────────

interface NodeProps {
  node: AstNode;
  depth: number;
  query: string;
  matchIds: Set<string>;
  ancestorIds: Set<string>;
  expandedIds: Set<string>;
  onToggle: (id: string) => void;
  onJumpToLine?: (line: number) => void;
}

function TreeNode({ node, depth, query, matchIds, ancestorIds, expandedIds, onToggle, onJumpToLine }: NodeProps) {
  const isMatch = matchIds.has(node.id);
  const isAncestor = ancestorIds.has(node.id);
  const isExpanded = expandedIds.has(node.id);
  const hasChildren = node.children.length > 0;

  // Skip pure whitespace nodes when no search
  const isWhitespace = node.type.toLowerCase().includes("whitespace") || node.type.toLowerCase().includes("newline");
  if (isWhitespace && !query && !node.raw?.trim()) return null;

  const indent = depth * 16;
  const nodeClass = getNodeClass(node.type);
  const badgeStyle = getNodeBadgeStyle(node.type);

  return (
    <div>
      <div
        className={`flex items-center gap-1.5 py-0.5 px-2 rounded cursor-pointer group transition-colors duration-100 ${
          isMatch
            ? "bg-[oklch(0.68_0.16_210_/_0.12)] border border-[oklch(0.68_0.16_210_/_0.3)]"
            : "hover:bg-[oklch(0.18_0.010_264)]"
        }`}
        style={{ paddingLeft: `${indent + 8}px` }}
        onClick={() => {
          if (hasChildren) onToggle(node.id);
          if (node.start_line && onJumpToLine) onJumpToLine(node.start_line);
        }}
      >
        {/* Expand/collapse chevron */}
        <span className="w-4 h-4 flex items-center justify-center flex-shrink-0 text-[oklch(0.45_0.010_264)]">
          {hasChildren ? (
            isExpanded ? <ChevronDown size={12} /> : <ChevronRight size={12} />
          ) : (
            <span className="w-1 h-1 rounded-full bg-[oklch(0.30_0.010_264)]" />
          )}
        </span>

        {/* Node type badge */}
        <span
          className="text-[10px] font-mono font-medium px-1.5 py-0.5 rounded flex-shrink-0"
          style={badgeStyle}
        >
          {highlightText(node.type, query)}
        </span>

        {/* Raw value for leaf nodes */}
        {node.is_leaf && node.raw && node.raw.trim() && (
          <span className={`text-xs font-mono truncate max-w-[200px] ${nodeClass}`}>
            {highlightText(JSON.stringify(node.raw), query)}
          </span>
        )}

        {/* Position info */}
        {node.start_line && (
          <span className="text-[10px] text-[oklch(0.38_0.010_264)] ml-auto flex-shrink-0 opacity-0 group-hover:opacity-100 transition-opacity">
            L{node.start_line}:{node.start_pos}
          </span>
        )}
      </div>

      {/* Children */}
      {hasChildren && isExpanded && (
        <div>
          {node.children.map((child) => (
            <TreeNode
              key={child.id}
              node={child}
              depth={depth + 1}
              query={query}
              matchIds={matchIds}
              ancestorIds={ancestorIds}
              expandedIds={expandedIds}
              onToggle={onToggle}
              onJumpToLine={onJumpToLine}
            />
          ))}
        </div>
      )}
    </div>
  );
}

// ── Main component ─────────────────────────────────────────────────────

interface AstTreeViewProps {
  tree: AstNode | null;
  tokenCount?: number;
  depth?: number;
  onJumpToLine?: (line: number) => void;
}

export function AstTreeView({ tree, tokenCount, depth, onJumpToLine }: AstTreeViewProps) {
  const [query, setQuery] = useState("");
  const [expandedIds, setExpandedIds] = useState<Set<string>>(new Set());

  // Auto-expand root on first load
  useEffect(() => {
    if (tree) {
      setExpandedIds(new Set([tree.id, ...tree.children.map((c) => c.id)]));
    }
  }, [tree]);

  const { matchIds, ancestorIds } = useMemo(() => {
    if (!tree || !query.trim()) return { matchIds: new Set<string>(), ancestorIds: new Set<string>() };
    const matchIds = new Set<string>();
    const ancestorIds = new Set<string>();
    collectMatchIds(tree, query.trim(), matchIds, ancestorIds, []);
    return { matchIds, ancestorIds };
  }, [tree, query]);

  // Auto-expand ancestors when searching
  useEffect(() => {
    if (query.trim() && ancestorIds.size > 0) {
      setExpandedIds((prev) => {
        const next = new Set(prev);
        ancestorIds.forEach((id) => next.add(id));
        matchIds.forEach((id) => next.add(id));
        return next;
      });
    }
  }, [ancestorIds, matchIds, query]);

  const handleToggle = useCallback((id: string) => {
    setExpandedIds((prev) => {
      const next = new Set(prev);
      if (next.has(id)) next.delete(id);
      else next.add(id);
      return next;
    });
  }, []);

  const expandAll = useCallback(() => {
    if (!tree) return;
    const ids = new Set<string>();
    const collect = (n: AstNode) => { ids.add(n.id); n.children.forEach(collect); };
    collect(tree);
    setExpandedIds(ids);
  }, [tree]);

  const collapseAll = useCallback(() => {
    if (!tree) return;
    setExpandedIds(new Set([tree.id]));
  }, [tree]);

  if (!tree) {
    return (
      <div className="flex flex-col items-center justify-center h-full gap-3 text-[oklch(0.45_0.010_264)]">
        <div className="text-4xl opacity-30">⟨/⟩</div>
        <p className="text-sm">Run analysis to view the AST</p>
      </div>
    );
  }

  return (
    <div className="flex flex-col h-full">
      {/* Toolbar */}
      <div className="flex items-center gap-2 px-3 py-2 border-b border-[oklch(0.22_0.010_264)] flex-shrink-0">
        <div className="relative flex-1">
          <Search size={12} className="absolute left-2.5 top-1/2 -translate-y-1/2 text-[oklch(0.45_0.010_264)]" />
          <input
            value={query}
            onChange={(e) => setQuery(e.target.value)}
            placeholder="Filter by node type…"
            className="h-7 pl-7 pr-7 text-xs w-full rounded"
            style={{ background: "var(--bg-elevated)", border: "1px solid var(--border)", color: "var(--text-primary)", fontFamily: "var(--font-mono)", outline: "none", paddingLeft: "28px", paddingRight: "28px" }}
          />
          {query && (
            <button
              onClick={() => setQuery("")}
              className="absolute right-2 top-1/2 -translate-y-1/2 text-[oklch(0.45_0.010_264)] hover:text-foreground"
            >
              <X size={12} />
            </button>
          )}
        </div>
        <button onClick={expandAll} className="h-7 px-2 text-xs rounded transition-colors" style={{ color: "var(--text-secondary)", background: "transparent" }}>Expand</button>
        <button onClick={collapseAll} className="h-7 px-2 text-xs rounded transition-colors" style={{ color: "var(--text-secondary)", background: "transparent" }}>Collapse</button>
      </div>

      {/* Stats */}
      <div className="flex items-center gap-3 px-3 py-1.5 border-b border-[oklch(0.22_0.010_264)] flex-shrink-0">
        <span className="text-[10px] text-[oklch(0.45_0.010_264)]">
          <span className="text-[oklch(0.68_0.16_210)]">{tokenCount}</span> tokens
        </span>
        <span className="text-[10px] text-[oklch(0.45_0.010_264)]">
          depth <span className="text-[oklch(0.68_0.16_210)]">{depth}</span>
        </span>
        {query && matchIds.size > 0 && (
          <span className="text-[10px] text-[oklch(0.75_0.16_145)]">
            {matchIds.size} match{matchIds.size !== 1 ? "es" : ""}
          </span>
        )}
        {query && matchIds.size === 0 && (
          <span className="text-[10px] text-[oklch(0.60_0.20_25)]">No matches</span>
        )}
      </div>

      {/* Tree */}
      <div className="flex-1 overflow-auto py-1 text-xs">
        <TreeNode
          node={tree}
          depth={0}
          query={query}
          matchIds={matchIds}
          ancestorIds={ancestorIds}
          expandedIds={expandedIds}
          onToggle={handleToggle}
          onJumpToLine={onJumpToLine}
        />
      </div>
    </div>
  );
}