import { eq, and, count, isNull, isNotNull, inArray } from "drizzle-orm";
import type { Database } from "../client.js";
import { agents } from "../schema/index.js";
import { insertReturning, updateReturning } from "../helpers.js";
import { DEFAULT_TENANT_ID } from "@opentrust/shared";

export function agentQueries(db: Database) {
  return {
    async findById(id: string, tenantId: string = DEFAULT_TENANT_ID) {
      const result = await db.select().from(agents).where(and(eq(agents.id, id), eq(agents.tenantId, tenantId))).limit(1);
      return result[0] ?? null;
    },

    async findByName(name: string, tenantId: string = DEFAULT_TENANT_ID) {
      const result = await db.select().from(agents).where(and(eq(agents.name, name), eq(agents.tenantId, tenantId))).limit(1);
      return result[0] ?? null;
    },

    async findAll(tenantId: string = DEFAULT_TENANT_ID, includeArchived = false) {
      const conditions = [eq(agents.tenantId, tenantId)];
      if (!includeArchived) conditions.push(isNull(agents.archivedAt));
      return db.select().from(agents).where(and(...conditions)).orderBy(agents.createdAt);
    },

    async findByHostId(hostId: string, tenantId: string = DEFAULT_TENANT_ID) {
      return db
        .select()
        .from(agents)
        .where(and(eq(agents.hostId, hostId), eq(agents.tenantId, tenantId)))
        .orderBy(agents.createdAt);
    },

    async countAll(tenantId: string = DEFAULT_TENANT_ID) {
      const result = await db.select({ count: count() }).from(agents).where(eq(agents.tenantId, tenantId));
      return result[0]?.count ?? 0;
    },

    async create(data: {
      name: string;
      description?: string | null;
      provider?: string;
      hostId?: string | null;
      metadata?: Record<string, unknown>;
      tenantId?: string;
    }) {
      return insertReturning(db, agents, {
        ...data,
        provider: data.provider ?? "custom",
        hostId: data.hostId ?? null,
        metadata: data.metadata ?? {},
        tenantId: data.tenantId ?? DEFAULT_TENANT_ID,
      });
    },

    async update(id: string, data: Partial<{
      name: string;
      description: string | null;
      provider: string;
      status: string;
      hostId: string | null;
      lastSeenAt: Date | string;
      metadata: Record<string, unknown>;
    }>, tenantId: string = DEFAULT_TENANT_ID) {
      return updateReturning(db, agents, and(eq(agents.id, id), eq(agents.tenantId, tenantId)), {
        ...data,
        updatedAt: new Date().toISOString(),
      });
    },

    async delete(id: string, tenantId: string = DEFAULT_TENANT_ID) {
      await db.delete(agents).where(and(eq(agents.id, id), eq(agents.tenantId, tenantId)));
    },

    async heartbeat(id: string, tenantId: string = DEFAULT_TENANT_ID) {
      await db
        .update(agents)
        .set({ status: "active", lastSeenAt: new Date().toISOString(), updatedAt: new Date().toISOString() })
        .where(and(eq(agents.id, id), eq(agents.tenantId, tenantId)));
    },

    async setHostId(id: string, hostId: string | null, tenantId: string = DEFAULT_TENANT_ID) {
      await db
        .update(agents)
        .set({ hostId, updatedAt: new Date().toISOString() })
        .where(and(eq(agents.id, id), eq(agents.tenantId, tenantId)));
    },

    async clearHostAgents(hostId: string, tenantId: string = DEFAULT_TENANT_ID) {
      await db
        .update(agents)
        .set({ hostId: null, updatedAt: new Date().toISOString() })
        .where(and(eq(agents.hostId, hostId), eq(agents.tenantId, tenantId)));
    },

    async archive(ids: string[]) {
      if (ids.length === 0) return;
      await db
        .update(agents)
        .set({ archivedAt: new Date().toISOString(), updatedAt: new Date().toISOString() })
        .where(inArray(agents.id, ids));
    },

    async unarchive(ids: string[]) {
      if (ids.length === 0) return;
      await db
        .update(agents)
        .set({ archivedAt: null, updatedAt: new Date().toISOString() })
        .where(inArray(agents.id, ids));
    },
  };
}