import {
  createPaginatedResponseSchema,
  PaginationQuerySchema,
  providerRequiresPerUserCredential,
  RouteId,
  type SupportedProvider,
  SupportedProvidersSchema,
} from "@archestra/shared";
import type { FastifyPluginAsyncZod } from "fastify-type-provider-zod";
import { z } from "zod";
import { userHasPermission } from "@/auth";
import {
  LlmProviderApiKeyModel,
  MemberModel,
  TeamModel,
  VirtualApiKeyModel,
} from "@/models";
import {
  ApiError,
  constructResponseSchema,
  type ResourceVisibilityScope,
  ResourceVisibilityScopeSchema,
  type User,
  VirtualApiKeyTypeSchema,
  VirtualApiKeyWithParentInfoSchema,
  VirtualApiKeyWithValueSchema,
} from "@/types";

const UpdateVirtualApiKeyResponseSchema = VirtualApiKeyWithValueSchema.omit({
  value: true,
});

const VirtualApiKeyBodyObjectSchema = z.object({
  name: z.string().min(1, "Name is required").max(256),
  keyType: VirtualApiKeyTypeSchema.default("standard"),
  expiresAt: z.coerce.date().nullable().optional(),
  scope: ResourceVisibilityScopeSchema.default("org"),
  teams: z.array(z.string()).default([]),
  providerApiKeys: z
    .array(
      z.object({
        provider: SupportedProvidersSchema,
        providerApiKeyId: z.string().uuid(),
      }),
    )
    .default([]),
});

/**
 * Contextual validation: which fields are accepted depends on the key type.
 * Standard keys map provider API keys; passthrough keys never carry provider
 * credentials or team/org scope.
 */
function refineVirtualApiKeyBody(
  value: z.infer<typeof VirtualApiKeyBodyObjectSchema>,
  ctx: z.RefinementCtx,
): void {
  if (value.keyType === "passthrough") {
    if (value.providerApiKeys.length > 0) {
      ctx.addIssue({
        code: "custom",
        path: ["providerApiKeys"],
        message: "Passthrough virtual keys cannot map provider API keys",
      });
    }
    if (value.teams.length > 0) {
      ctx.addIssue({
        code: "custom",
        path: ["teams"],
        message: "Passthrough virtual keys cannot be assigned to teams",
      });
    }
    return;
  }

  if (value.providerApiKeys.length === 0) {
    ctx.addIssue({
      code: "custom",
      path: ["providerApiKeys"],
      message: "At least one provider API key is required",
    });
  }
}

const CreateOrUpdateVirtualApiKeyBodySchema =
  VirtualApiKeyBodyObjectSchema.superRefine(refineVirtualApiKeyBody);

const CreateVirtualApiKeyBodySchema = VirtualApiKeyBodyObjectSchema.extend({
  /**
   * Owner the key is created on behalf of. Defaults to the creator. Setting
   * it to a different user requires llmVirtualKey:admin and that user must
   * belong to the organization.
   */
  ownerId: z.string().optional(),
}).superRefine(refineVirtualApiKeyBody);

const virtualApiKeysRoutes: FastifyPluginAsyncZod = async (fastify) => {
  fastify.get(
    "/api/llm-virtual-keys",
    {
      schema: {
        operationId: RouteId.GetAllVirtualApiKeys,
        description:
          "Get virtual API keys visible to the current user, with provider key mappings",
        tags: ["Virtual API Keys"],
        querystring: PaginationQuerySchema.extend({
          search: z.string().trim().min(1).optional(),
          providerApiKeyId: z.string().uuid().optional(),
          keyType: VirtualApiKeyTypeSchema.optional(),
        }),
        response: constructResponseSchema(
          createPaginatedResponseSchema(VirtualApiKeyWithParentInfoSchema),
        ),
      },
    },
    async (
      {
        query: { limit, offset, search, providerApiKeyId, keyType },
        organizationId,
        user,
      },
      reply,
    ) => {
      const [userTeamIds, isVirtualKeyAdmin] = await Promise.all([
        TeamModel.getUserTeamIds(user.id),
        userHasPermission(user.id, organizationId, "llmVirtualKey", "admin"),
      ]);

      const result = await VirtualApiKeyModel.findAllByOrganization({
        organizationId,
        pagination: { limit, offset },
        userId: user.id,
        userTeamIds,
        isAdmin: isVirtualKeyAdmin,
        search,
        providerApiKeyId,
        keyType,
      });
      return reply.send(result);
    },
  );

  fastify.post(
    "/api/llm-virtual-keys",
    {
      schema: {
        operationId: RouteId.CreateVirtualApiKey,
        description:
          "Create a new virtual API key. Returns the full token value once.",
        tags: ["Virtual API Keys"],
        body: CreateVirtualApiKeyBodySchema,
        response: constructResponseSchema(VirtualApiKeyWithValueSchema),
      },
    },
    async ({ body, organizationId, user }, reply) => {
      const response = await createVirtualApiKey({
        body,
        organizationId,
        user,
      });
      return reply.send(response);
    },
  );

  fastify.patch(
    "/api/llm-virtual-keys/:id",
    {
      schema: {
        operationId: RouteId.UpdateVirtualApiKey,
        description: "Update a virtual API key",
        tags: ["Virtual API Keys"],
        params: z.object({
          id: z.string().uuid(),
        }),
        body: CreateOrUpdateVirtualApiKeyBodySchema,
        response: constructResponseSchema(UpdateVirtualApiKeyResponseSchema),
      },
    },
    async ({ params, body, organizationId, user }, reply) => {
      const response = await updateVirtualApiKey({
        id: params.id,
        body,
        organizationId,
        user,
      });
      return reply.send(response);
    },
  );

  fastify.delete(
    "/api/llm-virtual-keys/:id",
    {
      schema: {
        operationId: RouteId.DeleteVirtualApiKey,
        description: "Delete a virtual API key",
        tags: ["Virtual API Keys"],
        params: z.object({
          id: z.string().uuid(),
        }),
        response: constructResponseSchema(z.object({ success: z.boolean() })),
      },
    },
    async ({ params, organizationId, user }, reply) => {
      const response = await deleteVirtualApiKey({
        id: params.id,
        organizationId,
        user,
      });
      return reply.send(response);
    },
  );
};

export default virtualApiKeysRoutes;

async function createVirtualApiKey(params: {
  body: z.infer<typeof CreateVirtualApiKeyBodySchema>;
  organizationId: string;
  user: User;
}): Promise<z.infer<typeof VirtualApiKeyWithValueSchema>> {
  const { body, organizationId, user } = params;

  if (body.expiresAt && body.expiresAt <= new Date()) {
    throw new ApiError(400, "Expiration date must be in the future");
  }

  const [userTeamIds, isVirtualKeyAdmin] = await Promise.all([
    TeamModel.getUserTeamIds(user.id),
    userHasPermission(user.id, organizationId, "llmVirtualKey", "admin"),
  ]);
  const ownerId = await resolveKeyOwner({
    requestedOwnerId: body.ownerId,
    creatorId: user.id,
    organizationId,
    isAdmin: isVirtualKeyAdmin,
  });

  // Passthrough keys are always personal and carry no provider keys; they only
  // authenticate the acting user.
  if (body.keyType === "passthrough") {
    const created = await VirtualApiKeyModel.create({
      organizationId,
      name: body.name,
      keyType: "passthrough",
      expiresAt: body.expiresAt ?? null,
      scope: "personal",
      authorId: ownerId,
    });

    return {
      ...created.virtualKey,
      value: created.value,
      teams: created.teams,
      authorName: created.authorName,
      providerApiKeys: created.providerApiKeys,
    };
  }

  await validateVirtualKeyScope({
    scope: body.scope,
    teamIds: body.teams,
    userId: user.id,
    organizationId,
    userTeamIds,
    isAdmin: isVirtualKeyAdmin,
  });
  await validateProviderApiKeys({
    mappings: body.providerApiKeys,
    organizationId,
    scope: body.scope,
    userId: user.id,
  });

  const { virtualKey, value, teams, authorName, providerApiKeys } =
    await VirtualApiKeyModel.create({
      organizationId,
      name: body.name,
      keyType: "standard",
      expiresAt: body.expiresAt ?? null,
      scope: body.scope,
      authorId: ownerId,
      teamIds: body.teams,
      providerApiKeys: body.providerApiKeys,
    });

  return {
    ...virtualKey,
    value,
    teams,
    authorName,
    providerApiKeys,
  };
}

async function updateVirtualApiKey(params: {
  id: string;
  body: z.infer<typeof CreateOrUpdateVirtualApiKeyBodySchema>;
  organizationId: string;
  user: User;
}): Promise<z.infer<typeof UpdateVirtualApiKeyResponseSchema>> {
  const { id, body, organizationId, user } = params;

  const accessContext = await VirtualApiKeyModel.findAccessContextById(id);

  if (!accessContext || accessContext.organizationId !== organizationId) {
    throw new ApiError(404, "Virtual API key not found");
  }

  if (body.expiresAt && body.expiresAt <= new Date()) {
    throw new ApiError(400, "Expiration date must be in the future");
  }

  const [userTeamIds, isVirtualKeyAdmin] = await Promise.all([
    TeamModel.getUserTeamIds(user.id),
    userHasPermission(user.id, organizationId, "llmVirtualKey", "admin"),
  ]);
  await requireVirtualKeyModifyPermission({
    virtualKey: accessContext,
    userId: user.id,
    organizationId,
    userTeamIds,
  });

  // The key type is fixed at creation; only its own configuration is editable.
  if (body.keyType !== accessContext.keyType) {
    throw new ApiError(400, "Virtual key type cannot be changed");
  }

  let updatedVirtualKey: Awaited<ReturnType<typeof VirtualApiKeyModel.update>>;
  if (accessContext.keyType === "passthrough") {
    updatedVirtualKey = await VirtualApiKeyModel.update({
      id,
      name: body.name,
      expiresAt: body.expiresAt ?? null,
      scope: "personal",
      // Preserve the key's owner; an edit must not transfer it to the editor.
      authorId: accessContext.authorId,
      teamIds: [],
      providerApiKeys: [],
    });
  } else {
    await validateVirtualKeyScope({
      scope: body.scope,
      teamIds: body.teams,
      userId: user.id,
      organizationId,
      userTeamIds,
      isAdmin: isVirtualKeyAdmin,
    });
    await validateProviderApiKeys({
      mappings: body.providerApiKeys,
      organizationId,
      scope: body.scope,
      userId: user.id,
    });
    updatedVirtualKey = await VirtualApiKeyModel.update({
      id,
      name: body.name,
      expiresAt: body.expiresAt ?? null,
      scope: body.scope,
      // Preserve the key's owner; an edit must not transfer it to the editor
      // (e.g. an admin editing a key minted on behalf of another user).
      authorId: accessContext.authorId,
      teamIds: body.teams,
      providerApiKeys: body.providerApiKeys,
    });
  }

  if (!updatedVirtualKey) {
    throw new ApiError(404, "Virtual API key not found");
  }

  const visibilityMetadata =
    await VirtualApiKeyModel.getVisibilityForVirtualApiKeyIds([id]);
  const providerApiKeys = await VirtualApiKeyModel.getProviderApiKeys(id);

  return {
    ...updatedVirtualKey,
    teams: visibilityMetadata.teams.get(id) ?? [],
    authorName: visibilityMetadata.authorName.get(id) ?? null,
    providerApiKeys,
  };
}

async function deleteVirtualApiKey(params: {
  id: string;
  organizationId: string;
  user: User;
}): Promise<{ success: boolean }> {
  const { id, organizationId, user } = params;

  const accessContext = await VirtualApiKeyModel.findAccessContextById(id);

  if (!accessContext || accessContext.organizationId !== organizationId) {
    throw new ApiError(404, "Virtual API key not found");
  }

  const userTeamIds = await TeamModel.getUserTeamIds(user.id);
  await requireVirtualKeyModifyPermission({
    virtualKey: accessContext,
    userId: user.id,
    organizationId,
    userTeamIds,
  });

  await VirtualApiKeyModel.delete(id);
  return { success: true };
}

async function resolveKeyOwner(params: {
  requestedOwnerId: string | undefined;
  creatorId: string;
  organizationId: string;
  isAdmin: boolean;
}): Promise<string> {
  const { requestedOwnerId, creatorId, organizationId, isAdmin } = params;

  if (!requestedOwnerId || requestedOwnerId === creatorId) {
    return creatorId;
  }

  if (!isAdmin) {
    throw new ApiError(
      403,
      "You need llmVirtualKey:admin permission to create a virtual key for another user",
    );
  }

  const member = await MemberModel.getByUserId(
    requestedOwnerId,
    organizationId,
  );
  if (!member) {
    throw new ApiError(404, "User is not a member of this organization");
  }

  return requestedOwnerId;
}

async function validateVirtualKeyScope(params: {
  scope: ResourceVisibilityScope;
  teamIds: string[];
  userId: string;
  organizationId: string;
  userTeamIds: string[];
  isAdmin: boolean;
}): Promise<void> {
  const { scope, teamIds, userTeamIds, isAdmin } = params;

  if (scope !== "team" && teamIds.length > 0) {
    throw new ApiError(400, "Teams can only be assigned to team-scoped keys");
  }

  if (scope === "team" && teamIds.length === 0) {
    throw new ApiError(400, "At least one team is required for team scope");
  }

  if (scope === "org") {
    if (!isAdmin) {
      throw new ApiError(
        403,
        "You need llmVirtualKey:admin permission to create org-scoped virtual keys",
      );
    }
    return;
  }

  if (scope === "team") {
    const uniqueTeamIds = [...new Set(teamIds)];
    const teams = await TeamModel.findByIds(uniqueTeamIds);
    if (teams.length !== uniqueTeamIds.length) {
      throw new ApiError(400, "One or more selected teams do not exist");
    }

    if (isAdmin) {
      return;
    }

    const userTeamIdSet = new Set(userTeamIds);
    const canManageAllTeams = uniqueTeamIds.every((teamId) =>
      userTeamIdSet.has(teamId),
    );
    if (!canManageAllTeams) {
      throw new ApiError(
        403,
        "You can only assign virtual keys to teams you are a member of",
      );
    }
  }
}

async function validateProviderApiKeys(params: {
  mappings: Array<{ provider: SupportedProvider; providerApiKeyId: string }>;
  organizationId: string;
  scope: ResourceVisibilityScope;
  userId: string;
}): Promise<void> {
  const { mappings, organizationId, scope, userId } = params;
  if (mappings.length === 0) {
    return;
  }

  const providers = new Set<SupportedProvider>();
  const apiKeys = await LlmProviderApiKeyModel.findByIds(
    mappings.map((mapping) => mapping.providerApiKeyId),
  );
  const apiKeysById = new Map(apiKeys.map((apiKey) => [apiKey.id, apiKey]));

  for (const mapping of mappings) {
    if (providers.has(mapping.provider)) {
      throw new ApiError(
        400,
        `Only one provider API key can be mapped for provider "${mapping.provider}".`,
      );
    }
    providers.add(mapping.provider);

    const apiKey = apiKeysById.get(mapping.providerApiKeyId);
    if (!apiKey || apiKey.organizationId !== organizationId) {
      throw new ApiError(404, "LLM provider API key not found");
    }
    if (apiKey.provider !== mapping.provider) {
      throw new ApiError(
        400,
        `Provider API key "${apiKey.name}" is for provider "${apiKey.provider}", not "${mapping.provider}".`,
      );
    }

    // Per-user-credential providers (GitHub Copilot) are individual tokens. A
    // virtual key is itself a shareable secret, so it may only wrap a per-user
    // key when it is the user's OWN personal key in their OWN personal virtual
    // key, and never bundled with other providers (a shared model-router key
    // would expose one user's token to everyone routing through it).
    if (providerRequiresPerUserCredential(mapping.provider)) {
      if (scope !== "personal" || mappings.length > 1) {
        throw new ApiError(
          400,
          `${mapping.provider} is per-user: it can only be wrapped in your own personal virtual key on its own, not in a shared or multi-provider (model-router) key.`,
        );
      }
      if (apiKey.scope !== "personal" || apiKey.userId !== userId) {
        throw new ApiError(
          403,
          `You can only map your own personal ${mapping.provider} key.`,
        );
      }
    }
  }
}

async function requireVirtualKeyModifyPermission(params: {
  virtualKey: {
    scope: ResourceVisibilityScope;
    authorId: string | null;
    teamIds: string[];
  };
  userId: string;
  organizationId: string;
  userTeamIds: string[];
}): Promise<void> {
  const { virtualKey, userId, organizationId, userTeamIds } = params;

  const isAdmin = await userHasPermission(
    userId,
    organizationId,
    "llmVirtualKey",
    "admin",
  );
  if (isAdmin) {
    return;
  }

  switch (virtualKey.scope) {
    case "org":
      throw new ApiError(
        403,
        "Only llmVirtualKey:admin users can manage org-scoped virtual keys",
      );
    case "team": {
      const userTeamIdSet = new Set(userTeamIds);
      const isMemberOfAnyTeam = virtualKey.teamIds.some((teamId) =>
        userTeamIdSet.has(teamId),
      );
      if (!isMemberOfAnyTeam) {
        throw new ApiError(
          403,
          "You can only manage virtual keys in teams you are a member of",
        );
      }
      return;
    }
    case "personal":
      if (virtualKey.authorId !== userId) {
        throw new ApiError(
          403,
          "You can only manage your own personal virtual keys",
        );
      }
      return;
  }
}
