import { desc, eq } from "drizzle-orm";
import type { Transaction } from "@platform-modules/db";
import type {
  AgentRunRepository,
  ApplicationTransaction,
  AttemptRepository,
  AuditRepository,
  BusinessEventRepository,
  ChangeSetRepository,
  FactoryRunRepository,
  GoalRepository,
  OutboxMessage,
  OutboxRepository,
  PlanRepository,
  PlanRevisionRepository,
  ProjectRepository,
  ProjectVisionRepository,
  ReviewRepository,
  TaskRepository,
  UnitOfWork,
  WorkspaceRepository,
} from "@awp/application";
import type {
  AgentRunId,
  AttemptId,
  AuditRecord,
  BusinessEvent,
  ChangeSetId,
  FactoryRunId,
  GoalId,
  PlanId,
  PlanRevisionId,
  ProjectId,
  ProjectVisionVersionId,
  ReviewId,
  TaskId,
  WorkspaceId,
} from "@awp/contracts";
import { unsafeOpaqueId } from "@awp/contracts";
import type {
  AgentRun,
  AgentRunStatus,
  Attempt,
  AttemptStatus,
  ChangeSet,
  ChangeSetStatus,
  FactoryRun,
  FactoryRunStatus,
  Goal,
  GoalStatus,
  Plan,
  PlanRevision,
  PlanStatus,
  Project,
  ProjectStatus,
  ProjectVisionVersion,
  Review,
  ReviewDisposition,
  ReviewStatus,
  Task,
  TaskStatus,
  Workspace,
} from "@awp/domain";
import {
  agentRuns,
  attempts,
  auditRecords,
  businessEvents,
  changeSets,
  factoryRuns,
  goals,
  outboxMessages,
  planRevisions,
  plans,
  projectVisionVersions,
  projects,
  reviews,
  tasks,
  workspaces,
} from "./schema.js";
import type { schema } from "./schema.js";

type Tx = Transaction<typeof schema>;
interface Database {
  transaction<T>(work: (tx: Tx) => Promise<T>): Promise<T>;
}

class PgProjectRepository implements ProjectRepository {
  constructor(private readonly tx: Tx) {}
  async getById(id: ProjectId): Promise<Project | undefined> {
    const row = (await this.tx.select().from(projects).where(eq(projects.id, id)).limit(1))[0];
    return row
      ? {
          id: unsafeOpaqueId<ProjectId>(row.id),
          name: row.name,
          status: row.status as ProjectStatus,
          revision: row.revision,
        }
      : undefined;
  }
  async insert(value: Project): Promise<void> {
    await this.tx.insert(projects).values(value);
  }
  async update(value: Project): Promise<void> {
    await this.tx
      .update(projects)
      .set({ name: value.name, status: value.status, revision: value.revision })
      .where(eq(projects.id, value.id));
  }
}
class PgProjectVisionRepository implements ProjectVisionRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof projectVisionVersions.$inferSelect): ProjectVisionVersion {
    return {
      id: unsafeOpaqueId<ProjectVisionVersionId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      sequence: row.sequence,
      summary: row.summary,
      ...(row.supersedesId === null
        ? {}
        : { supersedesId: unsafeOpaqueId<ProjectVisionVersionId>(row.supersedesId) }),
    };
  }
  async getById(id: ProjectVisionVersionId) {
    const row = (
      await this.tx
        .select()
        .from(projectVisionVersions)
        .where(eq(projectVisionVersions.id, id))
        .limit(1)
    )[0];
    return row ? this.map(row) : undefined;
  }
  async getCurrent(projectId: ProjectId) {
    const row = (
      await this.tx
        .select()
        .from(projectVisionVersions)
        .where(eq(projectVisionVersions.projectId, projectId))
        .orderBy(desc(projectVisionVersions.sequence))
        .limit(1)
    )[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: ProjectVisionVersion) {
    await this.tx.insert(projectVisionVersions).values(value);
  }
}
class PgGoalRepository implements GoalRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof goals.$inferSelect): Goal {
    return {
      id: unsafeOpaqueId<GoalId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      title: row.title,
      status: row.status as GoalStatus,
      successCriteria: row.successCriteria as string[],
      revision: row.revision,
      ...(row.description === null ? {} : { description: row.description }),
      ...(row.priority === null ? {} : { priority: row.priority }),
      ...(row.targetDate === null ? {} : { targetDate: row.targetDate }),
    };
  }
  async getById(id: GoalId) {
    const row = (await this.tx.select().from(goals).where(eq(goals.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async listByProject(id: ProjectId) {
    return (await this.tx.select().from(goals).where(eq(goals.projectId, id))).map((row) =>
      this.map(row),
    );
  }
  async insert(value: Goal) {
    await this.tx.insert(goals).values({ ...value, successCriteria: [...value.successCriteria] });
  }
  async update(value: Goal) {
    await this.tx
      .update(goals)
      .set({
        title: value.title,
        description: value.description,
        status: value.status,
        priority: value.priority,
        targetDate: value.targetDate,
        successCriteria: [...value.successCriteria],
        revision: value.revision,
      })
      .where(eq(goals.id, value.id));
  }
}
class PgPlanRepository implements PlanRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof plans.$inferSelect): Plan {
    return {
      id: unsafeOpaqueId<PlanId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      title: row.title,
      status: row.status as PlanStatus,
      revision: row.revision,
    };
  }
  async getById(id: PlanId) {
    const row = (await this.tx.select().from(plans).where(eq(plans.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async listByProject(id: ProjectId) {
    return (await this.tx.select().from(plans).where(eq(plans.projectId, id))).map((row) =>
      this.map(row),
    );
  }
  async insert(value: Plan) {
    await this.tx.insert(plans).values(value);
  }
  async update(value: Plan) {
    await this.tx
      .update(plans)
      .set({ title: value.title, status: value.status, revision: value.revision })
      .where(eq(plans.id, value.id));
  }
}
class PgPlanRevisionRepository implements PlanRevisionRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof planRevisions.$inferSelect): PlanRevision {
    return {
      id: unsafeOpaqueId<PlanRevisionId>(row.id),
      planId: unsafeOpaqueId<PlanId>(row.planId),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      sequence: row.sequence,
      title: row.title,
      goalIds: (row.goalIds as string[]).map((x) => unsafeOpaqueId<GoalId>(x)),
      ...(row.projectVisionVersionId === null
        ? {}
        : {
            projectVisionVersionId: unsafeOpaqueId<ProjectVisionVersionId>(
              row.projectVisionVersionId,
            ),
          }),
    };
  }
  async getById(id: PlanRevisionId) {
    const row = (
      await this.tx.select().from(planRevisions).where(eq(planRevisions.id, id)).limit(1)
    )[0];
    return row ? this.map(row) : undefined;
  }
  async getCurrent(planId: PlanId) {
    const row = (
      await this.tx
        .select()
        .from(planRevisions)
        .where(eq(planRevisions.planId, planId))
        .orderBy(desc(planRevisions.sequence))
        .limit(1)
    )[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: PlanRevision) {
    await this.tx.insert(planRevisions).values({ ...value, goalIds: [...value.goalIds] });
  }
}
class PgTaskRepository implements TaskRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof tasks.$inferSelect): Task {
    return {
      id: unsafeOpaqueId<TaskId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      planRevisionId: unsafeOpaqueId<PlanRevisionId>(row.planRevisionId),
      title: row.title,
      status: row.status as TaskStatus,
      dependencyIds: (row.dependencyIds as string[]).map((x) => unsafeOpaqueId<TaskId>(x)),
      goalIds: (row.goalIds as string[]).map((x) => unsafeOpaqueId<GoalId>(x)),
      revision: row.revision,
    };
  }
  async getById(id: TaskId) {
    const row = (await this.tx.select().from(tasks).where(eq(tasks.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async listByProject(id: ProjectId) {
    return (await this.tx.select().from(tasks).where(eq(tasks.projectId, id))).map((row) =>
      this.map(row),
    );
  }
  async insert(value: Task) {
    await this.tx.insert(tasks).values({
      ...value,
      dependencyIds: [...value.dependencyIds],
      goalIds: [...(value.goalIds ?? [])],
    });
  }
  async update(value: Task) {
    await this.tx
      .update(tasks)
      .set({
        title: value.title,
        status: value.status,
        dependencyIds: [...value.dependencyIds],
        goalIds: [...(value.goalIds ?? [])],
        revision: value.revision,
      })
      .where(eq(tasks.id, value.id));
  }
}
class PgFactoryRunRepository implements FactoryRunRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof factoryRuns.$inferSelect): FactoryRun {
    return {
      id: unsafeOpaqueId<FactoryRunId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      planRevisionId: unsafeOpaqueId<PlanRevisionId>(row.planRevisionId),
      status: row.status as FactoryRunStatus,
      revision: row.revision,
      ...(row.reason === null ? {} : { reason: row.reason }),
    };
  }
  async getById(id: FactoryRunId) {
    const row = (
      await this.tx.select().from(factoryRuns).where(eq(factoryRuns.id, id)).limit(1)
    )[0];
    return row ? this.map(row) : undefined;
  }
  async listByProject(id: ProjectId) {
    return (await this.tx.select().from(factoryRuns).where(eq(factoryRuns.projectId, id))).map(
      (row) => this.map(row),
    );
  }
  async insert(value: FactoryRun) {
    await this.tx.insert(factoryRuns).values(value);
  }
  async update(value: FactoryRun) {
    await this.tx
      .update(factoryRuns)
      .set({ status: value.status, reason: value.reason, revision: value.revision })
      .where(eq(factoryRuns.id, value.id));
  }
}
class PgAgentRunRepository implements AgentRunRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof agentRuns.$inferSelect): AgentRun {
    return {
      id: unsafeOpaqueId<AgentRunId>(row.id),
      factoryRunId: unsafeOpaqueId<FactoryRunId>(row.factoryRunId),
      taskId: unsafeOpaqueId<TaskId>(row.taskId),
      status: row.status as AgentRunStatus,
      revision: row.revision,
      ...(row.reason === null ? {} : { reason: row.reason }),
    };
  }
  async getById(id: AgentRunId) {
    const row = (await this.tx.select().from(agentRuns).where(eq(agentRuns.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: AgentRun) {
    await this.tx.insert(agentRuns).values(value);
  }
  async update(value: AgentRun) {
    await this.tx
      .update(agentRuns)
      .set({ status: value.status, reason: value.reason, revision: value.revision })
      .where(eq(agentRuns.id, value.id));
  }
}
class PgWorkspaceRepository implements WorkspaceRepository {
  constructor(private readonly tx: Tx) {}
  async getById(id: WorkspaceId) {
    const row = (await this.tx.select().from(workspaces).where(eq(workspaces.id, id)).limit(1))[0];
    return row
      ? {
          id: unsafeOpaqueId<WorkspaceId>(row.id),
          projectId: unsafeOpaqueId<ProjectId>(row.projectId),
          revision: row.revision,
        }
      : undefined;
  }
  async insert(value: Workspace) {
    await this.tx.insert(workspaces).values(value);
  }
  async update(value: Workspace) {
    await this.tx
      .update(workspaces)
      .set({ revision: value.revision })
      .where(eq(workspaces.id, value.id));
  }
}
class PgAttemptRepository implements AttemptRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof attempts.$inferSelect): Attempt {
    return {
      id: unsafeOpaqueId<AttemptId>(row.id),
      agentRunId: unsafeOpaqueId<AgentRunId>(row.agentRunId),
      workspaceId: unsafeOpaqueId<WorkspaceId>(row.workspaceId),
      status: row.status as AttemptStatus,
      revision: row.revision,
      ...(row.providerId === null ? {} : { providerId: unsafeOpaqueId(row.providerId) }),
      ...(row.accountId === null ? {} : { accountId: row.accountId }),
      ...(row.model === null ? {} : { model: row.model }),
    };
  }
  async getById(id: AttemptId) {
    const row = (await this.tx.select().from(attempts).where(eq(attempts.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: Attempt) {
    await this.tx.insert(attempts).values(value);
  }
  async update(value: Attempt) {
    await this.tx
      .update(attempts)
      .set({
        status: value.status,
        providerId: value.providerId,
        accountId: value.accountId,
        model: value.model,
        revision: value.revision,
      })
      .where(eq(attempts.id, value.id));
  }
}
class PgChangeSetRepository implements ChangeSetRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof changeSets.$inferSelect): ChangeSet {
    return {
      id: unsafeOpaqueId<ChangeSetId>(row.id),
      projectId: unsafeOpaqueId<ProjectId>(row.projectId),
      taskId: unsafeOpaqueId<TaskId>(row.taskId),
      producerAttemptId: unsafeOpaqueId(row.producerAttemptId),
      baseIdentity: row.baseIdentity,
      candidateDigest: row.candidateDigest,
      revision: row.revision,
      status: row.status as ChangeSetStatus,
      ...(row.repositoryId === null ? {} : { repositoryId: unsafeOpaqueId(row.repositoryId) }),
    };
  }
  async getById(id: ChangeSetId) {
    const row = (await this.tx.select().from(changeSets).where(eq(changeSets.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: ChangeSet) {
    await this.tx.insert(changeSets).values(value);
  }
  async update(value: ChangeSet) {
    await this.tx
      .update(changeSets)
      .set({
        candidateDigest: value.candidateDigest,
        revision: value.revision,
        status: value.status,
      })
      .where(eq(changeSets.id, value.id));
  }
}
class PgReviewRepository implements ReviewRepository {
  constructor(private readonly tx: Tx) {}
  private map(row: typeof reviews.$inferSelect): Review {
    return {
      id: unsafeOpaqueId<ReviewId>(row.id),
      changeSetId: unsafeOpaqueId<ChangeSetId>(row.changeSetId),
      candidateDigest: row.candidateDigest,
      reviewerPrincipalId: unsafeOpaqueId(row.reviewerPrincipalId),
      status: row.status as ReviewStatus,
      ...(row.disposition === null ? {} : { disposition: row.disposition as ReviewDisposition }),
    };
  }
  async getById(id: ReviewId) {
    const row = (await this.tx.select().from(reviews).where(eq(reviews.id, id)).limit(1))[0];
    return row ? this.map(row) : undefined;
  }
  async insert(value: Review) {
    await this.tx.insert(reviews).values(value);
  }
  async update(value: Review) {
    await this.tx
      .update(reviews)
      .set({ status: value.status, disposition: value.disposition })
      .where(eq(reviews.id, value.id));
  }
}
class PgBusinessEventRepository implements BusinessEventRepository {
  constructor(private readonly tx: Tx) {}
  async append(value: BusinessEvent) {
    await this.tx.insert(businessEvents).values({ ...value });
  }
}
class PgAuditRepository implements AuditRepository {
  constructor(private readonly tx: Tx) {}
  async append(value: AuditRecord) {
    await this.tx.insert(auditRecords).values({ ...value });
  }
}
class PgOutboxRepository implements OutboxRepository {
  constructor(private readonly tx: Tx) {}
  async append(value: OutboxMessage) {
    await this.tx.insert(outboxMessages).values({ ...value });
  }
}

export class PostgresUnitOfWork implements UnitOfWork {
  constructor(private readonly db: Database) {}
  async transaction<T>(work: (tx: ApplicationTransaction) => Promise<T>): Promise<T> {
    return this.db.transaction(async (tx) =>
      work({
        projects: new PgProjectRepository(tx),
        projectVisions: new PgProjectVisionRepository(tx),
        goals: new PgGoalRepository(tx),
        plans: new PgPlanRepository(tx),
        planRevisions: new PgPlanRevisionRepository(tx),
        tasks: new PgTaskRepository(tx),
        factoryRuns: new PgFactoryRunRepository(tx),
        agentRuns: new PgAgentRunRepository(tx),
        workspaces: new PgWorkspaceRepository(tx),
        attempts: new PgAttemptRepository(tx),
        changeSets: new PgChangeSetRepository(tx),
        reviews: new PgReviewRepository(tx),
        events: new PgBusinessEventRepository(tx),
        audit: new PgAuditRepository(tx),
        outbox: new PgOutboxRepository(tx),
      }),
    );
  }
}
