import { XMarkIcon } from "@heroicons/react/20/solid";
import { type LoaderFunctionArgs } from "@remix-run/node";
import { Form } from "@remix-run/react";
import type { TaskTriggerSource } from "@trigger.dev/database";
import { type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from "react";
import ReactGridLayout from "react-grid-layout";
import { typedjson, useTypedLoaderData } from "remix-typedjson";
import { z } from "zod";
import { PageBody, PageContainer } from "~/components/layout/AppLayout";
import { LogsTaskFilter } from "~/components/logs/LogsTaskFilter";
import { ModelsFilter, type ModelOption } from "~/components/metrics/ModelsFilter";
import { OperationsFilter } from "~/components/metrics/OperationsFilter";
import { PromptsFilter } from "~/components/metrics/PromptsFilter";
import { ProvidersFilter } from "~/components/metrics/ProvidersFilter";
import { type WidgetData } from "~/components/metrics/QueryWidget";
import { QueuesFilter } from "~/components/metrics/QueuesFilter";
import { ScopeFilter } from "~/components/metrics/ScopeFilter";
import { TitleWidget } from "~/components/metrics/TitleWidget";
import { CreateDashboardPageButton } from "~/components/navigation/DashboardDialogs";
import { Button } from "~/components/primitives/Buttons";
import { NavBar, PageTitle } from "~/components/primitives/PageHeader";
import { TimeFilter } from "~/components/runs/v3/SharedFilters";
import { useEnvironment } from "~/hooks/useEnvironment";
import { useOrganization } from "~/hooks/useOrganizations";
import { useProject } from "~/hooks/useProject";
import { useSearchParams } from "~/hooks/useSearchParam";
import { findProjectBySlug } from "~/models/project.server";
import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
import { getTaskIdentifiers } from "~/models/task.server";
import {
  type BuiltInDashboardFilter,
  type LayoutItem,
  type Widget,
  MetricDashboardPresenter,
} from "~/presenters/v3/MetricDashboardPresenter.server";
import { PromptPresenter } from "~/presenters/v3/PromptPresenter.server";
import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
import { requireUser } from "~/services/session.server";
import { cn } from "~/utils/cn";
import { EnvironmentParamSchema } from "~/utils/pathBuilder";
import { QueryScopeSchema } from "~/v3/querySchemas";
import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route";
import { MetricWidget } from "../resources.metric";

const ParamSchema = EnvironmentParamSchema.extend({
  dashboardKey: z.string(),
});

export const loader = async ({ request, params }: LoaderFunctionArgs) => {
  const user = await requireUser(request);
  const { projectParam, organizationSlug, envParam, dashboardKey } = ParamSchema.parse(params);

  const project = await findProjectBySlug(organizationSlug, projectParam, user.id);
  if (!project) {
    throw new Response(undefined, {
      status: 404,
      statusText: "Project not found",
    });
  }

  const environment = await findEnvironmentBySlug(project.id, envParam, user.id);
  if (!environment) {
    throw new Response(undefined, {
      status: 404,
      statusText: "Environment not found",
    });
  }

  const presenter = new MetricDashboardPresenter();
  const [dashboard, possibleTasks] = await Promise.all([
    presenter.builtInDashboard({
      organizationId: project.organizationId,
      key: dashboardKey,
    }),
    getTaskIdentifiers(environment.id),
  ]);

  const filters = dashboard.filters ?? ["tasks", "queues"];

  const clickhouse = await clickhouseFactory.getClickhouseForOrganization(
    project.organizationId,
    "standard"
  );

  // Load distinct models from ClickHouse if the dashboard has a models filter
  let possibleModels: { model: string; system: string }[] = [];
  if (filters.includes("models")) {
    const queryFn = clickhouse.reader.query({
      name: "getDistinctModels",
      query: `SELECT response_model, any(gen_ai_system) AS gen_ai_system FROM trigger_dev.llm_metrics_v1 WHERE organization_id = {organizationId: String} AND project_id = {projectId: String} AND environment_id = {environmentId: String} AND response_model != '' GROUP BY response_model ORDER BY response_model`,
      params: z.object({
        organizationId: z.string(),
        projectId: z.string(),
        environmentId: z.string(),
      }),
      schema: z.object({ response_model: z.string(), gen_ai_system: z.string() }),
    });
    const [error, rows] = await queryFn({
      organizationId: project.organizationId,
      projectId: project.id,
      environmentId: environment.id,
    });
    if (!error) {
      possibleModels = rows.map((r) => ({ model: r.response_model, system: r.gen_ai_system }));
    }
  }

  const promptPresenter = new PromptPresenter(clickhouse);
  const [possiblePrompts, possibleOperations, possibleProviders] = await Promise.all([
    filters.includes("prompts")
      ? promptPresenter.getDistinctPromptSlugs(project.organizationId, project.id, environment.id)
      : ([] as string[]),
    filters.includes("operations")
      ? promptPresenter.getDistinctOperations(project.organizationId, project.id, environment.id)
      : ([] as string[]),
    filters.includes("providers")
      ? promptPresenter.getDistinctProviders(project.organizationId, project.id, environment.id)
      : ([] as string[]),
  ]);

  return typedjson({
    ...dashboard,
    filters,
    possibleTasks,
    possibleModels,
    possiblePrompts,
    possibleOperations,
    possibleProviders,
  });
};

export default function Page() {
  const {
    key,
    title,
    layout: dashboardLayout,
    defaultPeriod,
    filters,
    possibleTasks,
    possibleModels,
    possiblePrompts,
    possibleOperations,
    possibleProviders,
  } = useTypedLoaderData<typeof loader>();

  const organization = useOrganization();
  const project = useProject();
  const environment = useEnvironment();

  return (
    <PageContainer>
      <NavBar>
        <PageTitle title={title} />
      </NavBar>
      <PageBody scrollable={false}>
        <div className="h-full">
          <MetricDashboard
            key={key}
            layout={dashboardLayout.layout}
            widgets={dashboardLayout.widgets}
            defaultPeriod={defaultPeriod}
            editable={false}
            filters={filters}
            possibleTasks={possibleTasks}
            possibleModels={possibleModels}
            possiblePrompts={possiblePrompts}
            possibleOperations={possibleOperations}
            possibleProviders={possibleProviders}
            filterAccessories={
              <CreateDashboardPageButton
                organization={organization}
                project={project}
                environment={environment}
                shortcut={{ key: "n" }}
              />
            }
          />
        </div>
      </PageBody>
    </PageContainer>
  );
}

export function MetricDashboard({
  layout,
  widgets,
  defaultPeriod,
  editable,
  filters: filterConfig,
  possibleTasks,
  possibleModels,
  possiblePrompts,
  possibleOperations,
  possibleProviders,
  onLayoutChange,
  onEditWidget,
  onRenameWidget,
  onDeleteWidget,
  onDuplicateWidget,
  filterAccessories,
}: {
  /** The layout items (positions/sizes) - fully controlled from parent */
  layout: LayoutItem[];
  /** The widget configurations keyed by widget ID - fully controlled from parent */
  widgets: Record<string, Widget>;
  defaultPeriod: string;
  editable: boolean;
  /** Which filters to show. Defaults to ["tasks", "queues"]. */
  filters?: BuiltInDashboardFilter[];
  /** Possible tasks for filtering */
  possibleTasks?: {
    slug: string;
    triggerSource: TaskTriggerSource;
    isInLatestDeployment: boolean;
  }[];
  /** Possible models for filtering */
  possibleModels?: ModelOption[];
  /** Possible prompt slugs for filtering */
  possiblePrompts?: string[];
  /** Possible operations for filtering */
  possibleOperations?: string[];
  /** Possible providers for filtering */
  possibleProviders?: string[];
  onLayoutChange?: (layout: LayoutItem[]) => void;
  onEditWidget?: (widgetId: string, widget: WidgetData) => void;
  onRenameWidget?: (widgetId: string, newTitle: string) => void;
  onDeleteWidget?: (widgetId: string) => void;
  onDuplicateWidget?: (widgetId: string, widget: WidgetData) => void;
  filterAccessories?: ReactNode;
}) {
  const { value, values } = useSearchParams();
  const { width, containerRef, mounted } = useContainerWidth();
  const [resizingItemId, setResizingItemId] = useState<string | null>(null);
  const [isDragging, setIsDragging] = useState(false);
  const isInteracting = resizingItemId !== null || isDragging;

  const organization = useOrganization();
  const project = useProject();
  const environment = useEnvironment();

  const plan = useCurrentPlan();
  const maxPeriodDays = plan?.v3Subscription?.plan?.limits?.queryPeriodDays?.number;

  const period = value("period");
  const from = value("from");
  const to = value("to");
  const parsedScope = QueryScopeSchema.safeParse(value("scope") ?? "environment");
  const scope = parsedScope.success ? parsedScope.data : "environment";
  const tasks = values("tasks").filter((v) => v !== "");
  const queues = values("queues").filter((v) => v !== "");
  const models = values("models").filter((v) => v !== "");
  const prompts = values("prompts").filter((v) => v !== "");
  const operations = values("operations").filter((v) => v !== "");
  const providers = values("providers").filter((v) => v !== "");

  const activeFilters = filterConfig ?? ["tasks", "queues"];
  const hasAppliedFilters =
    tasks.length > 0 ||
    queues.length > 0 ||
    models.length > 0 ||
    prompts.length > 0 ||
    operations.length > 0 ||
    providers.length > 0;

  const handleLayoutChange = useCallback(
    (newLayout: readonly LayoutItem[]) => {
      const mutableLayout = [...newLayout];
      onLayoutChange?.(mutableLayout);
    },
    [onLayoutChange]
  );

  // Apply constraints for title widgets: fixed height of 2, allow horizontal resize only
  const constrainedLayout = useMemo(
    () =>
      layout.map((item) => {
        const widget = widgets[item.i];
        if (widget?.display.type === "title") {
          return { ...item, h: 2, minH: 2, maxH: 2 };
        }
        return item;
      }),
    [layout, widgets]
  );

  return (
    <div className="grid max-h-full grid-rows-[auto_1fr] overflow-hidden">
      <div className="flex items-center justify-between gap-x-2 border-b border-b-grid-bright py-2 pl-2 pr-3">
        <div className="flex flex-wrap items-center gap-x-1.5 gap-y-1">
          <ScopeFilter shortcut={{ key: "s" }} />
          {activeFilters.includes("tasks") && (
            <LogsTaskFilter possibleTasks={possibleTasks ?? []} />
          )}
          {activeFilters.includes("queues") && <QueuesFilter />}
          {activeFilters.includes("models") && (
            <ModelsFilter possibleModels={possibleModels ?? []} />
          )}
          {activeFilters.includes("prompts") && (
            <PromptsFilter possiblePrompts={possiblePrompts ?? []} />
          )}
          {activeFilters.includes("operations") && (
            <OperationsFilter possibleOperations={possibleOperations ?? []} />
          )}
          {activeFilters.includes("providers") && (
            <ProvidersFilter possibleProviders={possibleProviders ?? []} />
          )}
          <TimeFilter
            defaultPeriod={defaultPeriod}
            labelName="Period"
            hideLabel
            maxPeriodDays={maxPeriodDays}
            valueClassName="text-text-bright"
            shortcut={{ key: "d" }}
          />
          {hasAppliedFilters && (
            <Form className="-ml-1 h-6">
              <Button
                variant="minimal/small"
                LeadingIcon={XMarkIcon}
                tooltip="Clear all filters"
                className="group-hover/button:bg-transparent"
                leadingIconClassName="group-hover/button:text-text-bright"
              />
            </Form>
          )}
        </div>
        {filterAccessories && <div className="flex shrink-0 items-center">{filterAccessories}</div>}
      </div>
      <div
        ref={containerRef}
        className={cn(
          "overflow-y-auto scrollbar-thin scrollbar-track-background-bright scrollbar-thumb-background-raised",
          isInteracting && "select-none"
        )}
      >
        {mounted && (
          <ReactGridLayout
            layout={constrainedLayout}
            width={width}
            gridConfig={{ cols: 12, rowHeight: 30 }}
            resizeConfig={{
              enabled: editable,
              handles: ["se"],
            }}
            dragConfig={{ enabled: editable, handle: ".drag-handle" }}
            onLayoutChange={handleLayoutChange}
            onResizeStart={(_layout, oldItem) => setResizingItemId(oldItem?.i ?? null)}
            onResizeStop={() => setResizingItemId(null)}
            onDragStart={() => setIsDragging(true)}
            onDragStop={() => setIsDragging(false)}
          >
            {Object.entries(widgets).map(([key, widget]) => (
              <div key={key}>
                {widget.display.type === "title" ? (
                  <TitleWidget
                    title={widget.title}
                    isDraggable={editable}
                    isResizing={resizingItemId === key}
                    onRename={
                      onRenameWidget ? (newTitle) => onRenameWidget(key, newTitle) : undefined
                    }
                    onDelete={onDeleteWidget ? () => onDeleteWidget(key) : undefined}
                  />
                ) : (
                  <MetricWidget
                    widgetKey={key}
                    title={widget.title}
                    query={widget.query}
                    scope={scope}
                    period={period ?? defaultPeriod}
                    from={from ?? null}
                    to={to ?? null}
                    taskIdentifiers={tasks.length > 0 ? tasks : undefined}
                    queues={queues.length > 0 ? queues : undefined}
                    responseModels={models.length > 0 ? models : undefined}
                    promptSlugs={prompts.length > 0 ? prompts : undefined}
                    operations={operations.length > 0 ? operations : undefined}
                    providers={providers.length > 0 ? providers : undefined}
                    config={widget.display}
                    organizationId={organization.id}
                    projectId={project.id}
                    environmentId={environment.id}
                    refreshIntervalMs={60_000}
                    isResizing={resizingItemId === key}
                    isDraggable={editable}
                    onEdit={
                      onEditWidget
                        ? (resultData) => onEditWidget(key, { ...widget, resultData })
                        : undefined
                    }
                    onRename={
                      onRenameWidget ? (newTitle) => onRenameWidget(key, newTitle) : undefined
                    }
                    onDelete={onDeleteWidget ? () => onDeleteWidget(key) : undefined}
                    onDuplicate={
                      onDuplicateWidget
                        ? (resultData) => onDuplicateWidget(key, { ...widget, resultData })
                        : undefined
                    }
                  />
                )}
              </div>
            ))}
          </ReactGridLayout>
        )}
      </div>
    </div>
  );
}

function useContainerWidth(initialWidth = 1280) {
  const containerRef = useRef<HTMLDivElement>(null);
  const [width, setWidth] = useState(initialWidth);
  const [mounted, setMounted] = useState(false);

  const measureWidth = useCallback(() => {
    if (containerRef.current) {
      setWidth(containerRef.current.offsetWidth);
    }
  }, []);

  useEffect(() => {
    measureWidth();
    setMounted(true);

    const element = containerRef.current;
    if (!element) return;

    const resizeObserver = new ResizeObserver((entries) => {
      for (const entry of entries) {
        setWidth(entry.contentRect.width);
      }
    });

    resizeObserver.observe(element);
    return () => resizeObserver.disconnect();
  }, [measureWidth]);

  return { width, containerRef, mounted };
}