import { useState } from "react";
import {
  useDatabases,
  useSwitchDatabase,
  useSchemas,
  useTables,
  useViews,
  useFunctions,
  useSequences,
} from "@/hooks/use-schema";
import { useConnections } from "@/hooks/use-connections";
import { useAppStore } from "@/stores/app-store";
import { toast } from "sonner";
import {
  Database,
  Table2,
  Eye,
  FunctionSquare,
  Hash,
  FolderOpen,
  ChevronRight,
  ChevronDown,
  HardDrive,
} from "lucide-react";
import {
  ContextMenu,
  ContextMenuContent,
  ContextMenuItem,
  ContextMenuTrigger,
} from "@/components/ui/context-menu";
import type { Tab, SchemaObject } from "@/types";

function formatSize(bytes: number): string {
  if (bytes < 1024) return `${bytes} B`;
  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)} KB`;
  if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
  return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)} GB`;
}

function formatCount(n: number): string {
  if (n < 0) return "~0";
  if (n < 1000) return `~${n}`;
  if (n < 1_000_000) return `~${(n / 1000).toFixed(1)}k`;
  return `~${(n / 1_000_000).toFixed(1)}M`;
}

function TableSizeInfo({ item }: { item: SchemaObject }) {
  if (item.row_count == null && item.size_bytes == null) return null;
  const parts: string[] = [];
  if (item.row_count != null) parts.push(formatCount(item.row_count));
  if (item.size_bytes != null) parts.push(formatSize(item.size_bytes));
  return (
    <span className="ml-auto shrink-0 font-mono text-[10px] text-muted-foreground/50">
      {parts.join(", ")}
    </span>
  );
}

export function SchemaTree() {
  const { activeConnectionId, currentDatabase, setCurrentDatabase, addTab } =
    useAppStore();
  const { data: databases } = useDatabases(activeConnectionId);
  const { data: connections } = useConnections();
  const switchDbMutation = useSwitchDatabase();

  if (!activeConnectionId) {
    return (
      <div className="flex flex-col items-center gap-2 p-6 text-center">
        <Database className="h-8 w-8 text-muted-foreground/20" />
        <p className="text-sm text-muted-foreground/60">
          Connect to a database to browse schema
        </p>
      </div>
    );
  }

  if (!databases || databases.length === 0) {
    return (
      <div className="p-4 text-sm text-muted-foreground/60">
        No databases found.
      </div>
    );
  }

  const activeConfig = connections?.find((c) => c.id === activeConnectionId);

  const handleSwitchDb = (dbName: string) => {
    if (!activeConfig || dbName === currentDatabase) return;
    switchDbMutation.mutate(
      { config: activeConfig, database: dbName },
      {
        onSuccess: () => {
          setCurrentDatabase(dbName);
          toast.success(`Switched to database: ${dbName}`);
        },
        onError: (err) => {
          toast.error("Failed to switch database", {
            description: String(err),
          });
        },
      }
    );
  };

  return (
    <div className="space-y-0.5 p-2">
      {databases.map((db) => (
        <DatabaseNode
          key={db}
          name={db}
          isActive={db === currentDatabase}
          connectionId={activeConnectionId}
          onSwitch={() => handleSwitchDb(db)}
          isSwitching={switchDbMutation.isPending}
          onOpenTable={(schema, table) => {
            const tab: Tab = {
              id: crypto.randomUUID(),
              type: "table",
              title: table,
              connectionId: activeConnectionId,
              database: currentDatabase ?? undefined,
              schema,
              table,
            };
            addTab(tab);
          }}
          onViewStructure={(schema, table) => {
            const tab: Tab = {
              id: crypto.randomUUID(),
              type: "structure",
              title: `${table} (structure)`,
              connectionId: activeConnectionId,
              database: currentDatabase ?? undefined,
              schema,
              table,
            };
            addTab(tab);
          }}
        />
      ))}
    </div>
  );
}

function DatabaseNode({
  name,
  isActive,
  connectionId,
  onSwitch,
  isSwitching,
  onOpenTable,
  onViewStructure,
}: {
  name: string;
  isActive: boolean;
  connectionId: string;
  onSwitch: () => void;
  isSwitching: boolean;
  onOpenTable: (schema: string, table: string) => void;
  onViewStructure: (schema: string, table: string) => void;
}) {
  const [expanded, setExpanded] = useState(false);

  const handleClick = () => {
    if (!isActive) {
      onSwitch();
    }
    setExpanded(!expanded);
  };

  return (
    <div>
      <ContextMenu>
        <ContextMenuTrigger>
          <div
            className={`flex items-center gap-1.5 rounded-md px-1.5 py-1 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none font-medium ${
              isActive ? "text-foreground" : "text-muted-foreground"
            }`}
            onClick={handleClick}
          >
            <span className="text-muted-foreground/50">
              {expanded ? (
                <ChevronDown className="h-3.5 w-3.5" />
              ) : (
                <ChevronRight className="h-3.5 w-3.5" />
              )}
            </span>
            <HardDrive
              className={`h-3.5 w-3.5 shrink-0 ${isActive ? "text-primary" : "text-muted-foreground/50"}`}
            />
            <span className="truncate">{name}</span>
            {isActive && (
              <span className="ml-auto rounded-sm bg-primary/10 px-1 py-px text-[9px] font-semibold tracking-wider text-primary uppercase">
                active
              </span>
            )}
          </div>
        </ContextMenuTrigger>
        <ContextMenuContent>
          <ContextMenuItem
            onClick={() => toast.info(`Database: ${name}`, { description: isActive ? "Currently active" : "Not active" })}
          >
            Properties
          </ContextMenuItem>
        </ContextMenuContent>
      </ContextMenu>
      {expanded && isActive && (
        <div className="ml-4">
          <SchemasForCurrentDb
            connectionId={connectionId}
            onOpenTable={onOpenTable}
            onViewStructure={onViewStructure}
          />
        </div>
      )}
      {expanded && !isActive && (
        <div className="ml-6 py-1 text-xs text-muted-foreground/50">
          {isSwitching ? "Switching..." : "Click to switch to this database"}
        </div>
      )}
    </div>
  );
}

function SchemasForCurrentDb({
  connectionId,
  onOpenTable,
  onViewStructure,
}: {
  connectionId: string;
  onOpenTable: (schema: string, table: string) => void;
  onViewStructure: (schema: string, table: string) => void;
}) {
  const { data: schemas } = useSchemas(connectionId);

  if (!schemas || schemas.length === 0) {
    return (
      <div className="py-1 text-xs text-muted-foreground/50">No schemas found.</div>
    );
  }

  return (
    <>
      {schemas.map((schema) => (
        <SchemaNode
          key={schema}
          schema={schema}
          connectionId={connectionId}
          onOpenTable={(table) => onOpenTable(schema, table)}
          onViewStructure={(table) => onViewStructure(schema, table)}
        />
      ))}
    </>
  );
}

function SchemaNode({
  schema,
  connectionId,
  onOpenTable,
  onViewStructure,
}: {
  schema: string;
  connectionId: string;
  onOpenTable: (table: string) => void;
  onViewStructure: (table: string) => void;
}) {
  const [expanded, setExpanded] = useState(false);

  return (
    <div>
      <div
        className="flex items-center gap-1.5 rounded-md px-1.5 py-1 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none font-medium"
        onClick={() => setExpanded(!expanded)}
      >
        <span className="text-muted-foreground/50">
          {expanded ? (
            <ChevronDown className="h-3.5 w-3.5" />
          ) : (
            <ChevronRight className="h-3.5 w-3.5" />
          )}
        </span>
        {expanded ? (
          <FolderOpen className="h-3.5 w-3.5 text-primary/70" />
        ) : (
          <Database className="h-3.5 w-3.5 text-muted-foreground/50" />
        )}
        <span>{schema}</span>
      </div>
      {expanded && (
        <div className="ml-4">
          <CategoryNode
            label="Tables"
            connectionId={connectionId}
            schema={schema}
            category="tables"
            onOpenTable={onOpenTable}
            onViewStructure={onViewStructure}
          />
          <CategoryNode
            label="Views"
            connectionId={connectionId}
            schema={schema}
            category="views"
            onOpenTable={onOpenTable}
            onViewStructure={onViewStructure}
          />
          <CategoryNode
            label="Functions"
            connectionId={connectionId}
            schema={schema}
            category="functions"
            onOpenTable={onOpenTable}
            onViewStructure={onViewStructure}
          />
          <CategoryNode
            label="Sequences"
            connectionId={connectionId}
            schema={schema}
            category="sequences"
            onOpenTable={onOpenTable}
            onViewStructure={onViewStructure}
          />
        </div>
      )}
    </div>
  );
}

const categoryIcons = {
  tables: <Table2 className="h-3.5 w-3.5 text-data-table" />,
  views: <Eye className="h-3.5 w-3.5 text-data-view" />,
  functions: <FunctionSquare className="h-3.5 w-3.5 text-data-function" />,
  sequences: <Hash className="h-3.5 w-3.5 text-data-sequence" />,
};

function CategoryNode({
  label,
  connectionId,
  schema,
  category,
  onOpenTable,
  onViewStructure,
}: {
  label: string;
  connectionId: string;
  schema: string;
  category: "tables" | "views" | "functions" | "sequences";
  onOpenTable: (table: string) => void;
  onViewStructure: (table: string) => void;
}) {
  const [expanded, setExpanded] = useState(false);

  const tablesQuery = useTables(
    expanded && category === "tables" ? connectionId : null,
    schema
  );
  const viewsQuery = useViews(
    expanded && category === "views" ? connectionId : null,
    schema
  );
  const functionsQuery = useFunctions(
    expanded && category === "functions" ? connectionId : null,
    schema
  );
  const sequencesQuery = useSequences(
    expanded && category === "sequences" ? connectionId : null,
    schema
  );

  const items =
    category === "tables"
      ? tablesQuery.data
      : category === "views"
        ? viewsQuery.data
        : category === "functions"
          ? functionsQuery.data
          : sequencesQuery.data;

  const icon = categoryIcons[category];

  return (
    <div>
      <div
        className="flex items-center gap-1.5 rounded-md px-1.5 py-1 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none"
        onClick={() => setExpanded(!expanded)}
      >
        <span className="text-muted-foreground/50">
          {expanded ? (
            <ChevronDown className="h-3.5 w-3.5" />
          ) : (
            <ChevronRight className="h-3.5 w-3.5" />
          )}
        </span>
        {icon}
        <span className="truncate font-medium">
          {label}
          {items && (
            <span className="ml-1 font-normal text-muted-foreground/40">
              {items.length}
            </span>
          )}
        </span>
      </div>
      {expanded && (
        <div className="ml-4">
          {items?.map((item) => {
            const isTableOrView = category === "tables" || category === "views";

            if (isTableOrView) {
              return (
                <ContextMenu key={item.name}>
                  <ContextMenuTrigger>
                    <div
                      className="flex items-center gap-1.5 rounded-md px-1.5 py-0.5 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none"
                      onDoubleClick={() => onOpenTable(item.name)}
                    >
                      <span className="w-3.5 shrink-0" />
                      {icon}
                      <span className="truncate text-foreground/80">{item.name}</span>
                      {category === "tables" && <TableSizeInfo item={item} />}
                    </div>
                  </ContextMenuTrigger>
                  <ContextMenuContent>
                    <ContextMenuItem onClick={() => onOpenTable(item.name)}>
                      Open Data
                    </ContextMenuItem>
                    <ContextMenuItem
                      onClick={() => onViewStructure(item.name)}
                    >
                      View Structure
                    </ContextMenuItem>
                  </ContextMenuContent>
                </ContextMenu>
              );
            }

            return (
              <div
                key={item.name}
                className="flex items-center gap-1.5 rounded-md px-1.5 py-0.5 text-sm transition-colors hover:bg-accent/50 cursor-pointer select-none"
              >
                <span className="w-3.5 shrink-0" />
                {icon}
                <span className="truncate text-foreground/80">{item.name}</span>
              </div>
            );
          })}
        </div>
      )}
    </div>
  );
}
