import { Plan } from "../interiors/plan.ts";
import {
  createRobotActivity,
  type RobotActivityController,
  type RobotActivitySnapshot,
  type RobotActivityState,
} from "../interiors/robotActivity.ts";
import {
  resolveRobotOperations,
  type ResolvedRobotOperations,
  type RobotJobKind,
} from "../interiors/robotOperations.ts";
import { LUMBRIDGE_HQ } from "../offices/lumbridge-hq.ts";
import { MATEO_COURT } from "../offices/mateo-court.ts";
import { LUMBRIDGE_HQ_ROBOT_OPERATIONS } from "../offices/operations/lumbridge-hq.ts";
import { MATEO_COURT_ROBOT_OPERATIONS } from "../offices/operations/mateo-court.ts";
import { BaseArenaEnvironment, type SimulationTransition } from "./base.ts";
import { ARENA_SOURCE_HASHES } from "./sourceHashes.ts";
import { ArenaScenarioRegistry } from "./scenarios.ts";
import { ARENA_API_VERSION, type ArenaFieldSpec, type ArenaManifest, type ArenaScenario } from "./types.ts";

export interface OfficeJobsAction {
  x: number;
  z: number;
  interact: boolean;
}

export interface OfficeJobsObservation {
  officeId: string;
  robotId: string;
  levelId: string;
  x: number;
  z: number;
  mode: string;
  phase: string;
  jobKind: RobotJobKind | "none";
  payload: "parcel" | null;
  battery: number;
  jobProgress: number;
  nextStationId: string;
  nextX: number;
  nextZ: number;
  deltaX: number;
  deltaZ: number;
  distanceToNextM: number;
  canInteract: boolean;
  blockedStreak: number;
  recoveryCount: number;
  completedJobs: number;
}

export type OfficeJobsReward = Record<
  "navigation" | "job" | "success" | "time" | "control" | "collision" | "interaction" | "safety",
  number
>;

interface OfficeJobsScenarioParameters {
  officeId: "lumbridge-hq" | "mateo-court";
  robotId: string;
  jobId: string;
}

interface OfficeJobsSimulationSnapshot {
  activity: RobotActivitySnapshot;
  previousDistanceM: number;
  previousProgress: number;
  blockedStreak: number;
  wrongInteractions: number;
  initialCompletedJobs: number;
}

const MAX_STEPS = 720;
// Slightly inside robotActivity's 0.30 m arrival radius so an agent never sees
// an interaction affordance the authoritative controller would still reject.
const INTERACTION_RADIUS_M = 0.295;
const WRONG_INTERACTION_LIMIT = 8;

const SF_PLAN = new Plan(LUMBRIDGE_HQ, { depth: "public", warn: false });
const LA_PLAN = new Plan(MATEO_COURT, { depth: "public", warn: false });
const SF_OPERATIONS = resolveRobotOperations(SF_PLAN, LUMBRIDGE_HQ_ROBOT_OPERATIONS);
const LA_OPERATIONS = resolveRobotOperations(LA_PLAN, MATEO_COURT_ROBOT_OPERATIONS);

const DEFINITIONS = [
  {
    id: "train-sf-display-inspection",
    split: "train" as const,
    parameters: { officeId: "lumbridge-hq" as const, robotId: "sf-studio-activity-01", jobId: "sf-studio-inspect-display" },
  },
  {
    id: "train-la-directory-inspection",
    split: "train" as const,
    parameters: { officeId: "mateo-court" as const, robotId: "la-office-activity-01", jobId: "la-l1-inspect-directory" },
  },
  {
    id: "dev-sf-parcel-delivery",
    split: "dev" as const,
    parameters: { officeId: "lumbridge-hq" as const, robotId: "sf-studio-activity-01", jobId: "sf-studio-deliver" },
  },
  {
    id: "dev-la-loft-delivery",
    split: "dev" as const,
    parameters: { officeId: "mateo-court" as const, robotId: "la-office-activity-02", jobId: "la-l2-deliver" },
  },
] as const;

export const OFFICE_JOBS_SCENARIOS = new ArenaScenarioRegistry<OfficeJobsScenarioParameters>(
  "office-jobs-v1",
  DEFINITIONS,
  (parameters) => ({ ...parameters }),
);

export const OFFICE_JOBS_MANIFEST: ArenaManifest = Object.freeze({
  apiVersion: ARENA_API_VERSION,
  id: "office-jobs-v1",
  // 2 rather than 1: `ArenaScenarioRegistry` now selects a bare `{ split }`
  // request by hashing the seed against each scenario id instead of indexing
  // definition order, so a seed run before that change may resolve to a
  // different scenario after it. `version` is what a snapshot, a trace and a
  // results table are pinned to, and a selection change that nothing recorded
  // is exactly the silent remap the new selector exists to prevent.
  version: 2,
  title: "Seeded office robot jobs",
  description: "Headless job execution over explicit simulated SF/LA operations and resolved office collision.",
  simulator: "Plan + robotRoutes + fixed-step robotActivity",
  fixedStepSeconds: 0.1,
  maxSteps: MAX_STEPS,
  actionFields: ["x", "z", "interact"],
  observationFields: [
    "officeId", "robotId", "levelId", "x", "z", "mode", "phase", "jobKind", "payload",
    "battery", "jobProgress", "nextStationId", "nextX", "nextZ", "deltaX", "deltaZ",
    "distanceToNextM", "canInteract", "blockedStreak", "recoveryCount", "completedJobs",
  ],
  actionSpace: [
    { name: "x", kind: "float", low: -1, high: 1, unit: "normalized drive demand" },
    { name: "z", kind: "float", low: -1, high: 1, unit: "normalized drive demand" },
    { name: "interact", kind: "bool" },
  ] satisfies readonly ArenaFieldSpec[],
  // `mode`, `phase`, `jobKind` and `payload` are closed vocabularies in
  // `robotActivity.ts` and `robotOperations.ts` and are one-hot here rather than
  // hashed: a controller with five modes that a trainer sees as five unrelated
  // reals is a controller a trainer cannot condition on.
  observationSpace: [
    { name: "officeId", kind: "enum", values: ["lumbridge-hq", "mateo-court"] },
    { name: "robotId", kind: "id" },
    { name: "levelId", kind: "enum", values: ["level-1", "level-2"] },
    { name: "x", kind: "float", low: 0, high: 40, unit: "m" },
    { name: "z", kind: "float", low: 0, high: 40, unit: "m" },
    // `RobotActivityMode` in robotOperations.ts, member for member.
    {
      name: "mode",
      kind: "enum",
      values: ["idle", "patrol", "deliver", "inspect", "charge", "blocked-recovery"],
    },
    // Every string `navigationPhase`, `activityPhase` and the recovery paths in
    // robotActivity.ts can assign. `phase` is typed `string` there rather than
    // as a union, so this list is a transcription and the spaces test walks a
    // rollout of every scenario asserting nothing outside it is ever observed.
    {
      name: "phase",
      kind: "enum",
      values: [
        "scheduled-idle", "patrolling", "patrol-check", "to-pickup", "loading-parcel",
        "to-dropoff", "delivering-parcel", "to-inspection", "inspecting", "to-charge",
        "charging", "awaiting-interaction", "replanning", "backoff-and-replan", "terminal",
      ],
    },
    // `RobotJobKind` plus the `"none"` this observation substitutes for `null`.
    {
      name: "jobKind",
      kind: "enum",
      values: ["none", "patrol", "deliver", "inspect", "charge"],
    },
    // One member and a legal absence: `payload` is `"parcel" | null`, and a null
    // encodes as the zero vector rather than as a second category. That is the
    // documented meaning of an out-of-vocabulary value and it is the right one
    // here — "carrying nothing" is not a thing being carried.
    { name: "payload", kind: "enum", values: ["parcel"] },
    { name: "battery", kind: "float", low: 0, high: 1, unit: "fraction" },
    { name: "jobProgress", kind: "float", low: 0, high: 1, unit: "fraction" },
    { name: "nextStationId", kind: "id" },
    { name: "nextX", kind: "float", low: 0, high: 40, unit: "m" },
    { name: "nextZ", kind: "float", low: 0, high: 40, unit: "m" },
    { name: "deltaX", kind: "float", low: -40, high: 40, unit: "m" },
    { name: "deltaZ", kind: "float", low: -40, high: 40, unit: "m" },
    { name: "distanceToNextM", kind: "float", low: 0, high: 60, unit: "m" },
    { name: "canInteract", kind: "bool" },
    { name: "blockedStreak", kind: "float", low: 0, high: 720, unit: "steps" },
    { name: "recoveryCount", kind: "float", low: 0, high: 8, unit: "count" },
    { name: "completedJobs", kind: "float", low: 0, high: 32, unit: "count" },
  ] satisfies readonly ArenaFieldSpec[],
  rewardComponents: {
    navigation: "Bounded reduction in distance to the current authored job station.",
    job: "Dense progress through pickup, delivery, inspection, or charge phases.",
    success: "Sparse bonus for completing the assigned job.",
    time: "Per-step pressure that keeps inaction below zero.",
    control: "Small cost on planar action magnitude.",
    collision: "Cost for commanded movement that fails against resolved collision.",
    interaction: "Penalty for interaction away from the current work station.",
    safety: "Terminal collision-stall or repeated wrong-interaction penalty.",
  },
  safetyTerminals: ["collision-stall", "wrong-interaction-limit"],
  scenarioIds: {
    train: OFFICE_JOBS_SCENARIOS.ids("train"),
    dev: OFFICE_JOBS_SCENARIOS.ids("dev"),
  },
  baselines: {
    inaction: "x=0,z=0,interact=false never completes a job and accumulates the time floor.",
    scripted: "Follow the exposed current route waypoint and interact only inside the station radius.",
  },
});

export const OFFICE_JOBS_INACTION: Readonly<OfficeJobsAction> = Object.freeze({
  x: 0,
  z: 0,
  interact: false,
});

export function officeJobsScriptedBaseline(observation: OfficeJobsObservation): OfficeJobsAction {
  if (observation.canInteract) return { x: 0, z: 0, interact: true };
  const length = Math.sqrt(
    observation.deltaX * observation.deltaX + observation.deltaZ * observation.deltaZ,
  );
  if (length <= 1e-9) return { x: 0, z: 0, interact: false };
  const magnitude = Math.min(1, length / (1.05 * OFFICE_JOBS_MANIFEST.fixedStepSeconds));
  return {
    x: observation.deltaX / length * magnitude,
    z: observation.deltaZ / length * magnitude,
    interact: false,
  };
}

export class OfficeJobsEnvironment extends BaseArenaEnvironment<
  OfficeJobsAction,
  OfficeJobsObservation,
  OfficeJobsReward,
  OfficeJobsSimulationSnapshot,
  OfficeJobsScenarioParameters
> {
  readonly manifest = OFFICE_JOBS_MANIFEST;
  protected readonly registry = OFFICE_JOBS_SCENARIOS;
  protected readonly sourceHashes = ARENA_SOURCE_HASHES["office-jobs-v1"]!;
  private activity: RobotActivityController | null = null;
  private operations: ResolvedRobotOperations | null = null;
  private robotId = "";
  private previousDistanceM = 0;
  private previousProgress = 0;
  private blockedStreak = 0;
  private wrongInteractions = 0;
  private initialCompletedJobs = 0;

  protected resetSimulation(
    scenario: ArenaScenario<OfficeJobsScenarioParameters>,
  ): OfficeJobsObservation {
    const selected = setupFor(scenario.parameters.officeId);
    this.operations = selected.operations;
    this.robotId = scenario.parameters.robotId;
    this.activity = createRobotActivity(selected.plan, selected.operations, {
      seed: scenario.seed,
      robotIds: [this.robotId],
      controlledRobotIds: [this.robotId],
      initialJobIds: { [this.robotId]: scenario.parameters.jobId },
    });
    const robot = this.robot();
    this.previousDistanceM = this.distanceToNext(robot);
    this.previousProgress = robot.progress;
    this.blockedStreak = 0;
    this.wrongInteractions = 0;
    this.initialCompletedJobs = robot.completedJobs;
    return this.observation();
  }

  protected normalizeAction(action: OfficeJobsAction): OfficeJobsAction {
    const x = Number.isFinite(action?.x) ? action.x : 0;
    const z = Number.isFinite(action?.z) ? action.z : 0;
    const length = Math.sqrt(x * x + z * z);
    return {
      x: length > 1 ? x / length : x,
      z: length > 1 ? z / length : z,
      interact: action?.interact === true,
    };
  }

  protected advanceSimulation(
    action: OfficeJobsAction,
  ): SimulationTransition<OfficeJobsObservation, OfficeJobsReward> {
    const before = this.robot();
    const beforePosition = { ...before.position };
    const canInteract = this.canInteract(before);
    if (action.interact && !canInteract) this.wrongInteractions += 1;
    const after = this.requireActivity().step({ [this.robotId]: action }).robots[0]!;
    const movedX = after.position.x - beforePosition.x;
    const movedZ = after.position.z - beforePosition.z;
    const moved = Math.sqrt(movedX * movedX + movedZ * movedZ);
    const demand = Math.sqrt(action.x * action.x + action.z * action.z);
    const blocked = demand > 0.2 && moved < 0.012;
    this.blockedStreak = blocked ? this.blockedStreak + 1 : 0;
    const distance = this.distanceToNext(after);
    const navigationProgress = Math.max(-0.3, Math.min(0.3, this.previousDistanceM - distance));
    const jobProgress = Math.max(-0.5, Math.min(0.5, after.progress - this.previousProgress));
    this.previousDistanceM = distance;
    this.previousProgress = after.progress;
    const success = after.completedJobs > this.initialCompletedJobs;
    const collisionStall = after.mode === "blocked-recovery";
    const wrongInteractionFailure = this.wrongInteractions >= WRONG_INTERACTION_LIMIT;
    const controllerFailure = after.terminalReason;
    const terminated = success || collisionStall || wrongInteractionFailure || controllerFailure !== null;
    const terminalReason = success
      ? "job-complete"
      : collisionStall
        ? "collision-stall"
        : wrongInteractionFailure
          ? "wrong-interaction-limit"
          : controllerFailure ?? undefined;
    return {
      observation: this.observation(),
      rewardComponents: {
        navigation: navigationProgress * 0.18,
        job: jobProgress * 1.2,
        success: success ? 8 : 0,
        time: -0.006,
        control: -0.001 * demand ** 2,
        collision: blocked ? -0.08 : 0,
        interaction: action.interact && !canInteract ? -0.12 : 0,
        safety: collisionStall || wrongInteractionFailure || controllerFailure ? -2 : 0,
      },
      terminated,
      terminalReason,
    };
  }

  protected simulationSnapshot(): OfficeJobsSimulationSnapshot {
    return {
      activity: this.requireActivity().snapshot(),
      previousDistanceM: this.previousDistanceM,
      previousProgress: this.previousProgress,
      blockedStreak: this.blockedStreak,
      wrongInteractions: this.wrongInteractions,
      initialCompletedJobs: this.initialCompletedJobs,
    };
  }

  protected restoreSimulation(snapshot: OfficeJobsSimulationSnapshot): OfficeJobsObservation {
    this.resetSimulation(this.currentScenario());
    if (
      !Number.isFinite(snapshot.previousDistanceM) || !Number.isFinite(snapshot.previousProgress) ||
      !Number.isSafeInteger(snapshot.blockedStreak) || snapshot.blockedStreak < 0 ||
      !Number.isSafeInteger(snapshot.wrongInteractions) || snapshot.wrongInteractions < 0 ||
      !Number.isSafeInteger(snapshot.initialCompletedJobs) || snapshot.initialCompletedJobs < 0
    ) throw new Error("office jobs simulation snapshot is invalid");
    this.requireActivity().restore(snapshot.activity);
    this.previousDistanceM = snapshot.previousDistanceM;
    this.previousProgress = snapshot.previousProgress;
    this.blockedStreak = snapshot.blockedStreak;
    this.wrongInteractions = snapshot.wrongInteractions;
    this.initialCompletedJobs = snapshot.initialCompletedJobs;
    return this.observation();
  }

  private observation(): OfficeJobsObservation {
    const robot = this.robot();
    const station = this.nextStation(robot);
    const waypoint = robot.route[robot.routeIndex] ?? station?.position ?? robot.position;
    const distance = station ? this.distanceToNext(robot) : 0;
    return {
      officeId: this.requireOperations().definition.officeId,
      robotId: robot.id,
      levelId: robot.levelId,
      x: robot.position.x,
      z: robot.position.z,
      mode: robot.mode,
      phase: robot.phase,
      jobKind: robot.activeJobKind ?? "none",
      payload: robot.payload,
      battery: robot.battery,
      jobProgress: robot.progress,
      nextStationId: station?.id ?? "",
      nextX: waypoint.x,
      nextZ: waypoint.z,
      deltaX: waypoint.x - robot.position.x,
      deltaZ: waypoint.z - robot.position.z,
      distanceToNextM: distance,
      canInteract: this.canInteract(robot),
      blockedStreak: this.blockedStreak,
      recoveryCount: robot.recoveryCount,
      completedJobs: robot.completedJobs,
    };
  }

  private nextStation(robot: RobotActivityState) {
    const id = robot.activeStationIds[robot.stationIndex];
    return id ? this.requireOperations().stations.get(id) : undefined;
  }

  private distanceToNext(robot: RobotActivityState): number {
    const station = this.nextStation(robot);
    if (!station) return 0;
    const dx = station.position.x - robot.position.x;
    const dz = station.position.z - robot.position.z;
    return Math.sqrt(dx * dx + dz * dz);
  }

  private canInteract(robot: RobotActivityState): boolean {
    return this.distanceToNext(robot) <= INTERACTION_RADIUS_M;
  }

  private robot(): RobotActivityState {
    const robot = this.requireActivity().states()[0];
    if (!robot) throw new Error("office jobs environment has no controlled robot");
    return robot;
  }

  private requireActivity(): RobotActivityController {
    if (!this.activity) throw new Error("office jobs environment has not been reset");
    return this.activity;
  }

  private requireOperations(): ResolvedRobotOperations {
    if (!this.operations) throw new Error("office jobs environment has not been reset");
    return this.operations;
  }
}

function setupFor(officeId: OfficeJobsScenarioParameters["officeId"]) {
  return officeId === "lumbridge-hq"
    ? { plan: SF_PLAN, operations: SF_OPERATIONS }
    : { plan: LA_PLAN, operations: LA_OPERATIONS };
}