tusk/src/components/workspace/WorkspacePanel.tsx

import { useState, useCallback } from "react";
import {
  ResizableHandle,
  ResizablePanel,
  ResizablePanelGroup,
} from "@/components/ui/resizable";
import { SqlEditor } from "@/components/editor/SqlEditor";
import { ResultsPanel } from "@/components/results/ResultsPanel";
import { ExplainView } from "@/components/results/ExplainView";
import { useQueryExecution } from "@/hooks/use-query-execution";
import { useAddHistory } from "@/hooks/use-history";
import { useCompletionSchema } from "@/hooks/use-completion-schema";
import { useConnections } from "@/hooks/use-connections";
import { useAppStore } from "@/stores/app-store";
import { Button } from "@/components/ui/button";
import { Play, Loader2, Lock, BarChart3, Download, AlignLeft, Bookmark, Table2, Braces, Sparkles, BrainCircuit } from "lucide-react";
import { format as formatSql } from "sql-formatter";
import { SaveQueryDialog } from "@/components/saved-queries/SaveQueryDialog";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { exportCsv, exportJson } from "@/lib/tauri";
import { save } from "@tauri-apps/plugin-dialog";
import { toast } from "sonner";
import { AiBar } from "@/components/ai/AiBar";
import { useExplainSql, useFixSqlError } from "@/hooks/use-ai";
import type { QueryResult, ExplainResult } from "@/types";

interface Props {
  connectionId: string;
  initialSql?: string;
  onSqlChange?: (sql: string) => void;
  onResult?: (result: QueryResult | null, error: string | null) => void;
}

export function WorkspacePanel({
  connectionId,
  initialSql = "",
  onSqlChange,
  onResult,
}: Props) {
  const readOnlyMap = useAppStore((s) => s.readOnlyMap);
  const currentDatabase = useAppStore((s) => s.currentDatabase);
  const isReadOnly = readOnlyMap[connectionId] ?? true;

  const [sqlValue, setSqlValue] = useState(initialSql);
  const [result, setResult] = useState<QueryResult | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [explainData, setExplainData] = useState<ExplainResult | null>(null);
  const [resultView, setResultView] = useState<"results" | "explain">("results");
  const [resultViewMode, setResultViewMode] = useState<"table" | "json">("table");
  const [saveDialogOpen, setSaveDialogOpen] = useState(false);
  const [aiBarOpen, setAiBarOpen] = useState(false);
  const [aiExplanation, setAiExplanation] = useState<string | null>(null);

  const queryMutation = useQueryExecution();
  const explainMutation = useExplainSql();
  const fixMutation = useFixSqlError();
  const addHistoryMutation = useAddHistory();
  const { data: connections } = useConnections();
  const { data: completionSchema } = useCompletionSchema(connectionId);

  const connName =
    connections?.find((c) => c.id === connectionId)?.name ?? "unknown";

  const handleChange = useCallback(
    (val: string) => {
      setSqlValue(val);
      onSqlChange?.(val);
    },
    [onSqlChange]
  );

  const recordHistory = useCallback(
    (
      sql: string,
      status: "success" | "error",
      executionTimeMs: number,
      rowCount?: number,
      errorMessage?: string
    ) => {
      addHistoryMutation.mutate({
        id: crypto.randomUUID(),
        connection_id: connectionId,
        connection_name: connName,
        database: currentDatabase ?? "",
        sql,
        status,
        error_message: errorMessage,
        row_count: rowCount,
        execution_time_ms: executionTimeMs,
        executed_at: new Date().toISOString(),
      });
    },
    [addHistoryMutation, connectionId, connName, currentDatabase]
  );

  const handleExecute = useCallback(() => {
    if (!sqlValue.trim() || !connectionId) return;
    setError(null);
    setExplainData(null);
    setAiExplanation(null);
    setResultView("results");
    queryMutation.mutate(
      { connectionId, sql: sqlValue },
      {
        onSuccess: (data) => {
          setResult(data);
          setError(null);
          onResult?.(data, null);
          recordHistory(
            sqlValue,
            "success",
            data.execution_time_ms,
            data.row_count
          );
        },
        onError: (err) => {
          setResult(null);
          setError(String(err));
          onResult?.(null, String(err));
          recordHistory(sqlValue, "error", 0, undefined, String(err));
        },
      }
    );
  }, [connectionId, sqlValue, queryMutation, onResult, recordHistory]);

  const handleExplain = useCallback(() => {
    if (!sqlValue.trim() || !connectionId) return;
    setError(null);
    const explainSql = `EXPLAIN (ANALYZE, COSTS, BUFFERS, FORMAT JSON) ${sqlValue}`;
    queryMutation.mutate(
      { connectionId, sql: explainSql },
      {
        onSuccess: (data) => {
          try {
            const raw = data.rows[0]?.[0];
            const parsed =
              typeof raw === "string" ? JSON.parse(raw) : raw;
            const plan: ExplainResult = Array.isArray(parsed)
              ? parsed[0]
              : parsed;
            setExplainData(plan);
            setResultView("explain");
            setResult(null);
            setError(null);
          } catch {
            setError("Failed to parse EXPLAIN output");
            setExplainData(null);
          }
        },
        onError: (err) => {
          setResult(null);
          setError(String(err));
          setExplainData(null);
        },
      }
    );
  }, [connectionId, sqlValue, queryMutation]);

  const handleFormat = useCallback(() => {
    if (!sqlValue.trim()) return;
    try {
      const formatted = formatSql(sqlValue, { language: "postgresql" });
      setSqlValue(formatted);
      onSqlChange?.(formatted);
    } catch {
      // Silently ignore format errors on invalid SQL
    }
  }, [sqlValue, onSqlChange]);

  const handleExport = useCallback(
    async (format: "csv" | "json") => {
      if (!result || result.columns.length === 0) return;
      const ext = format === "csv" ? "csv" : "json";
      const path = await save({
        title: `Export as ${ext.toUpperCase()}`,
        defaultPath: `query_result.${ext}`,
        filters: [{ name: ext.toUpperCase(), extensions: [ext] }],
      });
      if (!path) return;
      try {
        if (format === "csv") {
          await exportCsv(path, result.columns, result.rows as unknown[][]);
        } else {
          await exportJson(path, result.columns, result.rows as unknown[][]);
        }
        toast.success(`Exported to ${path}`);
      } catch (err) {
        toast.error("Export failed", { description: String(err) });
      }
    },
    [result]
  );

  const isAiLoading = explainMutation.isPending || fixMutation.isPending;

  const handleAiExplain = useCallback(() => {
    if (!sqlValue.trim() || !connectionId) return;
    setAiExplanation(null);
    setResultView("results");
    explainMutation.mutate(
      { connectionId, sql: sqlValue },
      {
        onSuccess: (explanation) => {
          setAiExplanation(explanation);
        },
        onError: (err) => {
          toast.error("AI Explain failed", { description: String(err) });
        },
      }
    );
  }, [connectionId, sqlValue, explainMutation]);

  const handleExplainError = useCallback(() => {
    if (!sqlValue.trim() || !connectionId || !error) return;
    setAiExplanation(null);
    explainMutation.mutate(
      { connectionId, sql: `${sqlValue}\n\n-- Error: ${error}` },
      {
        onSuccess: (explanation) => {
          setAiExplanation(explanation);
        },
        onError: (err) => {
          toast.error("AI Explain failed", { description: String(err) });
        },
      }
    );
  }, [connectionId, sqlValue, error, explainMutation]);

  const handleFixError = useCallback(() => {
    if (!sqlValue.trim() || !connectionId || !error) return;
    fixMutation.mutate(
      { connectionId, sql: sqlValue, errorMessage: error },
      {
        onSuccess: (fixedSql) => {
          setSqlValue(fixedSql);
          onSqlChange?.(fixedSql);
          setError(null);
          setAiExplanation(null);
          toast.success("SQL replaced by AI suggestion");
        },
        onError: (err) => {
          toast.error("AI Fix failed", { description: String(err) });
        },
      }
    );
  }, [connectionId, sqlValue, error, fixMutation, onSqlChange]);

  return (
    <>
    <ResizablePanelGroup orientation="vertical">
      <ResizablePanel id="editor" defaultSize="40%" minSize="15%">
        <div className="flex h-full flex-col">
          <div className="flex items-center gap-2 border-b px-2 py-1">
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={handleExecute}
              disabled={queryMutation.isPending || !sqlValue.trim()}
            >
              {queryMutation.isPending && resultView === "results" ? (
                <Loader2 className="h-3 w-3 animate-spin" />
              ) : (
                <Play className="h-3 w-3" />
              )}
              Run
            </Button>
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={handleExplain}
              disabled={queryMutation.isPending || !sqlValue.trim()}
            >
              {queryMutation.isPending && resultView === "explain" ? (
                <Loader2 className="h-3 w-3 animate-spin" />
              ) : (
                <BarChart3 className="h-3 w-3" />
              )}
              Explain
            </Button>
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={handleFormat}
              disabled={!sqlValue.trim()}
              title="Format SQL (Shift+Alt+F)"
            >
              <AlignLeft className="h-3 w-3" />
              Format
            </Button>
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={() => setSaveDialogOpen(true)}
              disabled={!sqlValue.trim()}
              title="Save query"
            >
              <Bookmark className="h-3 w-3" />
              Save
            </Button>
            <Button
              size="sm"
              variant={aiBarOpen ? "secondary" : "ghost"}
              className="h-6 gap-1 text-xs"
              onClick={() => setAiBarOpen(!aiBarOpen)}
              title="AI SQL Generator"
            >
              <Sparkles className="h-3 w-3" />
              AI
            </Button>
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={handleAiExplain}
              disabled={isAiLoading || !sqlValue.trim()}
              title="Explain query with AI"
            >
              {isAiLoading ? (
                <Loader2 className="h-3 w-3 animate-spin" />
              ) : (
                <BrainCircuit className="h-3 w-3" />
              )}
              AI Explain
            </Button>
            {result && result.columns.length > 0 && (
              <DropdownMenu>
                <DropdownMenuTrigger asChild>
                  <Button
                    size="sm"
                    variant="ghost"
                    className="h-6 gap-1 text-xs"
                  >
                    <Download className="h-3 w-3" />
                    Export
                  </Button>
                </DropdownMenuTrigger>
                <DropdownMenuContent align="start">
                  <DropdownMenuItem onClick={() => handleExport("csv")}>
                    Export CSV
                  </DropdownMenuItem>
                  <DropdownMenuItem onClick={() => handleExport("json")}>
                    Export JSON
                  </DropdownMenuItem>
                </DropdownMenuContent>
              </DropdownMenu>
            )}
            <span className="text-[11px] text-muted-foreground">
              Ctrl+Enter to execute
            </span>
            {isReadOnly && (
              <span className="flex items-center gap-1 rounded bg-yellow-500/10 px-1.5 py-0.5 text-[10px] font-semibold text-yellow-600 dark:text-yellow-500">
                <Lock className="h-3 w-3" />
                Read-Only
              </span>
            )}
          </div>
          {aiBarOpen && (
            <AiBar
              connectionId={connectionId}
              onSqlGenerated={(sql) => {
                setSqlValue(sql);
                onSqlChange?.(sql);
              }}
              onClose={() => setAiBarOpen(false)}
              onExecute={handleExecute}
            />
          )}
          <div className="min-h-0 flex-1">
            <SqlEditor
              value={sqlValue}
              onChange={handleChange}
              onExecute={handleExecute}
              onFormat={handleFormat}
              schema={completionSchema}
            />
          </div>
        </div>
      </ResizablePanel>
      <ResizableHandle withHandle />
      <ResizablePanel id="results" defaultSize="60%" minSize="15%">
        <div className="flex h-full flex-col overflow-hidden">
          {(explainData || result || error || aiExplanation) && (
            <div className="flex shrink-0 items-center border-b text-xs">
              <button
                className={`px-3 py-1 font-medium ${
                  resultView === "results"
                    ? "bg-background text-foreground"
                    : "text-muted-foreground hover:text-foreground"
                }`}
                onClick={() => setResultView("results")}
              >
                Results
              </button>
              {explainData && (
                <button
                  className={`px-3 py-1 font-medium ${
                    resultView === "explain"
                      ? "bg-background text-foreground"
                      : "text-muted-foreground hover:text-foreground"
                  }`}
                  onClick={() => setResultView("explain")}
                >
                  Explain
                </button>
              )}
              {resultView === "results" && result && result.columns.length > 0 && (
                <div className="ml-auto mr-2 flex items-center rounded-md border">
                  <button
                    className={`flex items-center gap-1 px-2 py-0.5 font-medium ${
                      resultViewMode === "table"
                        ? "bg-muted text-foreground"
                        : "text-muted-foreground hover:text-foreground"
                    }`}
                    onClick={() => setResultViewMode("table")}
                    title="Table view"
                  >
                    <Table2 className="h-3 w-3" />
                    Table
                  </button>
                  <button
                    className={`flex items-center gap-1 px-2 py-0.5 font-medium ${
                      resultViewMode === "json"
                        ? "bg-muted text-foreground"
                        : "text-muted-foreground hover:text-foreground"
                    }`}
                    onClick={() => setResultViewMode("json")}
                    title="JSON view"
                  >
                    <Braces className="h-3 w-3" />
                    JSON
                  </button>
                </div>
              )}
            </div>
          )}
          <div className="min-h-0 flex-1">
            {resultView === "explain" && explainData ? (
              <ExplainView data={explainData} />
            ) : (
              <ResultsPanel
                result={result}
                error={error}
                isLoading={queryMutation.isPending && resultView === "results"}
                viewMode={resultViewMode}
                aiExplanation={aiExplanation}
                isAiLoading={isAiLoading}
                onExplainError={error ? handleExplainError : undefined}
                onFixError={error ? handleFixError : undefined}
              />
            )}
          </div>
        </div>
      </ResizablePanel>
    </ResizablePanelGroup>

    <SaveQueryDialog
      open={saveDialogOpen}
      onOpenChange={setSaveDialogOpen}
      sql={sqlValue}
      connectionId={connectionId}
    />
    </>
  );
}