import { first } from "../db";
import { accountDb } from "../scoped-db";
import { locate, locatedByPrefix, registerLocator } from "../scoped-db/locators";
import { scopedId } from "../scoped-db/ids";
import { accountsById } from "../people";
import { billingSettings, commitUsage, hardCap, periodUsage, releaseUsage, reserveUsage } from "../metering/quota";
import type { Reservation } from "../metering/quota";
import { recordUsage } from "../metering/record";
import { planAllowsEditorContainers } from "~/lib/plans";
import { containerLimits, containerStub } from "./container-config";
import { accountDisableState, assertAccountActive } from "../account-disable";

export const EDITOR_INSTANCE_TYPE = "standard-1";
export const EDITOR_METER = "compute.editor";
const HOUR_MS = 60 * 60 * 1000;
const HOUR_S = 60 * 60;
const SECOND_MS = 1000;
const MINUTE_S = 60;

export type ContainerUsage = {
  /** Seconds committed this month plus the elapsed time of sessions still running. */
  usedSeconds: number;
  includedSeconds: number;
  hardCapSeconds: number | null;
  /** False on plans with no editor allowance and a zero hard cap (free tier without a spending limit). */
  allowed: boolean;
};

export const CONTAINERS_PAID_ONLY = "Containers are available on paid plans";

/** The account's editor-container usage this month against its plan, for status bars and listings. */
export async function containerUsage(accountId: string): Promise<ContainerUsage> {
  const now = Date.now();
  const [settings, period, running] = await Promise.all([
    billingSettings(accountId),
    periodUsage(accountId, EDITOR_METER),
    accountDb(accountId).all<{ started_at: number }>(
      "SELECT started_at FROM editor_container_sessions WHERE account_id = ? AND stopped_at IS NULL",
      accountId,
    ),
  ]);
  const elapsed = running.reduce((sum, row) => sum + (now - row.started_at) / SECOND_MS, 0);
  const includedSeconds = settings.plan.included[EDITOR_METER] ?? 0;
  const hardCapSeconds = hardCap(settings, EDITOR_METER);
  return {
    usedSeconds: period.quantity + elapsed,
    includedSeconds,
    hardCapSeconds,
    allowed: planAllowsEditorContainers(settings.plan),
  };
}

/** Whether `accountId`'s plan lets it start editor containers at all (independent of EDITOR_CONTAINERS). */
export async function containersAllowedFor(accountId: string): Promise<boolean> {
  if (await accountDisableState(accountId)) return false;
  return planAllowsEditorContainers((await billingSettings(accountId)).plan);
}

export class ContainerCapError extends Error {}

function editorRepoKey(repoId: string, accountId: string): string {
  return `${repoId}:${accountId}`;
}

/** Backfill visit: indexes the repositories an account opened editor containers on, for per-repo listings. */
export async function indexEditorRepos(accountId: string): Promise<void> {
  const rows = await accountDb(accountId).all<{ repo_id: string }>(
    "SELECT DISTINCT repo_id FROM editor_container_sessions",
  );
  for (const row of rows)
    await registerLocator("editor_repo_account", editorRepoKey(row.repo_id, accountId), accountId);
}

/**
 * Holds `budgetSeconds` of compute.editor for a new session and records its row. The reservation id is
 * the session id, so a retried start of the same session never holds twice. Throws QuotaExceededError
 * when the hold doesn't fit. The per-account concurrency cap is enforced inside the insert itself, so two
 * simultaneous starts cannot both slip under it.
 */
export async function openContainerSession(
  accountId: string,
  repoId: string,
  branch: string,
  budgetSeconds: number,
): Promise<Reservation> {
  await assertAccountActive(accountId);
  const cap = containerLimits().perAccount;
  const id = scopedId(accountId);
  const reservation = await reserveUsage({
    id,
    accountId,
    meter: EDITOR_METER,
    quantity: budgetSeconds,
    ttlMs: budgetSeconds * SECOND_MS + HOUR_MS,
  });
  try {
    await registerLocator("editor_session", id, accountId);
    await registerLocator("editor_repo_account", editorRepoKey(repoId, accountId), accountId);
    const result = await accountDb(accountId).run(
      `INSERT INTO editor_container_sessions (id, account_id, repo_id, branch, instance_type, started_at)
       SELECT ?1, ?2, ?3, ?4, ?5, ?6
       WHERE (SELECT COUNT(*) FROM editor_container_sessions WHERE account_id = ?2 AND stopped_at IS NULL) < ?7`,
      id,
      accountId,
      repoId,
      branch,
      EDITOR_INSTANCE_TYPE,
      Date.now(),
      cap,
    );
    if ((result.meta.changes ?? 0) === 0) {
      throw new ContainerCapError(`This account already has ${cap} containers running; stop one first`);
    }
  } catch (error) {
    await releaseUsage(id, accountId).catch(() => undefined);
    throw error;
  }
  return reservation;
}

type ClosedRow = { account_id: string; repo_id: string; started_at: number; stopped_at: number; stop_reason: string };

/**
 * Closes the row (a no-op when already closed) and returns it as stored, so a retried settle after a
 * failed commit bills the same seconds again and the keyed commit dedupes it.
 */
async function closeRow(sessionId: string, reason: string, stoppedAt: number): Promise<ClosedRow | null> {
  const accountId = await locate("editor_session", sessionId);
  if (!accountId) return null;
  const db = accountDb(accountId);
  await db.run(
    `UPDATE editor_container_sessions
     SET stopped_at = MAX(started_at, ?1), minutes = CAST(ROUND((MAX(started_at, ?1) - started_at) / ${MINUTE_S * SECOND_MS}.0) AS INTEGER), stop_reason = ?2
     WHERE id = ?3 AND stopped_at IS NULL`,
    stoppedAt,
    reason,
    sessionId,
  );
  const row = await db.first<ClosedRow>(
    "SELECT account_id, repo_id, started_at, stopped_at, stop_reason FROM editor_container_sessions WHERE id = ?",
    sessionId,
  );
  return row && row.stopped_at !== null && row.stop_reason !== "failed" ? row : null;
}

function billedSeconds(row: ClosedRow, heldSeconds: number): number {
  return Math.max(0, Math.min((row.stopped_at - row.started_at) / SECOND_MS, heldSeconds));
}

/**
 * Closes the session row and settles its hold with the seconds it ran, never more than was held.
 * Idempotent: safe to retry after any partial failure. Returns the seconds billed.
 */
export async function closeContainerSession(
  reservation: Reservation,
  reason: string,
  stoppedAt = Date.now(),
): Promise<number> {
  const row = await closeRow(reservation.id, reason, stoppedAt);
  if (!row) return 0;
  const quantity = billedSeconds(row, reservation.quantity);
  await commitUsage(reservation, {
    quantity,
    repoId: row.repo_id,
    attributes: { instanceType: EDITOR_INSTANCE_TYPE, reason: row.stop_reason },
  });
  return quantity;
}

/** A boot that never produced a usable container: nothing is billed and the hold is released. */
export async function abandonContainerSession(reservation: Reservation): Promise<void> {
  const now = Date.now();
  await accountDb(reservation.accountId).run(
    `UPDATE editor_container_sessions SET stopped_at = ?1, minutes = 0, stop_reason = 'failed'
     WHERE id = ?2 AND stopped_at IS NULL`,
    now,
    reservation.id,
  );
  await releaseUsage(reservation.id, reservation.accountId);
}

/**
 * Settles a session whose Durable Object no longer remembers it (its storage was reset, or the row was
 * left open by a crash). Bills against the original hold when it still exists, else up to the session cap.
 */
export async function closeOrphanSession(sessionId: string, reason: string, now = Date.now()): Promise<number> {
  const row = await closeRow(sessionId, reason, now);
  if (!row) return 0;
  const attributes = { instanceType: EDITOR_INSTANCE_TYPE, reason: row.stop_reason };
  const held = await accountDb(row.account_id).first<{ period: string; quantity: number }>(
    "SELECT period, quantity FROM usage_reservations WHERE id = ?",
    sessionId,
  );
  if (held) {
    const reservation: Reservation = { id: sessionId, accountId: row.account_id, meter: EDITOR_METER, ...held };
    const quantity = billedSeconds(row, held.quantity);
    await commitUsage(reservation, { quantity, repoId: row.repo_id, attributes });
    return quantity;
  }
  const quantity = billedSeconds(row, containerLimits().capHours * HOUR_S);
  await recordUsage({
    accountId: row.account_id,
    meter: EDITOR_METER,
    quantity,
    repoId: row.repo_id,
    attributes,
    idempotencyKey: `reservation:${sessionId}`,
    at: now,
  });
  return quantity;
}

const EXPIRED_BATCH = 100;
/** A row is inserted only after its hold, so a hold gone this long means the session was never settled. */
const MISSING_HOLD_GRACE_MS = 10 * 60 * 1000;

/** When the account's next open session runs out of hold (or lost it): its AccountState alarm closes it then. */
export const EXPIRING_SESSION_DUE = `SELECT MIN(COALESCE(u.expires_at, s.started_at + ${MISSING_HOLD_GRACE_MS})) AS due
  FROM editor_container_sessions s LEFT JOIN usage_reservations u ON u.id = s.id WHERE s.stopped_at IS NULL`;

/**
 * AccountState alarm: closes the account's rows whose hold expired or vanished without the Durable Object
 * settling them, so they stop counting toward the concurrency cap and running usage. Each row is billed at most
 * what it held.
 */
export async function closeExpiredContainerSessions(accountId: string, now = Date.now()): Promise<number> {
  const rows = await accountDb(accountId).all<{
    id: string;
    account_id: string;
    repo_id: string;
    branch: string;
    expires_at: number | null;
  }>(
    `SELECT s.id, s.account_id, s.repo_id, s.branch, u.expires_at FROM editor_container_sessions s LEFT JOIN usage_reservations u ON u.id = s.id
     WHERE s.stopped_at IS NULL AND ((u.id IS NULL AND s.started_at < ?2) OR u.expires_at < ?1) LIMIT ?3`,
    now,
    now - MISSING_HOLD_GRACE_MS,
    EXPIRED_BATCH,
  );
  let closed = 0;
  for (const row of rows) {
    try {
      // The object still running this session stops its container and settles from its own hold first.
      await containerStub(row.account_id, row.repo_id, row.branch).stop("expired", row.id);
      await closeOrphanSession(row.id, "expired", row.expires_at ?? now);
      closed++;
    } catch (error) {
      console.error("Couldn't close an expired container session", error);
    }
  }
  // Still due: failing makes the alarm back off instead of retrying at its floor interval.
  if (closed < rows.length) throw new Error(`${rows.length - closed} expired container sessions stayed open`);
  return closed;
}

export type ContainerSessionRow = {
  id: string;
  accountId: string;
  accountName: string | null;
  repoId: string;
  repoOwner: string | null;
  repoName: string | null;
  branch: string;
  startedAt: number;
  stoppedAt: number | null;
  stopReason: string | null;
  /** Last terminal or preview I/O, from the Durable Object; null once stopped or when it isn't reachable. */
  lastActivity: number | null;
  /** Derived from the same start and stop times the meter is billed on. */
  minutes: number;
};

const RECENT_LIMIT = 30;

type SessionDbRow = {
  id: string;
  account_id: string;
  account_name: string | null;
  repo_id: string;
  owner_name: string | null;
  repo_name: string | null;
  branch: string;
  started_at: number;
  stopped_at: number | null;
  stop_reason: string | null;
};

function toSession(row: SessionDbRow, now: number): ContainerSessionRow {
  return {
    id: row.id,
    accountId: row.account_id,
    accountName: row.account_name,
    repoId: row.repo_id,
    repoOwner: row.owner_name,
    repoName: row.repo_name,
    branch: row.branch,
    startedAt: row.started_at,
    stoppedAt: row.stopped_at,
    stopReason: row.stop_reason,
    lastActivity: null,
    minutes: ((row.stopped_at ?? now) - row.started_at) / (MINUTE_S * SECOND_MS),
  };
}

const SESSION_SELECT = `SELECT s.id, s.account_id, s.repo_id, s.branch, s.started_at, s.stopped_at, s.stop_reason
  FROM editor_container_sessions s`;

type ScopedSessionRow = Omit<SessionDbRow, "account_name" | "owner_name" | "repo_name">;

/** Account and repository names come from the D1 directory. */
async function withNames(rows: ScopedSessionRow[]): Promise<SessionDbRow[]> {
  const repoIds = [...new Set(rows.map((row) => row.repo_id))];
  const repos = new Map<string, { name: string; owner_id: string }>();
  for (const repoId of repoIds) {
    const repo = await first<{ name: string; owner_id: string }>(
      "SELECT name, owner_id FROM repos WHERE id = ?",
      repoId,
    );
    if (repo) repos.set(repoId, repo);
  }
  const accounts = await accountsById([
    ...rows.map((row) => row.account_id),
    ...[...repos.values()].map((repo) => repo.owner_id),
  ]);
  return rows.map((row) => {
    const repo = repos.get(row.repo_id);
    return {
      ...row,
      account_name: accounts.get(row.account_id)?.name ?? null,
      repo_name: repo?.name ?? null,
      owner_name: repo ? (accounts.get(repo.owner_id)?.name ?? null) : null,
    };
  });
}

const ORDER = "ORDER BY (s.stopped_at IS NULL) DESC, s.started_at DESC";

/** Running sessions first, then the most recent stopped ones, for one billed account or one repo. */
export async function listContainerSessions(
  scope: { accountId: string } | { repoId: string },
): Promise<ContainerSessionRow[]> {
  let rows: ScopedSessionRow[];
  if ("accountId" in scope) {
    rows = await accountDb(scope.accountId).all<ScopedSessionRow>(
      `${SESSION_SELECT} WHERE s.account_id = ? ${ORDER} LIMIT ${RECENT_LIMIT}`,
      scope.accountId,
    );
  } else {
    // Sessions are billed to whoever opened them: fan out over the accounts that opened one on this repository.
    rows = [];
    for (const accountId of await locatedByPrefix("editor_repo_account", `${scope.repoId}:`)) {
      rows.push(
        ...(await accountDb(accountId).all<ScopedSessionRow>(
          `${SESSION_SELECT} WHERE s.repo_id = ? ${ORDER} LIMIT ${RECENT_LIMIT}`,
          scope.repoId,
        )),
      );
    }
    rows.sort(
      (left, right) =>
        Number(left.stopped_at !== null) - Number(right.stopped_at !== null) || right.started_at - left.started_at,
    );
    rows = rows.slice(0, RECENT_LIMIT);
  }
  const now = Date.now();
  return (await withNames(rows)).map((row) => toSession(row, now));
}

export async function loadContainerSession(sessionId: string): Promise<ContainerSessionRow | null> {
  const accountId = await locate("editor_session", sessionId);
  if (!accountId) return null;
  const row = await accountDb(accountId).first<ScopedSessionRow>(`${SESSION_SELECT} WHERE s.id = ?`, sessionId);
  if (!row) return null;
  const [named] = await withNames([row]);
  return named ? toSession(named, Date.now()) : null;
}