import { createHash } from "node:crypto"

export type AgentStatus =
  | "planning"
  | "validating"
  | "awaiting_approval"
  | "executing"
  | "succeeded"
  | "failed"
  | "cancelled"
  | "exhausted"

export type Proposal =
  | { kind: "tool"; name: string; args: Record<string, unknown> }
  | { kind: "finish"; answer: string }

export type ToolPolicy = {
  requiresApproval: boolean
}

export type AgentState = {
  status: AgentStatus
  steps: number
  maxSteps: number
  spentUsd: number
  maxCostUsd: number
  seen: Set<string>
}

export type Validation =
  | { ok: true; fingerprint?: string; requiresApproval?: boolean }
  | {
      ok: false
      reason: "not_planning" | "step_budget" | "cost_budget" | "tool_denied" | "repeated_action"
    }

type PlainRecord = Record<string, unknown>

const STATUSES = new Set<AgentStatus>([
  "planning",
  "validating",
  "awaiting_approval",
  "executing",
  "succeeded",
  "failed",
  "cancelled",
  "exhausted",
])
const CONTROL_CHARACTER = /[\u0000-\u001f\u007f]/u
const MAX_COLLECTION_ITEMS = 256
const MAX_CANONICAL_NODES = 2_048
const MAX_CANONICAL_DEPTH = 12
const MAX_CANONICAL_CODE_UNITS = 65_536
const MAX_ARGUMENT_STRING_LENGTH = 4_096
const MAX_ARGUMENT_KEY_LENGTH = 256
const USD_MICRO_SCALE = 1_000_000

type CanonicalState = {
  nodes: number
  codeUnits: number
  ancestors: Set<object>
}

type ValidationProof = {
  state: AgentState
  stateRevision: string
  costMicros: number
  fingerprint?: string
  requiresApproval: boolean
}

const validationProofs = new WeakMap<object, ValidationProof>()

function plainRecord(value: unknown, label: string): PlainRecord {
  if (value === null || typeof value !== "object" || Array.isArray(value)) {
    throw new TypeError(`${label} must be a plain record`)
  }
  const prototype = Object.getPrototypeOf(value)
  if (prototype !== Object.prototype && prototype !== null) {
    throw new TypeError(`${label} must be a plain record`)
  }
  for (const key of Reflect.ownKeys(value)) {
    if (typeof key !== "string") throw new TypeError(`${label} cannot contain symbol keys`)
    const descriptor = Object.getOwnPropertyDescriptor(value, key)
    if (!descriptor?.enumerable || !("value" in descriptor)) {
      throw new TypeError(`${label} must contain enumerable data properties only`)
    }
  }
  return value as PlainRecord
}

function snapshotRecord(value: unknown, label: string): PlainRecord {
  const record = plainRecord(value, label)
  let snapshot: unknown
  try {
    snapshot = structuredClone(record)
  } catch {
    throw new TypeError(`${label} must be cloneable data, not a proxy or host object`)
  }
  return plainRecord(snapshot, label)
}

function setValues(value: unknown, label: string): string[] {
  if (!(value instanceof Set) || Object.getPrototypeOf(value) !== Set.prototype) {
    throw new TypeError(`${label} must be a built-in Set`)
  }
  if (Reflect.ownKeys(value).length !== 0) {
    throw new TypeError(`${label} cannot override built-in Set behavior`)
  }
  try {
    return [...Set.prototype.values.call(value)] as string[]
  } catch {
    throw new TypeError(`${label} must have built-in Set internal slots`)
  }
}

function mapEntries(value: unknown, label: string): Array<[unknown, unknown]> {
  if (!(value instanceof Map) || Object.getPrototypeOf(value) !== Map.prototype) {
    throw new TypeError(`${label} must be a built-in Map`)
  }
  if (Reflect.ownKeys(value).length !== 0) {
    throw new TypeError(`${label} cannot override built-in Map behavior`)
  }
  try {
    return [...Map.prototype.entries.call(value)] as Array<[unknown, unknown]>
  } catch {
    throw new TypeError(`${label} must have built-in Map internal slots`)
  }
}

function exactKeys(
  value: PlainRecord,
  required: readonly string[],
  optional: readonly string[],
  label: string
): void {
  const allowed = new Set([...required, ...optional])
  for (const key of required) {
    if (!Object.prototype.hasOwnProperty.call(value, key)) {
      throw new TypeError(`${label}.${key} is required`)
    }
  }
  for (const key of Object.keys(value)) {
    if (!allowed.has(key)) throw new TypeError(`${label}.${key} is not allowed`)
  }
}

function text(value: unknown, label: string, maxLength = 4_096): string {
  if (
    typeof value !== "string" ||
    value.length === 0 ||
    value.length > maxLength ||
    value.trim() !== value ||
    CONTROL_CHARACTER.test(value)
  ) {
    throw new TypeError(`${label} must be bounded non-empty text without control characters`)
  }
  return value
}

function numberValue(
  value: unknown,
  label: string,
  { integer = false, positive = false }: { integer?: boolean; positive?: boolean } = {}
): number {
  if (
    typeof value !== "number" ||
    !Number.isFinite(value) ||
    value < 0 ||
    (positive && value === 0) ||
    (integer && !Number.isSafeInteger(value))
  ) {
    throw new TypeError(`${label} must be a finite ${integer ? "integer" : "number"}`)
  }
  return value
}

function usdMicros(value: unknown, label: string, positive = false): number {
  const amount = numberValue(value, label, { positive })
  const micros = Math.round(amount * USD_MICRO_SCALE)
  const normalized = micros / USD_MICRO_SCALE
  const tolerance = Number.EPSILON * Math.max(1, Math.abs(amount)) * 8
  if (
    !Number.isSafeInteger(micros) ||
    (amount > 0 && micros === 0) ||
    Math.abs(amount - normalized) > tolerance
  ) {
    throw new TypeError(`${label} must be an exact non-negative micro-USD amount`)
  }
  return micros
}

function canonicalToken(value: string, state: CanonicalState): string {
  state.codeUnits += value.length
  if (state.codeUnits > MAX_CANONICAL_CODE_UNITS) {
    throw new RangeError("tool arguments exceed the canonicalization boundary")
  }
  return value
}

function compareCodeUnits(left: string, right: string): number {
  return left < right ? -1 : left > right ? 1 : 0
}

function canonicalize(
  value: unknown,
  depth = 0,
  state: CanonicalState = { nodes: 0, codeUnits: 0, ancestors: new Set<object>() }
): string {
  state.nodes += 1
  if (state.nodes > MAX_CANONICAL_NODES || depth > MAX_CANONICAL_DEPTH) {
    throw new RangeError("tool arguments exceed the canonicalization boundary")
  }
  if (value === null) return canonicalToken("null", state)
  if (typeof value === "string") {
    if (value.length > MAX_ARGUMENT_STRING_LENGTH) {
      throw new RangeError("tool argument strings are too large")
    }
    return canonicalToken(JSON.stringify(value), state)
  }
  if (typeof value === "boolean") return canonicalToken(value ? "true" : "false", state)
  if (typeof value === "number") {
    if (!Number.isFinite(value)) throw new TypeError("tool arguments require finite numbers")
    return canonicalToken(Object.is(value, -0) ? "0" : JSON.stringify(value), state)
  }
  if (typeof value !== "object") {
    throw new TypeError("tool arguments must be JSON-compatible data")
  }

  if (state.ancestors.has(value)) throw new TypeError("tool arguments cannot be cyclic")
  state.ancestors.add(value)
  try {
    if (Array.isArray(value)) {
      if (Object.getPrototypeOf(value) !== Array.prototype) {
        throw new TypeError("tool argument arrays must use the built-in Array prototype")
      }
      if (value.length > MAX_COLLECTION_ITEMS) {
        throw new RangeError("tool argument arrays are too large")
      }
      const expectedKeys = new Set([
        ...Array.from({ length: value.length }, (_, index) => String(index)),
        "length",
      ])
      const ownKeys = Reflect.ownKeys(value)
      const invalidDescriptor = Array.from({ length: value.length }, (_, index) =>
        Object.getOwnPropertyDescriptor(value, String(index))
      ).some((descriptor) => !descriptor?.enumerable || !("value" in descriptor))
      if (
        ownKeys.length !== expectedKeys.size ||
        ownKeys.some((key) => typeof key !== "string" || !expectedKeys.has(key)) ||
        invalidDescriptor
      ) {
        throw new TypeError("tool argument arrays must be dense indexed data")
      }
      const children = value.map((child) => canonicalize(child, depth + 1, state))
      canonicalToken("[", state)
      canonicalToken("]", state)
      if (children.length > 1) canonicalToken(",".repeat(children.length - 1), state)
      return `[${children.join(",")}]`
    }

    const record = plainRecord(value, "tool argument")
    const entries = Object.entries(record)
    if (entries.length > MAX_COLLECTION_ITEMS) {
      throw new RangeError("tool argument records are too large")
    }
    const fields = entries
      .sort(([left], [right]) => compareCodeUnits(left, right))
      .map(([key, child]) => {
        if (key.length > MAX_ARGUMENT_KEY_LENGTH) {
          throw new RangeError("tool argument keys are too large")
        }
        const encodedKey = canonicalToken(JSON.stringify(key), state)
        canonicalToken(":", state)
        return `${encodedKey}:${canonicalize(child, depth + 1, state)}`
      })
    canonicalToken("{", state)
    canonicalToken("}", state)
    if (fields.length > 1) canonicalToken(",".repeat(fields.length - 1), state)
    return `{${fields.join(",")}}`
  } finally {
    state.ancestors.delete(value)
  }
}

function toolProposal(value: unknown): Extract<Proposal, { kind: "tool" }> {
  const source = plainRecord(value, "proposal")
  exactKeys(source, ["kind", "name", "args"], [], "proposal")
  const sourceArgs = Object.getOwnPropertyDescriptor(source, "args")?.value
  canonicalize(sourceArgs)
  const record = snapshotRecord(source, "proposal")
  exactKeys(record, ["kind", "name", "args"], [], "proposal")
  if (record.kind !== "tool") throw new TypeError("proposal.kind must be tool")
  const name = text(record.name, "proposal.name", 128)
  const args = plainRecord(record.args, "proposal.args")
  canonicalize(args)
  return { kind: "tool", name, args }
}

function proposalValue(value: unknown): Proposal {
  const source = plainRecord(value, "proposal")
  const record = snapshotRecord(source, "proposal")
  if (record.kind === "tool") return toolProposal(source)
  exactKeys(record, ["kind", "answer"], [], "proposal")
  if (record.kind !== "finish") throw new TypeError("proposal.kind is unsupported")
  return { kind: "finish", answer: text(record.answer, "proposal.answer", 16_384) }
}

function stateValue(value: unknown): AgentState {
  const source = plainRecord(value, "state")
  const record = snapshotRecord(source, "state")
  exactKeys(
    record,
    ["status", "steps", "maxSteps", "spentUsd", "maxCostUsd", "seen"],
    [],
    "state"
  )
  if (typeof record.status !== "string" || !STATUSES.has(record.status as AgentStatus)) {
    throw new TypeError("state.status is unsupported")
  }
  const steps = numberValue(record.steps, "state.steps", { integer: true })
  const maxSteps = numberValue(record.maxSteps, "state.maxSteps", {
    integer: true,
    positive: true,
  })
  const spentMicros = usdMicros(record.spentUsd, "state.spentUsd")
  const maxCostMicros = usdMicros(record.maxCostUsd, "state.maxCostUsd", true)
  if (steps > maxSteps || spentMicros > maxCostMicros) {
    throw new RangeError("state counters cannot exceed their declared budgets")
  }
  const fingerprints = setValues(source.seen, "state.seen")
  if (fingerprints.length > MAX_COLLECTION_ITEMS) throw new RangeError("state.seen is too large")
  for (const fingerprint of fingerprints) {
    text(fingerprint, "state.seen fingerprint", 64)
    if (!/^[a-f0-9]{64}$/u.test(fingerprint)) {
      throw new TypeError("state.seen fingerprints must be SHA-256 digests")
    }
  }
  return source as AgentState
}

function toolPolicies(value: unknown): ReadonlyMap<string, ToolPolicy> {
  const entries = mapEntries(value, "toolPolicies")
  if (entries.length > MAX_COLLECTION_ITEMS) throw new RangeError("toolPolicies is too large")
  const sanitized = new Map<string, ToolPolicy>()
  for (const [nameValue, policyValue] of entries) {
    const name = text(nameValue, "tool policy name", 128)
    const policy = snapshotRecord(policyValue, "tool policy")
    exactKeys(policy, ["requiresApproval"], [], "tool policy")
    if (typeof policy.requiresApproval !== "boolean") {
      throw new TypeError("tool policy requiresApproval must be a boolean")
    }
    sanitized.set(name, Object.freeze({ requiresApproval: policy.requiresApproval }))
  }
  return sanitized
}

function stateRevision(state: AgentState): string {
  return canonicalize([
    state.status,
    state.steps,
    state.maxSteps,
    usdMicros(state.spentUsd, "state.spentUsd"),
    usdMicros(state.maxCostUsd, "state.maxCostUsd", true),
    setValues(state.seen, "state.seen").sort(compareCodeUnits),
  ])
}

function acceptedValidation(
  validation: Extract<Validation, { ok: true }>,
  proof: ValidationProof
): Extract<Validation, { ok: true }> {
  const result = Object.freeze(validation)
  validationProofs.set(result, proof)
  return result
}

export function actionFingerprint(
  tenantId: string,
  proposal: Extract<Proposal, { kind: "tool" }>
): string {
  const tenant = text(tenantId, "tenantId", 256)
  const validated = toolProposal(proposal)
  const payload = canonicalize([tenant, validated.name, validated.args])
  return createHash("sha256").update(payload, "utf8").digest("hex")
}

export function validateProposal(
  state: AgentState,
  tenantId: string,
  proposal: Proposal,
  policies: ReadonlyMap<string, ToolPolicy>,
  estimatedCostUsd = 0
): Validation {
  const current = stateValue(state)
  const tenant = text(tenantId, "tenantId", 256)
  const candidate = proposalValue(proposal)
  const policyMap = toolPolicies(policies)
  const estimatedCostMicros = usdMicros(estimatedCostUsd, "estimatedCostUsd")
  const revision = stateRevision(current)

  if (current.status !== "planning") return { ok: false, reason: "not_planning" }
  if (current.steps >= current.maxSteps) return { ok: false, reason: "step_budget" }
  if (
    usdMicros(current.spentUsd, "state.spentUsd") + estimatedCostMicros >
    usdMicros(current.maxCostUsd, "state.maxCostUsd", true)
  ) {
    return { ok: false, reason: "cost_budget" }
  }
  if (candidate.kind === "finish") {
    return acceptedValidation(
      { ok: true },
      {
        state: current,
        stateRevision: revision,
        costMicros: estimatedCostMicros,
        requiresApproval: false,
      }
    )
  }
  const policy = Map.prototype.get.call(policyMap, candidate.name) as ToolPolicy | undefined
  if (!policy) return { ok: false, reason: "tool_denied" }

  const fingerprint = actionFingerprint(tenant, candidate)
  if (Set.prototype.has.call(current.seen, fingerprint)) {
    return { ok: false, reason: "repeated_action" }
  }
  return acceptedValidation(
    { ok: true, fingerprint, requiresApproval: policy.requiresApproval },
    {
      state: current,
      stateRevision: revision,
      costMicros: estimatedCostMicros,
      fingerprint,
      requiresApproval: policy.requiresApproval,
    }
  )
}

export function recordAcceptedProposal(
  state: AgentState,
  validation: Extract<Validation, { ok: true }>,
  costUsd: number
): AgentState {
  const current = stateValue(state)
  if (current.status !== "planning" || current.steps >= current.maxSteps) {
    throw new RangeError("accepted proposals require available planning state")
  }
  const record = plainRecord(validation, "validation")
  exactKeys(record, ["ok"], ["fingerprint", "requiresApproval"], "validation")
  if (record.ok !== true) throw new TypeError("validation.ok must be true")
  const fingerprint =
    record.fingerprint === undefined
      ? undefined
      : text(record.fingerprint, "validation.fingerprint", 128)
  if (fingerprint !== undefined && !/^[a-f0-9]{64}$/u.test(fingerprint)) {
    throw new TypeError("validation.fingerprint must be a SHA-256 digest")
  }
  if (record.requiresApproval !== undefined && typeof record.requiresApproval !== "boolean") {
    throw new TypeError("validation.requiresApproval must be a boolean")
  }
  if (record.requiresApproval === true && fingerprint === undefined) {
    throw new TypeError("approval requires a tool fingerprint")
  }
  const proof = validationProofs.get(record)
  if (!proof) throw new TypeError("validation was not issued by validateProposal")
  if (
    proof.fingerprint !== fingerprint ||
    proof.requiresApproval !== (record.requiresApproval === true)
  ) {
    throw new TypeError("validation does not match its issued proof")
  }
  if (current !== proof.state || stateRevision(current) !== proof.stateRevision) {
    throw new RangeError("validation was issued for a different state revision")
  }
  const costMicros = usdMicros(costUsd, "costUsd")
  if (costMicros !== proof.costMicros) {
    throw new RangeError("recorded cost must match the validated estimate")
  }
  const spentMicros = usdMicros(current.spentUsd, "state.spentUsd")
  if (spentMicros + costMicros > usdMicros(current.maxCostUsd, "state.maxCostUsd", true)) {
    throw new RangeError("accepted proposal exceeds the cost budget")
  }
  if (fingerprint !== undefined && Set.prototype.has.call(current.seen, fingerprint)) {
    throw new RangeError("accepted proposal fingerprint was already recorded")
  }

  const seen = new Set(setValues(current.seen, "state.seen"))
  if (fingerprint !== undefined) seen.add(fingerprint)
  validationProofs.delete(record)
  return {
    ...current,
    status:
      fingerprint === undefined
        ? "succeeded"
        : record.requiresApproval === true
          ? "awaiting_approval"
          : "executing",
    steps: current.steps + 1,
    spentUsd: (spentMicros + costMicros) / USD_MICRO_SCALE,
    seen,
  }
}