import type {
  GeneratorMetadata,
  PostgresColumn,
  PostgresMaterializedView,
  PostgresSchema,
  PostgresTable,
  PostgresType,
  PostgresView,
} from "../types.ts";

/**
 * Emits tables/views/materialized views in `GeneratorMetadata` order. Pass
 * input pre-sorted with `sortGeneratorMetadata` for deterministic output.
 */
export const generatePython = ({
  schemas,
  tables,
  views,
  materializedViews,
  columns,
  types,
}: GeneratorMetadata): string => {
  const ctx = new PythonContext(types, columns, schemas);
  // Used for efficient lookup of types by schema name
  const schemasNames = new Set(schemas.map((schema) => schema.name));
  const py_tables = tables.flatMap((table) => {
    const py_class_and_methods = ctx.tableToClass(table);
    return py_class_and_methods;
  });
  const composite_types = types
    // We always include system schemas, so we need to filter out types that are not in the included schemas
    .filter(
      (type) => type.attributes.length > 0 && schemasNames.has(type.schema),
    )
    .map((type) => ctx.typeToClass(type));
  const py_views = views.map((view) => ctx.viewToClass(view));
  const py_matviews = materializedViews.map((matview) =>
    ctx.matViewToClass(matview),
  );

  let output = `from __future__ import annotations

import datetime
import uuid
from typing import (
    Annotated,
    Any,
    List,
    Literal,
    Optional,
    TypedDict,
)
from typing_extensions import NotRequired, TypeAlias

from pydantic import BaseModel, Field, JsonValue

${concatLines(Object.values(ctx.user_enums))}

${concatLines(py_tables)}

${concatLines(py_views)}

${concatLines(py_matviews)}

${concatLines(composite_types)}

`.trimEnd();

  return `${output}\n`;
};

interface Serializable {
  serialize(): string;
}

class PythonContext {
  types: { [k: string]: PostgresType };
  user_enums: { [k: string]: PythonEnum };
  columns: Record<number, PostgresColumn[]>;
  schemas: { [k: string]: PostgresSchema };

  constructor(
    types: PostgresType[],
    columns: PostgresColumn[],
    schemas: PostgresSchema[],
  ) {
    this.schemas = Object.fromEntries(
      schemas.map((schema) => [schema.name, schema]),
    );
    this.types = Object.fromEntries(
      types.map((type) => [qualifiedTypeName(type.schema, type.name), type]),
    );
    this.columns = columns.reduce(
      (acc, curr) => {
        (acc[curr.table_id] ??= []).push(curr);
        return acc;
      },
      {} as Record<number, PostgresColumn[]>,
    );
    this.user_enums = Object.fromEntries(
      types
        .filter((type) => type.enums.length > 0)
        .map((type) => [
          qualifiedTypeName(type.schema, type.name),
          new PythonEnum(type),
        ]),
    );
  }

  getSchema(name: string): PostgresSchema {
    const schema = this.schemas[name];
    if (!schema) {
      throw new Error(
        `Schema ${JSON.stringify(name)} is missing from the generator metadata`,
      );
    }
    return schema;
  }

  resolveTypeName(name: string, typeSchema: string): string {
    const qualified = qualifiedTypeName(typeSchema, name);
    const userEnum = this.user_enums[qualified];
    if (userEnum) {
      return userEnum.name;
    }
    const mappedType = PY_TYPE_MAP[name];
    if (mappedType !== undefined) {
      return mappedType;
    }
    const type = this.types[qualified];
    if (type) {
      return `${formatForPyClassName(type.schema)}${formatForPyClassName(type.name)}`;
    }
    return "Any";
  }

  // Array types live in the schema of their element type, so `typeSchema`
  // carries over to the element.
  parsePgType(pg_type: string, typeSchema: string): PythonType {
    if (pg_type.startsWith("_")) {
      const inner_str = pg_type.slice(1);
      const inner = this.parsePgType(inner_str, typeSchema);
      return new PythonListType(inner);
    } else {
      const type_name = this.resolveTypeName(pg_type, typeSchema);
      return new PythonSimpleType(type_name);
    }
  }

  typeToClass(type: PostgresType): PythonBaseModel {
    const types = Object.values(this.types);
    const attributes = type.attributes.map((attribute) => {
      const type = types.find((type) => type.id === attribute.type_id);
      return {
        ...attribute,
        type,
      };
    });
    const attributeEntries: PythonBaseModelAttr[] = attributes.map(
      (attribute) => {
        const type = this.parsePgType(
          attribute.type!.name,
          attribute.type!.schema,
        );
        // Composite type attributes cannot carry NOT NULL constraints in
        // Postgres, so every field is inherently nullable.
        return new PythonBaseModelAttr(attribute.name, type, true);
      },
    );

    const schema = this.getSchema(type.schema);
    return new PythonBaseModel(type.name, schema, attributeEntries);
  }

  columnsToClassAttrs(table_id: number): PythonBaseModelAttr[] {
    const attrs = this.columns[table_id] ?? [];
    return attrs.map((col) => {
      const type = this.parsePgType(col.format, col.type_schema);
      return new PythonBaseModelAttr(col.name, type, col.is_nullable);
    });
  }

  columnsToDictAttrs(
    table_id: number,
    not_required: boolean,
  ): PythonTypedDictAttr[] {
    const attrs = this.columns[table_id] ?? [];
    return attrs.map((col) => {
      const type = this.parsePgType(col.format, col.type_schema);
      return new PythonTypedDictAttr(
        col.name,
        type,
        col.is_nullable,
        not_required ||
          col.is_nullable ||
          col.is_identity ||
          col.default_value !== null,
      );
    });
  }

  tableToClass(
    table: PostgresTable,
  ): [PythonBaseModel, PythonTypedDict, PythonTypedDict] {
    const schema = this.getSchema(table.schema);
    const select = new PythonBaseModel(
      table.name,
      schema,
      this.columnsToClassAttrs(table.id),
    );
    const insert = new PythonTypedDict(
      table.name,
      "Insert",
      schema,
      this.columnsToDictAttrs(table.id, false),
    );
    const update = new PythonTypedDict(
      table.name,
      "Update",
      schema,
      this.columnsToDictAttrs(table.id, true),
    );
    return [select, insert, update];
  }

  viewToClass(view: PostgresView): PythonBaseModel {
    const attributes = this.columnsToClassAttrs(view.id);
    return new PythonBaseModel(
      view.name,
      this.getSchema(view.schema),
      attributes,
    );
  }

  matViewToClass(matview: PostgresMaterializedView): PythonBaseModel {
    const attributes = this.columnsToClassAttrs(matview.id);
    return new PythonBaseModel(
      matview.name,
      this.getSchema(matview.schema),
      attributes,
    );
  }
}

const qualifiedTypeName = (schema: string, name: string): string =>
  `${schema}.${name}`;

class PythonEnum implements Serializable {
  name: string;
  variants: string[];
  constructor(type: PostgresType) {
    this.name = `${formatForPyClassName(type.schema)}${formatForPyClassName(type.name)}`;
    this.variants = type.enums;
  }
  serialize(): string {
    const variants = this.variants
      .map((item) => escapePythonString(item))
      .join(", ");
    return `${this.name}: TypeAlias = Literal[${variants}]`;
  }
}

type PythonType = PythonListType | PythonSimpleType;

class PythonSimpleType implements Serializable {
  name: string;
  constructor(name: string) {
    this.name = name;
  }
  serialize(): string {
    return this.name;
  }
}

class PythonListType implements Serializable {
  inner: PythonType;
  constructor(inner: PythonType) {
    this.inner = inner;
  }
  serialize(): string {
    return `List[${this.inner.serialize()}]`;
  }
}

class PythonBaseModelAttr implements Serializable {
  name: string;
  pg_name: string;
  py_type: PythonType;
  nullable: boolean;

  constructor(name: string, py_type: PythonType, nullable: boolean) {
    this.name = formatForPyAttributeName(name);
    this.pg_name = name;
    this.py_type = py_type;
    this.nullable = nullable;
  }

  serialize(): string {
    const py_type = this.nullable
      ? `Optional[${this.py_type.serialize()}]`
      : this.py_type.serialize();
    return `    ${this.name}: ${py_type} = Field(alias=${escapePythonString(this.pg_name)})`;
  }
}

class PythonBaseModel implements Serializable {
  name: string;
  table_name: string;
  schema: PostgresSchema;
  class_attributes: PythonBaseModelAttr[];

  constructor(
    name: string,
    schema: PostgresSchema,
    class_attributes: PythonBaseModelAttr[],
  ) {
    this.schema = schema;
    this.class_attributes = class_attributes;
    this.table_name = name;
    this.name = `${formatForPyClassName(schema.name)}${formatForPyClassName(name)}`;
  }
  serialize(): string {
    const attributes =
      this.class_attributes.length > 0
        ? this.class_attributes.map((attr) => attr.serialize()).join("\n")
        : "    pass";
    return `class ${this.name}(BaseModel):\n${attributes}`;
  }
}

class PythonTypedDictAttr implements Serializable {
  name: string;
  pg_name: string;
  py_type: PythonType;
  nullable: boolean;
  not_required: boolean;

  constructor(
    name: string,
    py_type: PythonType,
    nullable: boolean,
    required: boolean,
  ) {
    this.name = formatForPyAttributeName(name);
    this.pg_name = name;
    this.py_type = py_type;
    this.nullable = nullable;
    this.not_required = required;
  }

  serialize(): string {
    const py_type = this.nullable
      ? `Optional[${this.py_type.serialize()}]`
      : this.py_type.serialize();
    const annotation = `Annotated[${py_type}, Field(alias=${escapePythonString(this.pg_name)})]`;
    const rhs = this.not_required ? `NotRequired[${annotation}]` : annotation;
    return `    ${this.name}: ${rhs}`;
  }
}

class PythonTypedDict implements Serializable {
  name: string;
  table_name: string;
  parent_class: string;
  schema: PostgresSchema;
  dict_attributes: PythonTypedDictAttr[];
  operation: "Insert" | "Update";

  constructor(
    name: string,
    operation: "Insert" | "Update",
    schema: PostgresSchema,
    dict_attributes: PythonTypedDictAttr[],
    parent_class: string = "BaseModel",
  ) {
    this.schema = schema;
    this.dict_attributes = dict_attributes;
    this.table_name = name;
    this.name = `${formatForPyClassName(schema.name)}${formatForPyClassName(name)}`;
    this.parent_class = parent_class;
    this.operation = operation;
  }
  serialize(): string {
    const attributes =
      this.dict_attributes.length > 0
        ? this.dict_attributes.map((attr) => attr.serialize()).join("\n")
        : "    pass";
    return `class ${this.name}${this.operation}(TypedDict):\n${attributes}`;
  }
}

function concatLines(items: Serializable[]): string {
  return items.map((item) => item.serialize()).join("\n\n");
}

/**
 * Emits a Postgres name as a double-quoted Python string literal. JSON string
 * escaping is a subset of Python's (`\"`, `\\`, `\n`, `\uXXXX`), so
 * `JSON.stringify` yields a literal Python parses to the exact original name
 * and keeps quotes, backslashes and newlines from breaking out of it.
 */
function escapePythonString(value: string): string {
  return JSON.stringify(value);
}

const PY_TYPE_MAP: Record<string, string> = {
  // Bool
  bool: "bool",

  // Numbers
  int2: "int",
  int4: "int",
  int8: "int",
  float4: "float",
  float8: "float",
  numeric: "float",

  // Strings
  bytea: "bytes",
  bpchar: "str",
  varchar: "str",
  string: "str",
  date: "datetime.date",
  text: "str",
  citext: "str",
  time: "datetime.time",
  timetz: "datetime.time",
  timestamp: "datetime.datetime",
  timestamptz: "datetime.datetime",
  uuid: "uuid.UUID",
  vector: "list[Any]",
  interval: "str",

  // JSON. PostgREST returns these columns already deserialized, so the
  // generated models must accept parsed values (JsonValue), not JSON strings
  // (which is what pydantic's Json[...] validates and parses).
  json: "JsonValue",
  jsonb: "JsonValue",

  // Range types (can be adjusted to more complex types if needed)
  int4range: "str",
  int4multirange: "str",
  int8range: "str",
  int8multirange: "str",
  numrange: "str",
  nummultirange: "str",
  tsrange: "str",
  tsmultirange: "str",
  tstzrange: "str",
  tstzmultirange: "str",
  daterange: "str",
  datemultirange: "str",

  // Miscellaneous types
  void: "None",
  record: "dict[str, Any]",
} as const;

/**
 * Converts a Postgres name to PascalCase.
 *
 * @example
 * ```ts
 * formatForPyTypeName('pokedex') // Pokedex
 * formatForPyTypeName('pokemon_center') // PokemonCenter
 * formatForPyTypeName('victory-road') // VictoryRoad
 * formatForPyTypeName('pokemon league') // PokemonLeague
 * ```
 */

function formatForPyClassName(name: string): string {
  return name
    .split(/[^a-zA-Z0-9]/)
    .map((word) => {
      if (word) {
        return `${word.charAt(0).toUpperCase()}${word.slice(1)}`;
      } else {
        return "";
      }
    })
    .join("");
}
/**
 * Converts a Postgres name to snake_case.
 *
 * @example
 * ```ts
 * formatForPyTypeName('Pokedex') // pokedex
 * formatForPyTypeName('PokemonCenter') // pokemon_enter
 * formatForPyTypeName('victory-road') // victory_road
 * formatForPyTypeName('pokemon league') // pokemon_league
 * ```
 */
function formatForPyAttributeName(name: string): string {
  return name
    .split(/[^a-zA-Z0-9]+/) // Split on non-alphanumeric characters (like spaces, dashes, etc.)
    .map((word) => word.toLowerCase()) // Convert each word to lowercase
    .join("_"); // Join with underscores
}