import {
    pgTable,
    text,
    timestamp,
    boolean,
    integer,
    jsonb,
    pgEnum,
    vector,
    index
} from 'drizzle-orm/pg-core';
import { nanoid } from 'nanoid';

export const roleEnum = pgEnum('role', ['user', 'assistant']);
export const partTypeEnum = pgEnum('part_type', ['reasoning', 'text', 'tool-call', 'file']);
export const statusEnum = pgEnum('status', ['pending', 'completed', 'failed', 'cancelled']);

export enum ToolCallType {
    Text = 'text',
    Json = 'json',
}

export const settings = pgTable('settings', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    systemAssistants: jsonb('system_assistants').$type<{
        rename: {
            enabled: boolean;
            modelId: string | null;
        };
        rerank: {
            enabled: boolean;
            modelId: string | null;
        };
    }>().notNull().default({} as any),
    appearance: jsonb('appearance').$type<{
        colorScheme?: 'light' | 'dark' | 'system';
        accent?: string;
        neutral?: string;
        hinting?: number;
        fontSize?: string;
    }>(),
}, (table) => [
    index('settings_userId_idx').on(table.userId),
]);

export const providers = pgTable('providers', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    type: text('type').notNull(),
    name: text('name').notNull(),
    enabled: boolean('enabled').notNull().default(true),
    config: jsonb('config').notNull().$type<{
        apiKey?: string;
        apiProxyUrl?: string;
    }>().default({}),
    createdAt: timestamp('created_at').notNull().defaultNow(),
    updatedAt: timestamp('updated_at').notNull().defaultNow().$onUpdate(() => /* @__PURE__ */ new Date()),
}, (table) => [
    index('providers_userId_idx').on(table.userId),
]);

export const models = pgTable('models', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    externalId: text('external_id').notNull(),
    providerId: text('provider_id').notNull().references(() => providers.id, { onDelete: 'cascade' }),
    name: text('name').notNull(),
    cost: jsonb('cost').notNull().$type<{
        prompt?: string;
        completion?: string;
        request?: string;
        image?: string;
        imageTokens?: string;
        imageOutput?: string;
        audio?: string;
        audioOutput?: string;
        inputAudioCache?: string;
        webSearch?: string;
        internalReasoning?: string;
        inputCacheRead?: string;
        inputCacheWrite?: string;
        discount?: string;
    }>().default({}),
    inputModalities: text('input_modalities').array().notNull().default([]),
    outputModalities: text('output_modalities').array().notNull().default([]),
    capabilities: text('capabilities').array().notNull().default([]),
    contextWindow: integer('context_window'),
    supportedParameters: text('supported_parameters').array(),
    isCustom: boolean('is_custom').notNull().default(false),
    enabled: boolean('enabled').notNull().default(true),
    releasedAt: timestamp('released_at'),
}, (table) => [
    index('models_providerId_idx').on(table.providerId),
    index('models_userId_idx').on(table.userId),
]);

export const agents = pgTable('agents', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    name: text('name').notNull(),
    systemPrompt: text('system_prompt'),
    imageUrl: text('image_url'),
    defaultModelId: text('default_model_id').references(() => models.id, { onDelete: 'set null' }),
    config: jsonb('config').$type<{
        search?: {
            enabled: boolean;
            rerank: boolean;
            maxResults: number;
        }
    }>().default({}),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('agents_userId_idx').on(table.userId),
]);

export const topics = pgTable('topics', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    agentId: text('agent_id').notNull().references(() => agents.id, { onDelete: 'cascade' }),
    name: text('name').notNull(),
    renaming: boolean('renaming').default(false),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('topics_agentId_idx').on(table.agentId),
    index('topics_userId_idx').on(table.userId),
]);

export const messages = pgTable('messages', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    topicId: text('topic_id').notNull().references(() => topics.id, { onDelete: 'cascade' }),
    parentMessageId: text('parent_message_id'),
    generationId: text('generation_id').references(() => generations.id, { onDelete: 'cascade' }),
    role: roleEnum('role').notNull(),
    content: text('content'),
    activeChildId: text('active_child_id'),
    deleted: boolean('deleted').default(false),
    updatedAt: timestamp('updated_at').notNull().defaultNow().$onUpdate(() => /* @__PURE__ */ new Date()),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('messages_topicId_idx').on(table.topicId),
    index('messages_parentMessageId_idx').on(table.parentMessageId),
    index('messages_userId_idx').on(table.userId),
]);

export const messageParts = pgTable('message_parts', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    topicId: text('topic_id').notNull().references(() => topics.id, { onDelete: 'cascade' }),
    messageId: text('message_id').notNull().references(() => messages.id, { onDelete: 'cascade' }),
    toolCallId: text('tool_call_id').references(() => toolCalls.id, { onDelete: 'cascade' }),
    type: partTypeEnum('type').notNull(),
    content: text('content'),
    providerOptions: jsonb('provider_options'),
    finished: boolean('finished').notNull().default(false),
    createdAt: timestamp('created_at').notNull().defaultNow(),
    lastUpdatedAt: timestamp('last_updated_at').notNull().defaultNow().$onUpdate(() => /* @__PURE__ */ new Date()),
}, (table) => [
    index('message_parts_topicId_idx').on(table.topicId),
    index('message_parts_messageId_idx').on(table.messageId),
    index('message_parts_userId_idx').on(table.userId),
]);

export const toolCalls = pgTable('tool_calls', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    toolName: text('tool_name').notNull(),
    status: statusEnum('status').notNull().default('pending'),
    input: jsonb('input').$type<{
        type: ToolCallType;
        value: string;
    }>(),
    output: jsonb('output').$type<{
        type: ToolCallType;
        value: string;
    }>(),
    error: jsonb('error').$type<{
        type: ToolCallType;
        value: string;
    }>(),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('tool_calls_userId_idx').on(table.userId),
]);

export const generations = pgTable('generations', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    topicId: text('topic_id').notNull().references(() => topics.id, { onDelete: 'cascade' }),
    modelId: text('model_id').notNull(),
    status: statusEnum('status').notNull(),
    tokens: jsonb('tokens').$type<{
        input?: number;
        cache?: {
            read?: number;
            write?: number;
        };
        output?: number;
        thinking?: number;
        ttft?: number;
        tps?: number;
    }>(),
    error: text('error'),
}, (table) => [
    index('generations_topicId_idx').on(table.topicId),
    index('generations_userId_idx').on(table.userId),
]);

export const files = pgTable('files', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    name: text('name').notNull(),
    mimeType: text('mime_type').notNull(),
    size: integer('size').notNull().default(0),
    url: text('url').notNull(),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('files_userId_idx').on(table.userId),
]);

export const embeddings = pgTable('embeddings', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    fileId: text('file_id').notNull().references(() => files.id, { onDelete: 'cascade' }),
    content: text('content').notNull(),
    embedding: vector('embedding', { dimensions: 1536 }),
    metadata: jsonb('metadata'),
}, (table) => [
    index('embeddings_fileId_idx').on(table.fileId),
]);

export const attachments = pgTable('attachments', {
    id: text('id').primaryKey().$defaultFn(nanoid),
    userId: text('user_id').notNull(),
    topicId: text('topic_id').notNull().references(() => topics.id, { onDelete: 'cascade' }),
    messageId: text('message_id').notNull().references(() => messages.id, { onDelete: 'cascade' }),
    fileId: text('file_id').notNull().references(() => files.id, { onDelete: 'cascade' }),
    createdAt: timestamp('created_at').notNull().defaultNow(),
}, (table) => [
    index('attachments_topicId_idx').on(table.topicId),
    index('attachments_messageId_idx').on(table.messageId),
    index('attachments_userId_idx').on(table.userId),
]);

export const users = pgTable("users", {
    id: text("id").primaryKey(),
    name: text("name").notNull(),
    email: text("email").notNull().unique(),
    emailVerified: boolean("email_verified").default(false).notNull(),
    image: text("image"),
    createdAt: timestamp("created_at").defaultNow().notNull(),
    updatedAt: timestamp("updated_at")
        .defaultNow()
        .$onUpdate(() => /* @__PURE__ */ new Date())
        .notNull(),
});

export const sessions = pgTable(
    "sessions",
    {
        id: text("id").primaryKey(),
        expiresAt: timestamp("expires_at").notNull(),
        token: text("token").notNull().unique(),
        createdAt: timestamp("created_at").defaultNow().notNull(),
        updatedAt: timestamp("updated_at")
            .$onUpdate(() => /* @__PURE__ */ new Date())
            .notNull(),
        ipAddress: text("ip_address"),
        userAgent: text("user_agent"),
        userId: text("user_id")
            .notNull()
            .references(() => users.id, { onDelete: "cascade" }),
    },
    (table) => [index("sessions_userId_idx").on(table.userId)],
);

export const accounts = pgTable(
    "accounts",
    {
        id: text("id").primaryKey(),
        accountId: text("account_id").notNull(),
        providerId: text("provider_id").notNull(),
        userId: text("user_id")
            .notNull()
            .references(() => users.id, { onDelete: "cascade" }),
        accessToken: text("access_token"),
        refreshToken: text("refresh_token"),
        idToken: text("id_token"),
        accessTokenExpiresAt: timestamp("access_token_expires_at"),
        refreshTokenExpiresAt: timestamp("refresh_token_expires_at"),
        scope: text("scope"),
        password: text("password"),
        createdAt: timestamp("created_at").defaultNow().notNull(),
        updatedAt: timestamp("updated_at")
            .$onUpdate(() => /* @__PURE__ */ new Date())
            .notNull(),
    },
    (table) => [index("accounts_userId_idx").on(table.userId)],
);

export const verifications = pgTable(
    "verifications",
    {
        id: text("id").primaryKey(),
        identifier: text("identifier").notNull(),
        value: text("value").notNull(),
        expiresAt: timestamp("expires_at").notNull(),
        createdAt: timestamp("created_at").defaultNow().notNull(),
        updatedAt: timestamp("updated_at")
            .defaultNow()
            .$onUpdate(() => /* @__PURE__ */ new Date())
            .notNull(),
    },
    (table) => [index("verifications_identifier_idx").on(table.identifier)],
);