import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../../infrastructure/database/prisma/prisma.service';
import { type LocalModel, type Prisma, type RuntimeType } from '../../../generated/prisma';
import {
  type CreateLocalModelData,
  type InstalledModelDedupRef,
  type LocalModelFilters,
} from '../types/ollama.types';
import { DEPRECATED_DEFAULT_LOCAL_MODEL_KEYS } from '../constants/default-models.constants';

@Injectable()
export class LocalModelsRepository {
  constructor(private readonly prisma: PrismaService) {}

  async create(data: CreateLocalModelData): Promise<LocalModel> {
    return this.prisma.localModel.create({ data });
  }

  async findById(id: string): Promise<LocalModel | null> {
    return this.prisma.localModel.findUnique({ where: { id } });
  }

  async findAll(filters: LocalModelFilters, page: number, limit: number): Promise<LocalModel[]> {
    const where = this.buildWhereClause(filters);
    const skip = (page - 1) * limit;

    return this.prisma.localModel.findMany({
      where,
      skip,
      take: limit,
      orderBy: { createdAt: 'desc' },
      include: { roles: true },
    });
  }

  async countAll(filters: LocalModelFilters): Promise<number> {
    const where = this.buildWhereClause(filters);
    return this.prisma.localModel.count({ where });
  }

  async upsertByNameTagRuntime(data: CreateLocalModelData): Promise<LocalModel> {
    return this.prisma.localModel.upsert({
      where: {
        name_tag_runtime: {
          name: data.name,
          tag: data.tag,
          runtime: data.runtime,
        },
      },
      update: {
        sizeBytes: data.sizeBytes,
        family: data.family,
        parameters: data.parameters,
        quantization: data.quantization,
        category: data.category,
        isInstalled: data.isInstalled ?? true,
      },
      create: data,
    });
  }

  async findByRuntime(runtime: RuntimeType): Promise<LocalModel[]> {
    return this.prisma.localModel.findMany({
      where: { runtime, isInstalled: true },
      include: { roles: true },
    });
  }

  async findAllInstalled(): Promise<LocalModel[]> {
    return this.prisma.localModel.findMany({
      where: { isInstalled: true },
      orderBy: { name: 'asc' },
    });
  }

  async findInstalledForDedup(): Promise<InstalledModelDedupRef[]> {
    return this.prisma.localModel.findMany({
      where: { isInstalled: true },
      select: { name: true, tag: true },
    });
  }

  async findAllInstalledWithRoles(): Promise<LocalModel[]> {
    return this.prisma.localModel.findMany({
      where: { isInstalled: true },
      include: { roles: true },
      orderBy: { name: 'asc' },
    });
  }

  async markMissingAsUninstalled(presentKeys: string[], runtime: RuntimeType): Promise<number> {
    const allInstalled = await this.prisma.localModel.findMany({
      where: { runtime, isInstalled: true },
    });

    let count = 0;
    for (const model of allInstalled) {
      const key = `${model.name}:${model.tag}`;
      if (!presentKeys.includes(key)) {
        await this.prisma.localModel.update({
          where: { id: model.id },
          data: { isInstalled: false },
        });
        count++;
      }
    }
    return count;
  }

  async delete(id: string): Promise<LocalModel> {
    return this.prisma.localModel.delete({ where: { id } });
  }

  async deleteDeprecatedDefaults(runtime: RuntimeType): Promise<number> {
    const keys = [...DEPRECATED_DEFAULT_LOCAL_MODEL_KEYS];
    if (keys.length === 0) {
      return 0;
    }

    const result = await this.prisma.localModel.deleteMany({
      where: {
        runtime,
        OR: keys.map((key) => {
          const [name, tag] = key.split(':');
          if (tag) {
            return { name, tag };
          }
          return { name: key };
        }),
      },
    });

    return result.count;
  }

  private buildWhereClause(filters: LocalModelFilters): Prisma.LocalModelWhereInput {
    const where: Prisma.LocalModelWhereInput = {};

    if (filters.runtime !== undefined) {
      where.runtime = filters.runtime;
    }

    if (filters.isInstalled !== undefined) {
      where.isInstalled = filters.isInstalled;
    }

    return where;
  }
}
