import { useMemo } from "react";
import { useTranslation } from "react-i18next";
import { useController, useFormContext } from "react-hook-form";
import { TABS_WORKSPACE_NAMESPACE } from "@/i18n/resources";
import { useGetAgentsSuspenseInfinite } from "@/api/agent";
import type { AgentBrief } from "@/api/generated/schemas";
import { AsyncBoundary } from "@/components/custom/AsyncBoundary";
import {
  SelectDialog,
  SelectDialogContent,
  SelectDialogEmpty,
  SelectDialogFooter,
  SelectDialogGroup,
  SelectDialogItem,
  SelectDialogList,
  SelectDialogSearch,
  SelectDialogSkeleton,
  SelectDialogTrigger,
} from "@/components/custom/dialog/SelectDialog";
import { InfiniteScroll } from "@/components/custom/InfiniteScroll";
import { FieldItem } from "@/components/custom/item/FieldItem";
import { Button } from "@/components/ui/button";
import { Item, ItemContent, ItemTitle } from "@/components/ui/item";
import { PAGINATED_QUERY_DEFAULT_OPTIONS } from "@/constants/query-options";
import { useWorkspaceStore } from "@/stores/workspace-store";
import type {
  WorkspaceCreateFormValues,
  WorkspaceEditFormValues,
} from "../form-types";

function AgentQueryList() {
  const query = useGetAgentsSuspenseInfinite(undefined, {
    query: PAGINATED_QUERY_DEFAULT_OPTIONS,
  });
  return (
    <InfiniteScroll
      query={query}
      selectItems={(page) => page.items}
      itemRender={(agent) => (
        <SelectDialogItem<number>
          key={agent.id}
          value={agent.id}
          keywords={[agent.name]}
        >
          {agent.name}
        </SelectDialogItem>
      )}
    />
  );
}

function AgentSelectedList({
  selectedAgentIds,
}: {
  selectedAgentIds: number[];
}) {
  const usableAgents = useWorkspaceStore(
    (state) => state.current?.usable_agents,
  );
  const { data } = useGetAgentsSuspenseInfinite(undefined, {
    query: PAGINATED_QUERY_DEFAULT_OPTIONS,
  });
  const usableAgentMap = useMemo(() => {
    const map = new Map<number, AgentBrief>();
    for (const agent of usableAgents ?? []) {
      map.set(agent.id, agent);
    }
    return map;
  }, [usableAgents]);
  const fetchedAgentMap = useMemo(() => {
    if (!data) {
      return new Map<number, AgentBrief>();
    }
    const map = new Map<number, AgentBrief>();
    for (const agent of data.pages.flatMap((page) => page.items)) {
      map.set(agent.id, agent);
    }
    return map;
  }, [data]);

  const selectedAgents = useMemo(() => {
    const agents: AgentBrief[] = [];
    for (const id of selectedAgentIds) {
      const agent1 = usableAgentMap.get(id);
      if (agent1) {
        agents.push(agent1);
        continue;
      }
      const agent2 = fetchedAgentMap.get(id);
      if (agent2) {
        agents.push(agent2);
      }
    }
    return agents;
  }, [selectedAgentIds, usableAgentMap, fetchedAgentMap]);
  return (
    <div className="mt-2 space-y-2">
      {selectedAgents.map((agent) => (
        <Item key={agent.id} variant="outline" size="sm">
          <ItemContent>
            <ItemTitle>{agent.name}</ItemTitle>
          </ItemContent>
        </Item>
      ))}
    </div>
  );
}

export function AgentMultiSelectField() {
  const { t } = useTranslation(TABS_WORKSPACE_NAMESPACE);
  const { control } = useFormContext<
    WorkspaceCreateFormValues | WorkspaceEditFormValues
  >();
  const {
    field: { value, onChange },
    fieldState,
  } = useController({
    name: "usable_agent_ids",
    control,
  });

  return (
    <div>
      <FieldItem label={t("form.usable_agents.label")} fieldState={fieldState}>
        <SelectDialog<number> multiple value={value}>
          <SelectDialogTrigger>
            <Button type="button" variant="outline">
              {t("form.usable_agents.select")}
            </Button>
          </SelectDialogTrigger>
          <SelectDialogContent>
            <SelectDialogSearch
              placeholder={t("form.usable_agents.search_placeholder")}
            />
            <SelectDialogList>
              <SelectDialogEmpty>
                {t("form.usable_agents.empty")}
              </SelectDialogEmpty>
              <SelectDialogGroup>
                <AsyncBoundary skeleton={<SelectDialogSkeleton />}>
                  <AgentQueryList />
                </AsyncBoundary>
              </SelectDialogGroup>
            </SelectDialogList>
            <SelectDialogFooter
              onConfirm={onChange}
              confirmText={t("form.usable_agents.confirm")}
              cancelText={t("form.usable_agents.cancel")}
            />
          </SelectDialogContent>
        </SelectDialog>
      </FieldItem>

      <AsyncBoundary>
        <AgentSelectedList selectedAgentIds={value} />
      </AsyncBoundary>
    </div>
  );
}
