import type {
  AgentRuntime,
  Machine,
  MachineHosting,
  MachineRuntime,
  MachineRuntimeStatus,
  MachineWithAgents,
  UsageInfo,
} from "@vtit-agent-coding/shared";
import { MACHINE_STALE_TIMEOUT_MS, normalizeRuntime, RUNTIME_LABELS } from "@vtit-agent-coding/shared";
import { asJson, type D1, newId, parseDbTimestamp, parseJsonFields, queryDb } from "./db";

export interface MachineRecord extends Machine {
  ama_environment_id: string | null;
}

export interface MachineWithAgentsRecord extends MachineWithAgents {
  ama_environment_id: string | null;
}

export interface CreateMachineInfo {
  name: string;
  os: string;
  version: string;
  runtimes: MachineRuntime[];
  device_id: string;
}

export interface HeartbeatInfo {
  version?: string;
  runtimes?: MachineRuntime[];
  usage_info?: UsageInfo | null;
}

type CachedMachineSnapshot = {
  id: string;
  status: string;
  last_heartbeat_at: string | null;
  created_at: string;
  runtimes: Array<string | { name?: string; status?: string }>;
};

const machineSnapshotCache = new Map<string, Map<string, CachedMachineSnapshot>>();

function cacheMachineSnapshot(
  ownerId: string,
  machine: Pick<MachineRecord, "id" | "status" | "last_heartbeat_at" | "created_at" | "runtimes">,
): void {
  let ownerMachines = machineSnapshotCache.get(ownerId);
  if (!ownerMachines) {
    ownerMachines = new Map();
    machineSnapshotCache.set(ownerId, ownerMachines);
  }
  ownerMachines.set(machine.id, {
    id: machine.id,
    status: machine.status,
    last_heartbeat_at: machine.last_heartbeat_at,
    created_at: machine.created_at,
    runtimes: asJson<Array<string | { name?: string; status?: string }>>(machine.runtimes, []),
  });
}

function removeCachedMachineSnapshot(ownerId: string, machineId: string): void {
  const ownerMachines = machineSnapshotCache.get(ownerId);
  if (!ownerMachines) return;
  ownerMachines.delete(machineId);
  if (ownerMachines.size === 0) machineSnapshotCache.delete(ownerId);
}

function availableRuntimesFromSnapshots(
  rows: Array<Pick<CachedMachineSnapshot, "status" | "last_heartbeat_at" | "created_at" | "runtimes">>,
): Set<string> {
  const cutoff = Date.now() - MACHINE_STALE_TIMEOUT_MS;
  const available = new Set<string>();
  for (const row of rows) {
    if (row.status !== "online") continue;
    const heartbeatAt = parseDbTimestamp(row.last_heartbeat_at);
    if (!Number.isFinite(heartbeatAt) || heartbeatAt < cutoff) continue;
    for (const entry of row.runtimes) {
      if (typeof entry === "string") {
        available.add(normalizeMachineRuntimeName(entry));
        continue;
      }
      if (entry?.status === "ready" && entry?.name) {
        available.add(normalizeMachineRuntimeName(entry.name));
      }
    }
  }
  return available;
}

export function getCachedAvailableRuntimesByOwner(ownerId: string): Set<string> | null {
  const ownerMachines = machineSnapshotCache.get(ownerId);
  if (!ownerMachines) return null;
  return availableRuntimesFromSnapshots([...ownerMachines.values()]);
}

export async function upsertMachine(db: D1, ownerId: string, info: CreateMachineInfo): Promise<MachineRecord> {
  const id = newId();
  const now = new Date().toISOString();
  await queryDb(
    db,
    `INSERT INTO machines (id, owner_id, device_id, name, os, version, runtimes, status, created_at)
     VALUES ($1, $2, $3, $4, $5, $6, $7, 'offline', NOW())
     ON CONFLICT(owner_id, device_id) DO UPDATE SET name = EXCLUDED.name, os = EXCLUDED.os, version = EXCLUDED.version, runtimes = EXCLUDED.runtimes`,
    [id, ownerId, info.device_id, info.name, info.os, info.version, JSON.stringify(normalizeMachineRuntimes(info.runtimes, now))],
  );
  const res = await queryDb<MachineRecord>(db, "SELECT * FROM machines WHERE owner_id = $1 AND device_id = $2", [ownerId, info.device_id]);
  const machine = parseMachine(res.rows[0]);
  cacheMachineSnapshot(ownerId, machine);
  return machine;
}

export async function createCloudMachine(
  db: D1,
  ownerId: string,
  info: { name: string; runtimes: AgentRuntime[]; amaEnvironmentId: string },
): Promise<MachineRecord> {
  const id = newId();
  const now = new Date().toISOString();
  const runtimes = info.runtimes.map((runtime) => ({ name: runtime, status: "ready" as const, checked_at: now }));
  await queryDb(
    db,
    `INSERT INTO machines (id, owner_id, device_id, name, os, version, runtimes, status, hosting, ama_environment_id, last_heartbeat_at, created_at)
     VALUES ($1, $2, $3, $4, 'cloud', 'cloud', $5, 'online', 'cloud', $6, NOW(), NOW())`,
    [id, ownerId, `cloud-${id}`, info.name, JSON.stringify(runtimes), info.amaEnvironmentId],
  );
  const res = await queryDb<MachineRecord>(db, "SELECT * FROM machines WHERE id = $1", [id]);
  const machine = parseMachine(res.rows[0]);
  cacheMachineSnapshot(ownerId, machine);
  return machine;
}

export async function updateMachineAmaEnvironment(
  db: D1,
  machineId: string,
  ownerId: string,
  amaEnvironmentId: string,
): Promise<MachineRecord | null> {
  const result = await queryDb(db, "UPDATE machines SET ama_environment_id = $1 WHERE id = $2 AND owner_id = $3", [
    amaEnvironmentId,
    machineId,
    ownerId,
  ]);
  if ((result.rowCount ?? 0) === 0) return null;
  const res = await queryDb<MachineRecord>(db, "SELECT * FROM machines WHERE id = $1", [machineId]);
  return res.rows[0] ? parseMachine(res.rows[0]) : null;
}

export interface MachineEnvironmentCandidate {
  machineId: string;
  environmentId: string;
  hosting: MachineHosting;
}

export async function listMachineEnvironmentCandidatesForRuntime(db: D1, ownerId: string, runtime: string): Promise<MachineEnvironmentCandidate[]> {
  const values = runtimeMatchValues(runtime);
  const val1 = values[0] || runtime;
  const val2 = values[1] || val1;
  const res = await queryDb<{
    id: string;
    environment_id: string;
    hosting: MachineHosting;
    runtimes: unknown;
    last_heartbeat_at: string | null;
    created_at: string;
  }>(
    db,
    `
    SELECT m.id, m.ama_environment_id AS environment_id, m.hosting AS hosting, m.runtimes, m.last_heartbeat_at, m.created_at
    FROM machines m
    WHERE m.owner_id = $1
      AND m.ama_environment_id IS NOT NULL
    ORDER BY COALESCE(m.last_heartbeat_at, m.created_at) DESC
    `,
    [ownerId],
  );
  return res.rows
    .filter((row) => {
      const runtimes = asJson<Array<string | { name?: string; status?: string }>>(row.runtimes, []);
      return runtimes.some((entry: string | { name?: string; status?: string }) => {
        if (typeof entry === "string") return entry === val1 || entry === val2;
        return entry?.status === "ready" && (entry?.name === val1 || entry?.name === val2);
      });
    })
    .sort((a, b) => {
      if (a.hosting !== b.hosting) return a.hosting === "cloud" ? -1 : 1;
      return parseDbTimestamp(b.last_heartbeat_at ?? b.created_at) - parseDbTimestamp(a.last_heartbeat_at ?? a.created_at);
    })
    .map((row) => ({ machineId: row.id, environmentId: row.environment_id, hosting: row.hosting }));
}

export async function deleteMachine(db: D1, machineId: string, ownerId: string): Promise<boolean> {
  const result = await queryDb(db, "DELETE FROM machines WHERE id = $1 AND owner_id = $2", [machineId, ownerId]);
  if ((result.rowCount ?? 0) > 0) removeCachedMachineSnapshot(ownerId, machineId);
  return (result.rowCount ?? 0) > 0;
}

export async function updateMachine(db: D1, machineId: string, ownerId: string, info: HeartbeatInfo): Promise<MachineRecord | null> {
  const now = new Date().toISOString();
  const sets: string[] = ["status = 'online'", "last_heartbeat_at = NOW()"];
  const binds: any[] = [];
  let paramIdx = 1;

  if (info.version) {
    sets.push(`version = $${paramIdx++}`);
    binds.push(info.version);
  }
  if (info.runtimes) {
    sets.push(`runtimes = $${paramIdx++}`);
    binds.push(JSON.stringify(normalizeMachineRuntimes(info.runtimes, now)));
  }
  if ("usage_info" in info) {
    const usageInfo = info.usage_info;
    sets.push(`usage_info = $${paramIdx++}`);
    binds.push(usageInfo == null ? null : JSON.stringify(normalizeUsageInfo(usageInfo)));
  }

  binds.push(machineId, ownerId);
  const result = await queryDb(db, `UPDATE machines SET ${sets.join(", ")} WHERE id = $${paramIdx++} AND owner_id = $${paramIdx}`, binds);
  if ((result.rowCount ?? 0) === 0) return null;

  const res = await queryDb<MachineRecord>(db, "SELECT * FROM machines WHERE id = $1", [machineId]);
  if (!res.rows[0]) return null;
  const machine = parseMachine(res.rows[0]);
  cacheMachineSnapshot(ownerId, machine);
  return machine;
}

export async function listMachines(db: D1, ownerId: string): Promise<MachineWithAgentsRecord[]> {
  const res = await queryDb<MachineWithAgentsRecord>(
    db,
    `
    SELECT m.*,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id) as session_count,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id AND s.status = 'active') as active_session_count
    FROM machines m
    WHERE m.owner_id = $1
    ORDER BY m.last_heartbeat_at DESC NULLS LAST
    `,
    [ownerId],
  );
  return res.rows.map(parseMachine);
}

export async function listMachinesForRuntimeRouting(db: D1, ownerId: string): Promise<MachineRecord[]> {
  const res = await queryDb<MachineRecord>(db, "SELECT * FROM machines WHERE owner_id = $1", [ownerId]);
  return res.rows.map(parseMachine);
}

export async function getMachine(db: D1, machineId: string, ownerId: string): Promise<(MachineWithAgentsRecord & { agents: any[] }) | null> {
  const mRes = await queryDb<MachineWithAgentsRecord>(
    db,
    `
    SELECT m.*,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id) as session_count,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id AND s.status = 'active') as active_session_count
    FROM machines m WHERE m.id = $1 AND m.owner_id = $2
    `,
    [machineId, ownerId],
  );

  const machine = mRes.rows[0];
  if (!machine) return null;

  const aRes = await queryDb(
    db,
    `
    SELECT a.id, a.name, a.role, a.runtime,
      COALESCE(SUM(CASE WHEN s.status = 'active' AND s.machine_id = $1 THEN 1 ELSE 0 END), 0) as active_session_count,
      MAX(CASE WHEN s.machine_id = $1 THEN s.created_at ELSE NULL END) as last_session_at
    FROM agents a
    LEFT JOIN agent_sessions s ON s.agent_id = a.id
    WHERE a.owner_id = $2
    GROUP BY a.id, a.name, a.role, a.runtime
    ORDER BY active_session_count DESC, last_session_at DESC NULLS LAST, a.created_at DESC
    `,
    [machineId, ownerId],
  );

  return { ...parseMachine(machine), agents: aRes.rows };
}

export interface AdminMachine extends MachineWithAgentsRecord {
  owner_name: string | null;
  owner_email: string | null;
}

export async function listAllMachines(db: D1): Promise<AdminMachine[]> {
  const res = await queryDb<AdminMachine>(
    db,
    `
    SELECT m.*,
      u.name AS owner_name, u.email AS owner_email,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id) AS session_count,
      (SELECT COUNT(*) FROM agent_sessions s WHERE s.machine_id = m.id AND s.status = 'active') AS active_session_count
    FROM machines m
    LEFT JOIN "user" u ON u.id = m.owner_id
    ORDER BY m.last_heartbeat_at DESC NULLS LAST
    `,
  );
  return res.rows.map(parseMachine);
}

function parseMachine<T extends MachineRecord>(row: T): T {
  const parsed = parseJsonFields(row, ["runtimes", "usage_info"]);
  if (!parsed.hosting) parsed.hosting = "local";
  parsed.runtimes = normalizeMachineRuntimes(parsed.runtimes ?? [], parsed.last_heartbeat_at ?? parsed.created_at);
  if (parsed.usage_info) parsed.usage_info = normalizeUsageInfo(parsed.usage_info);
  return parsed;
}

function normalizeUsageInfo(info: UsageInfo): UsageInfo {
  return {
    ...info,
    windows: info.windows.map((window) => ({
      ...window,
      utilization: window.utilization < 1 ? window.utilization * 100 : window.utilization,
    })),
  };
}

const RUNTIME_BY_LABEL = Object.fromEntries(Object.entries(RUNTIME_LABELS).map(([runtime, label]) => [label, runtime])) as Record<
  string,
  AgentRuntime
>;

export function runtimeMatchValues(runtime: string): string[] {
  const normalized = normalizeRuntime(runtime);
  const canonical = (RUNTIME_BY_LABEL[normalized] ?? normalized) as AgentRuntime;
  const label = RUNTIME_LABELS[canonical];
  return label && label !== canonical ? [canonical, label] : [canonical];
}

export function runtimeReadyPredicateSql(runtimeExpr: string): string {
  return `
    (
      (
        jsonb_typeof(rt) = 'string'
        AND (rt#>>'{}' = ${runtimeExpr} OR rt#>>'{}' = ${runtimeLabelCaseSql(runtimeExpr)})
      )
      OR (
        (rt->>'status') = 'ready'
        AND (rt->>'name') = ${runtimeExpr}
      )
    )
  `;
}

function runtimeLabelCaseSql(runtimeExpr: string): string {
  const cases = Object.entries(RUNTIME_LABELS)
    .map(([runtime, label]) => `WHEN '${runtime}' THEN '${label.replace(/'/g, "''")}'`)
    .join(" ");
  return `CASE ${runtimeExpr} ${cases} END`;
}

const RUNTIME_STATUSES: readonly MachineRuntimeStatus[] = ["missing", "unauthorized", "unhealthy", "limited", "ready"];

export function normalizeMachineRuntimes(runtimes: MachineRuntime[] | string[], checkedAt: string): MachineRuntime[] {
  return runtimes.map((runtime) => {
    if (typeof runtime === "string") {
      return { name: normalizeMachineRuntimeName(runtime), status: "ready", checked_at: checkedAt };
    }
    const name = normalizeMachineRuntimeName(runtime.name);
    if (!RUNTIME_STATUSES.includes(runtime.status)) {
      throw new Error(`Invalid runtime status "${runtime.status}"`);
    }
    return {
      name,
      status: runtime.status,
      ...(runtime.detail ? { detail: runtime.detail } : {}),
      ...(runtime.reset_at ? { reset_at: runtime.reset_at } : {}),
      checked_at: runtime.checked_at || checkedAt,
    };
  });
}

function normalizeMachineRuntimeName(runtime: string): AgentRuntime {
  const normalized = normalizeRuntime(runtime);
  const canonical = RUNTIME_BY_LABEL[normalized] ?? normalized;
  if (!(canonical in RUNTIME_LABELS)) {
    throw new Error(`Invalid runtime "${runtime}"`);
  }
  return canonical as AgentRuntime;
}

export async function isRuntimeAvailable(db: D1, ownerId: string, runtime: string): Promise<boolean> {
  const availability = await listAvailableRuntimesByOwner(db, ownerId);
  return runtimeMatchValues(runtime).some((value) => availability.has(value));
}

export async function listAvailableRuntimesByOwner(db: D1, ownerId: string): Promise<Set<string>> {
  const queryable = db as D1 & {
    query?: <T = { runtimes: unknown; status: string; last_heartbeat_at: string | null }>(text: string, params?: unknown[]) => Promise<{ rows: T[] }>;
  };
  const res =
    typeof queryable.query === "function"
      ? await queryable.query<{ runtimes: unknown; status: string; last_heartbeat_at: string | null }>(
          "SELECT runtimes, status, last_heartbeat_at FROM machines WHERE owner_id = $1",
          [ownerId],
        )
      : await queryDb<{ runtimes: unknown; status: string; last_heartbeat_at: string | null }>(
          db,
          "SELECT runtimes, status, last_heartbeat_at FROM machines WHERE owner_id = $1",
          [ownerId],
        );

  return availableRuntimesFromSnapshots(
    res.rows.map((row) => ({
      status: row.status,
      last_heartbeat_at: row.last_heartbeat_at,
      created_at: row.last_heartbeat_at ?? new Date(0).toISOString(),
      runtimes: asJson<Array<string | { name?: string; status?: string }>>(row.runtimes, []),
    })),
  );
}

export async function detectStaleMachines(db: D1): Promise<void> {
  const cutoff = new Date(Date.now() - MACHINE_STALE_TIMEOUT_MS).toISOString();
  await queryDb(db, "UPDATE machines SET status = 'offline' WHERE status = 'online' AND last_heartbeat_at < $1", [cutoff]);
  machineSnapshotCache.clear();
}