import { createLanguageModelFromResolved } from "@tasks/ai";
import { seedDefaultRoleRoutes } from "@tasks/database/model-routing";
import {
  modelProviders,
  roleModelRoutes,
} from "@tasks/database/schema";
import {
  convoyModelPolicySchema,
  providerTypeSchema,
  resolveModel,
  type ModelProvider,
  type RoleModelRoute,
} from "@tasks/shared/model-routing";
import { TRPCError } from "@trpc/server";
import { generateText } from "ai";
import { and, asc, eq } from "drizzle-orm";
import { z } from "zod";

import { rateLimit } from "@/server/lib/rate-limit";
import { router, workspaceProcedure } from "@/server/trpc";
import type { WorkspaceProcedureContext } from "@/server/trpc";

type ProviderRow = typeof modelProviders.$inferSelect;
type RouteRow = typeof roleModelRoutes.$inferSelect;

function assertWorkspaceAdmin(ctx: WorkspaceProcedureContext): void {
  if (ctx.workspace.role !== "owner" && ctx.workspace.role !== "admin") {
    throw new TRPCError({
      code: "FORBIDDEN",
      message: "Only workspace owners and admins can manage model providers.",
    });
  }
}

function canManageWorkspace(ctx: WorkspaceProcedureContext): boolean {
  return ctx.workspace.role === "owner" || ctx.workspace.role === "admin";
}

function mapProviderRow(row: ProviderRow): ModelProvider {
  return {
    id: row.id,
    workspaceId: row.workspaceId,
    alias: row.alias,
    displayName: row.displayName,
    providerType: providerTypeSchema.parse(row.providerType),
    modelId: row.modelId,
    baseUrl: row.baseUrl,
    apiKeyEnvVar: row.apiKeyEnvVar,
    extraHeaders: row.extraHeaders,
    defaultParams: row.defaultParams,
    metadata: row.metadata,
    enabled: row.enabled,
  };
}

function mapRouteRow(row: RouteRow): RoleModelRoute {
  return {
    id: row.id,
    workspaceId: row.workspaceId,
    roleId: row.roleId,
    providerAliasId: row.providerAliasId,
    enabled: row.enabled,
    sortOrder: row.sortOrder,
  };
}

function readApiKey(envVarName: string): string | undefined {
  return process.env[envVarName]?.trim() || undefined;
}

function validateProviderFields(input: {
  providerType: z.infer<typeof providerTypeSchema>;
  baseUrl?: string | null;
}): void {
  if (input.providerType === "openai_compatible" && !input.baseUrl?.trim()) {
    throw new TRPCError({
      code: "BAD_REQUEST",
      message: "baseUrl is required for openai_compatible providers.",
    });
  }
}

function providerToResolved(row: ProviderRow): ReturnType<typeof resolveModel> {
  const provider = mapProviderRow(row);
  const result = resolveModel({
    roleId: "test",
    providers: [provider],
    routes: [],
    briefRecommendedModel: provider.alias,
    getApiKey: readApiKey,
  });
  return result;
}

const providerFieldsSchema = z.object({
  alias: z.string().min(1).max(120),
  displayName: z.string().min(1).max(200),
  providerType: providerTypeSchema,
  modelId: z.string().min(1).max(200),
  baseUrl: z.string().url().nullable().optional(),
  apiKeyEnvVar: z.string().min(1).max(120).nullable().optional(),
  extraHeaders: z.record(z.string()).nullable().optional(),
  defaultParams: z
    .object({
      temperature: z.number().min(0).max(2).optional(),
      maxTokens: z.number().int().positive().optional(),
    })
    .nullable()
    .optional(),
  metadata: z.record(z.unknown()).nullable().optional(),
  enabled: z.boolean().optional(),
});

export const aiModelsRouter = router({
  listProviders: workspaceProcedure.query(async ({ ctx }) => {
    const rows = await ctx.db
      .select()
      .from(modelProviders)
      .where(eq(modelProviders.workspaceId, ctx.workspace.id))
      .orderBy(asc(modelProviders.alias));

    return {
      providers: rows.map(mapProviderRow),
      canManage: canManageWorkspace(ctx),
    };
  }),

  createProvider: workspaceProcedure
    .input(providerFieldsSchema)
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);
      validateProviderFields(input);

      try {
        const [created] = await ctx.db
          .insert(modelProviders)
          .values({
            workspaceId: ctx.workspace.id,
            alias: input.alias.trim(),
            displayName: input.displayName.trim(),
            providerType: input.providerType,
            modelId: input.modelId.trim(),
            baseUrl: input.baseUrl?.trim() || null,
            apiKeyEnvVar: input.apiKeyEnvVar?.trim() || null,
            extraHeaders: input.extraHeaders ?? null,
            defaultParams: input.defaultParams ?? null,
            metadata: input.metadata ?? null,
            enabled: input.enabled ?? true,
            updatedAt: new Date(),
          })
          .returning();

        if (!created) {
          throw new TRPCError({
            code: "INTERNAL_SERVER_ERROR",
            message: "Failed to create provider.",
          });
        }

        return mapProviderRow(created);
      } catch (e) {
        if (isUniqueViolation(e)) {
          throw new TRPCError({
            code: "CONFLICT",
            message: `Alias "${input.alias}" already exists in this workspace.`,
          });
        }
        throw e;
      }
    }),

  updateProvider: workspaceProcedure
    .input(
      providerFieldsSchema.partial().extend({
        id: z.string().uuid(),
      }),
    )
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);

      const [existing] = await ctx.db
        .select()
        .from(modelProviders)
        .where(
          and(
            eq(modelProviders.id, input.id),
            eq(modelProviders.workspaceId, ctx.workspace.id),
          ),
        )
        .limit(1);

      if (!existing) {
        throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
      }

      const nextType = input.providerType ?? existing.providerType;
      const nextBaseUrl =
        input.baseUrl !== undefined ? input.baseUrl : existing.baseUrl;
      validateProviderFields({
        providerType: providerTypeSchema.parse(nextType),
        baseUrl: nextBaseUrl,
      });

      try {
        const [updated] = await ctx.db
          .update(modelProviders)
          .set({
            ...(input.alias !== undefined ? { alias: input.alias.trim() } : {}),
            ...(input.displayName !== undefined
              ? { displayName: input.displayName.trim() }
              : {}),
            ...(input.providerType !== undefined
              ? { providerType: input.providerType }
              : {}),
            ...(input.modelId !== undefined ? { modelId: input.modelId.trim() } : {}),
            ...(input.baseUrl !== undefined
              ? { baseUrl: input.baseUrl?.trim() || null }
              : {}),
            ...(input.apiKeyEnvVar !== undefined
              ? { apiKeyEnvVar: input.apiKeyEnvVar?.trim() || null }
              : {}),
            ...(input.extraHeaders !== undefined
              ? { extraHeaders: input.extraHeaders }
              : {}),
            ...(input.defaultParams !== undefined
              ? { defaultParams: input.defaultParams }
              : {}),
            ...(input.metadata !== undefined ? { metadata: input.metadata } : {}),
            ...(input.enabled !== undefined ? { enabled: input.enabled } : {}),
            updatedAt: new Date(),
          })
          .where(
            and(
              eq(modelProviders.id, input.id),
              eq(modelProviders.workspaceId, ctx.workspace.id),
            ),
          )
          .returning();

        if (!updated) {
          throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
        }

        return mapProviderRow(updated);
      } catch (e) {
        if (isUniqueViolation(e)) {
          throw new TRPCError({
            code: "CONFLICT",
            message: "Another provider already uses that alias.",
          });
        }
        throw e;
      }
    }),

  deleteProvider: workspaceProcedure
    .input(z.object({ id: z.string().uuid() }))
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);
      const [deleted] = await ctx.db
        .delete(modelProviders)
        .where(
          and(
            eq(modelProviders.id, input.id),
            eq(modelProviders.workspaceId, ctx.workspace.id),
          ),
        )
        .returning({ id: modelProviders.id });

      if (!deleted) {
        throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
      }

      return { id: deleted.id };
    }),

  testProvider: workspaceProcedure
    .input(z.object({ providerId: z.string().uuid() }))
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);

      const limit = rateLimit({
        key: `ai:testProvider:${ctx.workspace.id}`,
        limit: 10,
        windowMs: 60_000,
      });
      if (!limit.allowed) {
        throw new TRPCError({
          code: "TOO_MANY_REQUESTS",
          message: "Test rate limit exceeded. Try again shortly.",
        });
      }

      const [row] = await ctx.db
        .select()
        .from(modelProviders)
        .where(
          and(
            eq(modelProviders.id, input.providerId),
            eq(modelProviders.workspaceId, ctx.workspace.id),
          ),
        )
        .limit(1);

      if (!row) {
        throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
      }

      const resolved = providerToResolved(row);
      if (!resolved.ok) {
        return {
          ok: false as const,
          latencyMs: 0,
          error: resolved.error.message,
        };
      }

      const client = createLanguageModelFromResolved(resolved.model);
      if (!client.ok) {
        return { ok: false as const, latencyMs: 0, error: client.reason };
      }

      const started = Date.now();
      try {
        await generateText({
          model: client.model,
          prompt: "Reply with exactly: OK",
          maxTokens: 5,
          temperature: 0,
        });
        return {
          ok: true as const,
          latencyMs: Date.now() - started,
        };
      } catch (e) {
        const message = e instanceof Error ? e.message : String(e);
        return {
          ok: false as const,
          latencyMs: Date.now() - started,
          error: message,
        };
      }
    }),

  listRoleRoutes: workspaceProcedure.query(async ({ ctx }) => {
    const routes = await ctx.db
      .select({
        route: roleModelRoutes,
        providerAlias: modelProviders.alias,
        providerEnabled: modelProviders.enabled,
        providerMetadata: modelProviders.metadata,
      })
      .from(roleModelRoutes)
      .innerJoin(
        modelProviders,
        eq(roleModelRoutes.providerAliasId, modelProviders.id),
      )
      .where(eq(roleModelRoutes.workspaceId, ctx.workspace.id))
      .orderBy(asc(roleModelRoutes.sortOrder), asc(roleModelRoutes.roleId));

    return {
      routes: routes.map((row) => ({
        ...mapRouteRow(row.route),
        providerAlias: row.providerAlias,
        providerEnabled: row.providerEnabled,
        providerMetadata: row.providerMetadata,
      })),
      canManage: canManageWorkspace(ctx),
    };
  }),

  upsertRoleRoute: workspaceProcedure
    .input(
      z.object({
        roleId: z.string().min(1).max(120),
        providerAliasId: z.string().uuid(),
        enabled: z.boolean().optional(),
        sortOrder: z.number().int().optional(),
      }),
    )
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);

      const [provider] = await ctx.db
        .select({ id: modelProviders.id })
        .from(modelProviders)
        .where(
          and(
            eq(modelProviders.id, input.providerAliasId),
            eq(modelProviders.workspaceId, ctx.workspace.id),
          ),
        )
        .limit(1);

      if (!provider) {
        throw new TRPCError({ code: "NOT_FOUND", message: "Provider not found." });
      }

      const [existing] = await ctx.db
        .select({ id: roleModelRoutes.id })
        .from(roleModelRoutes)
        .where(
          and(
            eq(roleModelRoutes.workspaceId, ctx.workspace.id),
            eq(roleModelRoutes.roleId, input.roleId),
          ),
        )
        .limit(1);

      if (existing) {
        const [updated] = await ctx.db
          .update(roleModelRoutes)
          .set({
            providerAliasId: input.providerAliasId,
            enabled: input.enabled ?? true,
            ...(input.sortOrder !== undefined ? { sortOrder: input.sortOrder } : {}),
            updatedAt: new Date(),
          })
          .where(
            and(
              eq(roleModelRoutes.id, existing.id),
              eq(roleModelRoutes.workspaceId, ctx.workspace.id),
            ),
          )
          .returning();

        if (!updated) {
          throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
        }

        return mapRouteRow(updated);
      }

      const [created] = await ctx.db
        .insert(roleModelRoutes)
        .values({
          workspaceId: ctx.workspace.id,
          roleId: input.roleId,
          providerAliasId: input.providerAliasId,
          enabled: input.enabled ?? true,
          sortOrder: input.sortOrder ?? 0,
          updatedAt: new Date(),
        })
        .returning();

      if (!created) {
        throw new TRPCError({
          code: "INTERNAL_SERVER_ERROR",
          message: "Failed to create role route.",
        });
      }

      return mapRouteRow(created);
    }),

  deleteRoleRoute: workspaceProcedure
    .input(z.object({ id: z.string().uuid() }))
    .mutation(async ({ ctx, input }) => {
      assertWorkspaceAdmin(ctx);
      const [deleted] = await ctx.db
        .delete(roleModelRoutes)
        .where(
          and(
            eq(roleModelRoutes.id, input.id),
            eq(roleModelRoutes.workspaceId, ctx.workspace.id),
          ),
        )
        .returning({ id: roleModelRoutes.id });

      if (!deleted) {
        throw new TRPCError({ code: "NOT_FOUND", message: "Route not found." });
      }

      return { id: deleted.id };
    }),

  resetPipelineDefaults: workspaceProcedure.mutation(async ({ ctx }) => {
    assertWorkspaceAdmin(ctx);
    const result = await seedDefaultRoleRoutes(ctx.db, ctx.workspace.id);
    return result;
  }),

  resolveForRole: workspaceProcedure
    .input(
      z.object({
        roleId: z.string().min(1).max(120),
        briefRecommendedModel: z.string().optional(),
        convoyPolicy: convoyModelPolicySchema.optional(),
      }),
    )
    .query(async ({ ctx, input }) => {
      const providerRows = await ctx.db
        .select()
        .from(modelProviders)
        .where(eq(modelProviders.workspaceId, ctx.workspace.id));

      const routeRows = await ctx.db
        .select()
        .from(roleModelRoutes)
        .where(eq(roleModelRoutes.workspaceId, ctx.workspace.id));

      const result = resolveModel({
        roleId: input.roleId,
        providers: providerRows.map(mapProviderRow),
        routes: routeRows.map(mapRouteRow),
        briefRecommendedModel: input.briefRecommendedModel,
        convoyPolicy: input.convoyPolicy,
        getApiKey: readApiKey,
      });

      if (!result.ok) {
        return result;
      }

      const { apiKey: _removed, ...publicModel } = result.model;
      void _removed;
      return {
        ok: true as const,
        model: publicModel,
        source: result.source,
        resolvedAlias: result.resolvedAlias,
      };
    }),
});

function isUniqueViolation(error: unknown): boolean {
  return (
    typeof error === "object" &&
    error !== null &&
    "code" in error &&
    (error as { code?: string }).code === "23505"
  );
}