import fs from "node:fs/promises";

import path from "node:path";

import { fileURLToPath } from "node:url";

import { loadAllModels } from "../../../src/data/load.js";

import { modelId, type Model, type Parameter } from "../../../src/schema/model.js";

const here = path.dirname(fileURLToPath(import.meta.url));

const PACKAGE_DIR = path.resolve(here, "..");

const PYTHON_PACKAGE_DIR = path.join(PACKAGE_DIR, "src", "modelparams");

const GENERATED_DIR = path.join(PYTHON_PACKAGE_DIR, "_generated");

const TYPES_DIR = path.join(PYTHON_PACKAGE_DIR, "types");

const HEADER =

"# AUTO-GENERATED by packages/modelparams-python/scripts/codegen.ts.\n" +

"# Do not edit by hand; the YAML catalog under /models is the source of truth.\n\n";

function pythonString(value: string): string {

return JSON.stringify(value);

}

function pythonLiteral(value: string | number | boolean | null): string {

if (value === null) return "None";

if (value === true) return "True";

if (value === false) return "False";

if (typeof value === "number" && !Number.isInteger(value)) {

throw new Error(

`Python typing.Literal cannot represent the non-integer enum value ${JSON.stringify(value)}`,

);

}

return typeof value === "string" ? pythonString(value) : String(value);

}

function moduleName(provider: string): string {

return provider.replace(/-/g, "_");

}

function typeName(model: Pick<Model, "authType" | "model">): string {

const suffix = model.authType === "subscription" ? "-subscription" : "";

const words = `${model.model}${suffix}`.split(/[^A-Za-z0-9]+/).filter(Boolean);

const base = words

.map((word) => `${word[0]!.toUpperCase()}${word.slice(1).toLowerCase()}`)

.join("_");

return `${/^\d/.test(base) ? "Model_" : ""}${base}Params`;

}

function numericType(base: "int" | "float", parameter: Parameter): string {

if (parameter.type !== "integer" && parameter.type !== "number") return base;

const constraints: string[] = [];

if (parameter.range?.min !== undefined) constraints.push(`ge=${parameter.range.min}`);

if (parameter.range?.max !== undefined) constraints.push(`le=${parameter.range.max}`);

return constraints.length > 0 ? `Annotated[${base}, Field(${constraints.join(", ")})]` : base;

}

function pythonType(parameter: Parameter): string {

switch (parameter.type) {

case "boolean":

return "bool";

case "integer":

return numericType("int", parameter);

case "number":

return numericType("float", parameter);

case "string":

return "str";

case "enum":

return `Literal[${parameter.values.map(pythonLiteral).join(", ")}]`;

}

}

function emitProviderTypes(provider: string, models: Model[]): string {

const entries = models.map((model) => {

const name = typeName(model);

const fields = model.params

.map((parameter) => ` ${pythonString(parameter.path)}: ${pythonType(parameter)},`)

.join("\n");

return (

`${name} = TypedDict(\n` +

` ${pythonString(name)},\n` +

` {\n${fields}\n },\n` +

` total=False,\n` +

`)\n` +

`setattr(${name}, "__pydantic_config__", _PARAMS_CONFIG)\n`

);

});

return (

HEADER +

`# ruff: noqa: F401\n\n` +

`from typing import Annotated, Literal\n\n` +

`from pydantic import ConfigDict, Field\n` +

`from typing_extensions import TypedDict\n\n` +

`_PARAMS_CONFIG = ConfigDict(strict=True, extra="forbid")\n\n` +

entries.join("\n") +

`\n__all__ = [\n${models.map((model) => ` ${pythonString(typeName(model))},`).join("\n")}\n]\n`

);

}

function emitLiteralType(name: string, values: string[]): string {

return `${name} = Literal[\n${values.map((value) => ` ${pythonString(value)},`).join("\n")}\n]\n`;

}

function emitModelIds(models: Model[]): string {

const ids = models.map(modelId);

const providers = [...new Set(models.map((model) => model.provider))].sort();

return (

HEADER +

`from typing import Literal\n\n` +

emitLiteralType("ModelId", ids) +

`\nMODEL_IDS: tuple[ModelId, ...] = (\n${ids.map((id) => ` ${pythonString(id)},`).join("\n")}\n)\n\n` +

emitLiteralType("Provider", providers) +

`\nPROVIDERS: tuple[Provider, ...] = (\n${providers.map((provider) => ` ${pythonString(provider)},`).join("\n")}\n)\n`

);

}

function emitRegistry(models: Model[]): string {

const providers = [...new Set(models.map((model) => model.provider))].sort();

const imports = providers

.map((provider) => `from modelparams.types import ${moduleName(provider)}`)

.join("\n");

const entries = models

.map(

(model) =>

` ${pythonString(modelId(model))}: ${moduleName(model.provider)}.${typeName(model)},`,

)

.join("\n");

return (

HEADER +

`from typing import Any\n\n` +

`${imports}\n\n` +

`from .model_ids import ModelId\n\n` +

`PARAM_TYPES: dict[ModelId, Any] = {\n${entries}\n}\n`

);

}

function assertUniqueNames(models: Model[]): void {

const modules = new Map<string, string>();

const types = new Map<string, string>();

for (const model of models) {

const module = moduleName(model.provider);

const previousProvider = modules.get(module);

if (previousProvider && previousProvider !== model.provider) {

throw new Error(

`provider module collision: ${JSON.stringify(previousProvider)} and ${JSON.stringify(model.provider)} both map to ${JSON.stringify(module)}`,

);

}

modules.set(module, model.provider);

const key = `${module}.${typeName(model)}`;

const id = modelId(model);

const previousId = types.get(key);

if (previousId) {

throw new Error(

`generated type collision: ${JSON.stringify(previousId)} and ${JSON.stringify(id)} both map to ${key}`,

);

}

types.set(key, id);

}

}

async function removeStaleTypeModules(expected: Set<string>): Promise<void> {

let entries: string[] = [];

try {

entries = await fs.readdir(TYPES_DIR);

} catch (error) {

if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;

}

await Promise.all(

entries

.filter((entry) => entry.endsWith(".py") && !expected.has(entry))

.map((entry) => fs.unlink(path.join(TYPES_DIR, entry))),

);

}

async function main(): Promise<void> {

const { models, issues } = await loadAllModels();

if (issues.length > 0) {

console.error(`Catalog has ${issues.length} validation issue(s); refusing to codegen:`);

for (const issue of issues) console.error(` ${issue.file}: ${issue.message}`);

process.exit(1);

}

assertUniqueNames(models);

await fs.mkdir(GENERATED_DIR, { recursive: true });

await fs.mkdir(TYPES_DIR, { recursive: true });

const byProvider = new Map<string, Model[]>();

for (const model of models) {

const providerModels = byProvider.get(model.provider) ?? [];

providerModels.push(model);

byProvider.set(model.provider, providerModels);

}

const typeFiles = new Set<string>(["__init__.py"]);

for (const [provider, providerModels] of byProvider) {

const filename = `${moduleName(provider)}.py`;

typeFiles.add(filename);

await fs.writeFile(

path.join(TYPES_DIR, filename),

emitProviderTypes(provider, providerModels),

"utf8",

);

}

await removeStaleTypeModules(typeFiles);

const modules = [...byProvider.keys()].sort().map(moduleName);

await Promise.all([

fs.writeFile(

path.join(GENERATED_DIR, "catalog.json"),

`${JSON.stringify(models, null, 2)}\n`,

"utf8",

),

fs.writeFile(path.join(GENERATED_DIR, "model_ids.py"), emitModelIds(models), "utf8"),

fs.writeFile(path.join(GENERATED_DIR, "registry.py"), emitRegistry(models), "utf8"),

fs.writeFile(path.join(GENERATED_DIR, "__init__.py"), `${HEADER.trimEnd()}\n`, "utf8"),

fs.writeFile(

path.join(TYPES_DIR, "__init__.py"),

`${HEADER}__all__ = [${modules.map(pythonString).join(", ")}]\n`,

"utf8",

),

]);

console.log(

`python codegen: wrote catalog and request types for ${models.length} models across ${byProvider.size} providers`,

);

}

const isMain = process.argv[1] && path.resolve(process.argv[1]) === fileURLToPath(import.meta.url);

if (isMain) {

main().catch((error) => {

console.error("python codegen crashed:", error);

process.exit(1);

});

}

export { pythonLiteral, typeName };