import { CollectionRegistry, resolveCollectionRelations } from "@rebasepro/common";
import { type CollectionConfig } from "@rebasepro/types";
import { PgEnum, PgTable } from "drizzle-orm/pg-core";
import { Relations } from "drizzle-orm";
import { CollectionRegistryInterface } from "../interfaces";
import { getTableName } from "@rebasepro/common";

/**
 * PostgreSQL-specific collection registry.
 * Extends the base CollectionRegistry with support for Drizzle ORM tables, enums, and relations.
 *
 * Satisfies CollectionRegistryInterface through inheritance from CollectionRegistry.
 */
export class PostgresCollectionRegistry extends CollectionRegistry implements CollectionRegistryInterface {

    private tables = new Map<string, PgTable>();
    private enums = new Map<string, PgEnum<[string, ...string[]]>>();
    private relations = new Map<string, Relations>();

    registerTable(table: PgTable, tableName: string) {
        this.tables.set(tableName, table);
    }

    getTable(tableName: string): PgTable | undefined {
        return this.tables.get(tableName);
    }

    /**
     * Checks if a specific collection has a registered table
     */
    hasTableForCollection(tableName: string): boolean {
        return this.tables.has(tableName);
    }

    /**
     * Returns all registered table names.
     */
    getTableNames(): string[] {
        return Array.from(this.tables.keys());
    }

    /**
     * Finds collections assigned to a specific data source that do not have a registered table.
     */
    getCollectionsWithoutTables(dataSourceKey = "(default)"): CollectionConfig[] {
        const collections = this.getCollections().filter(
            c => c.dataSource === dataSourceKey || (!c.dataSource && dataSourceKey === "(default)")
        );
        return collections.filter(c => !this.tables.has(getTableName(c)));
    }

    registerEnums(enums: Record<string, PgEnum<[string, ...string[]]>>) {
        Object.entries(enums).forEach(([name, value]) => this.enums.set(name, value));
    }

    registerRelations(relations: Record<string, Relations>) {
        Object.entries(relations).forEach(([name, value]) => this.relations.set(name, value));
    }

    getEnum(name: string): PgEnum<[string, ...string[]]> | undefined {
        return this.enums.get(name);
    }

    getRelation(name: string): Relations | undefined {
        return this.relations.get(name);
    }

    getAllEnums(): Record<string, PgEnum<[string, ...string[]]>> {
        return Object.fromEntries(this.enums.entries());
    }

    getAllRelations(): Record<string, Relations> {
        return Object.fromEntries(this.relations.entries());
    }

    /**
     * Get the merged schema object (tables + relations) for use with Drizzle's
     * relational query API (`db.query`).
     */
    getMergedSchema(): Record<string, unknown> {
        const result: Record<string, unknown> = {};
        for (const [name, table] of this.tables.entries()) {
            result[name] = table;
        }
        for (const [name, relation] of this.relations.entries()) {
            result[name] = relation;
        }
        return result;
    }

    /**
     * Get the available Drizzle relation keys for a given collection path.
     * Maps from the collection's relation property names to the Drizzle relation names
     * defined in the schema.
     */
    getRelationKeysForCollection(collectionPath: string): string[] {
        const collection = this.getCollectionByPath(collectionPath);
        if (!collection) return [];
        // Resolved, not authored. `relationName` is optional at the authoring
        // surface and defaults to the property key or the target's slug, so
        // reading the raw field dropped every relation that relied on the
        // default — and never saw relations declared inline on a property at
        // all, since those are not in the `relations` array.
        return Object.keys(resolveCollectionRelations(collection));
    }

}