/* Copyright 2026 Marimo. All rights reserved. */

import { BoxIcon, CheckCircleIcon, XCircleIcon } from "lucide-react";
import React from "react";
import { Spinner } from "@/components/icons/spinner";
import { Button } from "@/components/ui/button";
import { Kbd } from "@/components/ui/kbd";
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from "@/components/ui/table";
import { toast } from "@/components/ui/use-toast";
import { useResolvedMarimoConfig } from "@/core/config/config";
import { useRequestClient } from "@/core/network/requests";
import { isWasm } from "@/core/wasm/utils";
import { useAsyncData } from "@/hooks/useAsyncData";
import { ErrorBanner } from "@/plugins/impl/common/error-banner";
import { cn } from "@/utils/cn";
import { semverSort } from "@/utils/versions";
import { SettingSubtitle } from "./common";

interface Package {
  name: string;
  minVersion?: string;
}

interface InstalledPackage {
  name: string;
  version: string;
}

const PEP_440_PRERELEASE =
  /^((?:\d+!)?\d+(?:\.\d+)*)(?:[-_.]?(?:a|b|c|rc|alpha|beta|pre|preview|dev)\d*)/i;

export function isPackageRequirementInstalled(
  requirement: Package,
  installedPackages: InstalledPackage[],
): boolean {
  const packageName = requirement.name.split("[")[0];
  const installedPackage = installedPackages.find(
    (pkg) => pkg.name === packageName,
  );
  if (!installedPackage) {
    return false;
  }
  if (!requirement.minVersion) {
    return true;
  }
  const prerelease = installedPackage.version.match(PEP_440_PRERELEASE);
  if (prerelease) {
    return semverSort(prerelease[1], requirement.minVersion) > 0;
  }
  return semverSort(installedPackage.version, requirement.minVersion) >= 0;
}

interface OptionalFeature {
  id: string;
  /**
   * Required packages to install for the feature to work.
   */
  packagesRequired: Package[];
  /**
   * Additional packages to install if installed through this UI.
   */
  additionalPackageInstalls: Package[];
  /**
   * Description of the feature.
   */
  description: string;
}

// Define the optional dependencies and their features
const OPTIONAL_DEPENDENCIES: OptionalFeature[] = [
  {
    id: "sql",
    packagesRequired: [{ name: "duckdb" }, { name: "sqlglot" }],
    additionalPackageInstalls: [{ name: "polars[pyarrow]" }],
    description: "SQL cells",
  },
  {
    id: "charts",
    packagesRequired: [{ name: "altair" }],
    additionalPackageInstalls: [],
    description: "Charts in datasource viewer",
  },
  {
    id: "fast-charts",
    packagesRequired: [{ name: "vegafusion" }, { name: "vl-convert-python" }],
    additionalPackageInstalls: [],
    description: "Fast server-side charts",
  },
  {
    id: "formatting",
    packagesRequired: [isWasm() ? { name: "black" } : { name: "ruff" }],
    additionalPackageInstalls: [],
    description: "Formatting",
  },
  {
    id: "ai",
    packagesRequired: [{ name: "openai" }],
    additionalPackageInstalls: [],
    description: "AI features",
  },
  {
    id: "lens",
    packagesRequired: [{ name: "marimo-lens", minVersion: "0.2.3" }],
    additionalPackageInstalls: [],
    description: "Select outputs for agents",
  },
  {
    id: "mcp",
    packagesRequired: [{ name: "mcp", minVersion: "2.0.0" }],
    additionalPackageInstalls: [],
    description: "Connect to MCP servers",
  },
  {
    id: "ipy-export",
    packagesRequired: [{ name: "nbformat" }],
    additionalPackageInstalls: [],
    description: "Export as IPYNB",
  },
  {
    id: "testing",
    packagesRequired: [{ name: "pytest" }],
    additionalPackageInstalls: [],
    description: "Autorun unit tests",
  },
];

// Only available outside wasm
if (!isWasm()) {
  OPTIONAL_DEPENDENCIES.push({
    id: "lsp",
    packagesRequired: [{ name: "python-lsp-server" }, { name: "websockets" }],
    additionalPackageInstalls: [{ name: "python-lsp-ruff" }],
    description: "Language Server Protocol*",
  });
}

export const OptionalFeatures: React.FC = () => {
  const [config] = useResolvedMarimoConfig();
  const packageManager = config.package_management.manager;
  const { getPackageList } = useRequestClient();
  const { data, error, refetch, isPending } = useAsyncData(
    () => getPackageList(),
    [packageManager],
  );

  if (isPending) {
    return <Spinner size="medium" centered={true} />;
  }

  if (error) {
    return <ErrorBanner error={error} />;
  }

  const installedPackages = data?.packages || [];

  return (
    <div className="flex-1 flex flex-col overflow-hidden gap-2">
      <SettingSubtitle>Optional Features</SettingSubtitle>
      <p className="text-sm text-muted-foreground">
        marimo is lightweight, with few dependencies, to maximize compatibility
        with your own environments.
        <br />
        To unlock additional features in the marimo editor, you can install
        these optional dependencies:
      </p>
      <Table>
        <TableHeader>
          <TableRow>
            <TableHead>Dependency</TableHead>
            <TableHead>Feature</TableHead>
            <TableHead>Status</TableHead>
            <TableHead />
          </TableRow>
        </TableHeader>
        <TableBody>
          {OPTIONAL_DEPENDENCIES.map((dep) => {
            const isInstalled = dep.packagesRequired.every((pkg) =>
              isPackageRequirementInstalled(pkg, installedPackages),
            );
            const packageSpec = dep.packagesRequired
              .map((pkg) => pkg.name)
              .join(", ");

            return (
              <TableRow key={dep.id} className="text-sm">
                <TableCell>{dep.description}</TableCell>
                <TableCell className="font-mono text-xs">
                  {packageSpec}
                </TableCell>
                <TableCell>
                  {isInstalled ? (
                    <div className="flex items-center">
                      <CheckCircleIcon className="h-4 w-4 text-(--grass-10) mr-2" />
                      <span>Installed</span>
                    </div>
                  ) : (
                    <div className="flex items-center">
                      <XCircleIcon className="h-4 w-4 text-(--red-10) mr-2" />
                      <InstallButton
                        packageSpecs={[
                          ...dep.packagesRequired,
                          ...dep.additionalPackageInstalls,
                        ]}
                        packageManager={packageManager}
                        onSuccess={refetch}
                      />
                    </div>
                  )}
                </TableCell>
              </TableRow>
            );
          })}
        </TableBody>
      </Table>

      <p className="text-muted-foreground mt-2">*Requires server restart</p>
    </div>
  );
};

const InstallButton: React.FC<{
  packageSpecs: Package[];
  packageManager: string;
  onSuccess: () => void;
}> = ({ packageSpecs, packageManager, onSuccess }) => {
  const [loading, setLoading] = React.useState(false);
  const { addPackage } = useRequestClient();

  const handleInstall = async () => {
    try {
      setLoading(true);
      const packageSpec = packageSpecs
        .map((pkg) => {
          if (pkg.minVersion) {
            return `${pkg.name}>=${pkg.minVersion}`;
          }
          return pkg.name;
        })
        .join(" ");
      const response = await addPackage({ package: packageSpec, group: "dev" });
      if (response.success) {
        onSuccess();
        toast({
          title: "Package installed",
          description: (
            <span>
              The packages{" "}
              <Kbd className="inline">
                {packageSpecs.map((pkg) => pkg.name).join(", ")}
              </Kbd>{" "}
              have been added to your environment.
            </span>
          ),
        });
      } else {
        toast({
          title: "Failed to install package",
          description: response.error,
          variant: "danger",
        });
      }
    } finally {
      setLoading(false);
    }
  };

  return (
    <Button
      size="xs"
      variant="outline"
      className={cn("text-xs", loading && "opacity-50 cursor-not-allowed")}
      onClick={handleInstall}
      disabled={loading}
    >
      {loading ? (
        <Spinner size="small" className="mr-2 h-3 w-3" />
      ) : (
        <BoxIcon className="mr-2 h-3 w-3" />
      )}
      Install with {packageManager}
    </Button>
  );
};