import { Plan } from "../interiors/plan.ts";
import {
  createWalker,
  type WalkerController,
  type WalkerState,
} from "../interiors/walker.ts";
import { FRONTIER_VALLEY } from "../offices/frontier-valley.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 OfficeNavAction {
  x: number;
  z: number;
}

export interface OfficeNavObservation {
  levelId: string;
  x: number;
  z: number;
  goalX: number;
  goalZ: number;
  deltaX: number;
  deltaZ: number;
  distanceToGoalM: number;
  travelledM: number;
  blockedStreak: number;
}

export type OfficeNavReward = Record<
  "progress" | "success" | "time" | "control" | "collision" | "safety",
  number
>;

interface OfficeScenarioParameters {
  levelId: string;
  startX: number;
  startZ: number;
  goalX: number;
  goalZ: number;
}

interface OfficeSimulationSnapshot {
  walker: WalkerState;
  previousGoalDistanceM: number;
  blockedStreak: number;
}

const FIXED_STEP = 0.1;
const MAX_STEPS = 160;
const SUCCESS_RADIUS_M = 0.45;
const PLAN = new Plan(FRONTIER_VALLEY, { depth: "public", warn: false });

const DEFINITIONS = [
  {
    id: "train-hangar-crossing",
    split: "train" as const,
    parameters: { levelId: "level-1", startX: 20, startZ: 18, goalX: 34, goalZ: 18 },
  },
  {
    id: "train-galley-aisle",
    split: "train" as const,
    parameters: { levelId: "level-1", startX: 18, startZ: 21, goalX: 18, goalZ: 27 },
  },
  {
    id: "dev-apron-approach",
    split: "dev" as const,
    parameters: { levelId: "level-1", startX: 36, startZ: 19, goalX: 46, goalZ: 19 },
  },
  {
    id: "dev-south-aisle",
    split: "dev" as const,
    parameters: { levelId: "level-1", startX: 22, startZ: 20, goalX: 32, goalZ: 25 },
  },
] as const;

export const OFFICE_NAV_SCENARIOS = new ArenaScenarioRegistry<OfficeScenarioParameters>(
  "office-nav-v1",
  DEFINITIONS,
  (base, random) => {
    const offsetX = random.between(-0.08, 0.08);
    const offsetZ = random.between(-0.08, 0.08);
    return {
      ...base,
      startX: base.startX + offsetX,
      startZ: base.startZ + offsetZ,
      goalX: base.goalX - offsetX,
      goalZ: base.goalZ - offsetZ,
    };
  },
);

export const OFFICE_NAV_MANIFEST: ArenaManifest = Object.freeze({
  apiVersion: ARENA_API_VERSION,
  id: "office-nav-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: "Frontier Valley office navigation",
  description: "Headless navigation through the resolved public office plan and exact wall collision.",
  simulator: "Plan(FRONTIER_VALLEY) + createWalker",
  fixedStepSeconds: FIXED_STEP,
  maxSteps: MAX_STEPS,
  actionFields: ["x", "z"],
  observationFields: [
    "levelId", "x", "z", "goalX", "goalZ", "deltaX", "deltaZ",
    "distanceToGoalM", "travelledM", "blockedStreak",
  ],
  actionSpace: [
    { name: "x", kind: "float", low: -1, high: 1, unit: "normalized walk demand" },
    { name: "z", kind: "float", low: -1, high: 1, unit: "normalized walk demand" },
  ] satisfies readonly ArenaFieldSpec[],
  // Office-world metres, and the bounds are the reference pack's floor plate
  // with slack: Frontier Valley's authored waypoints run from x=18 to x=46, and
  // a walker resolved against collision cannot leave the building.
  observationSpace: [
    { name: "levelId", kind: "id" },
    { name: "x", kind: "float", low: 0, high: 80, unit: "m" },
    { name: "z", kind: "float", low: 0, high: 80, unit: "m" },
    { name: "goalX", kind: "float", low: 0, high: 80, unit: "m" },
    { name: "goalZ", kind: "float", low: 0, high: 80, unit: "m" },
    { name: "deltaX", kind: "float", low: -80, high: 80, unit: "m" },
    { name: "deltaZ", kind: "float", low: -80, high: 80, unit: "m" },
    { name: "distanceToGoalM", kind: "float", low: 0, high: 120, unit: "m" },
    { name: "travelledM", kind: "float", low: 0, high: 400, unit: "m" },
    { name: "blockedStreak", kind: "float", low: 0, high: 160, unit: "steps" },
  ] satisfies readonly ArenaFieldSpec[],
  rewardComponents: {
    progress: "Reduction in Euclidean distance to the goal.",
    success: "Sparse arrival bonus.",
    time: "Negative inaction floor and path-efficiency pressure.",
    control: "Small cost on action magnitude.",
    collision: "Cost when commanded travel is blocked by the resolved plan.",
    safety: "Terminal repeated-collision penalty.",
  },
  safetyTerminals: ["collision-stall"],
  scenarioIds: {
    train: OFFICE_NAV_SCENARIOS.ids("train"),
    dev: OFFICE_NAV_SCENARIOS.ids("dev"),
  },
  baselines: {
    inaction: "x=0,z=0 reaches no goal and accumulates the time floor.",
    scripted: "A normalized direct-to-goal vector completes all collision-clear public scenarios.",
  },
});

export const OFFICE_NAV_INACTION: Readonly<OfficeNavAction> = Object.freeze({ x: 0, z: 0 });

export function officeNavScriptedBaseline(observation: OfficeNavObservation): OfficeNavAction {
  const length = Math.sqrt(
    observation.deltaX * observation.deltaX + observation.deltaZ * observation.deltaZ,
  );
  if (length === 0) return { x: 0, z: 0 };
  return { x: observation.deltaX / length, z: observation.deltaZ / length };
}

export class OfficeNavEnvironment extends BaseArenaEnvironment<
  OfficeNavAction,
  OfficeNavObservation,
  OfficeNavReward,
  OfficeSimulationSnapshot,
  OfficeScenarioParameters
> {
  readonly manifest = OFFICE_NAV_MANIFEST;
  protected readonly registry = OFFICE_NAV_SCENARIOS;
  protected readonly sourceHashes = ARENA_SOURCE_HASHES["office-nav-v1"]!;
  private walker: WalkerController | null = null;
  private previousGoalDistanceM = 0;
  private blockedStreak = 0;

  protected resetSimulation(scenario: ArenaScenario<OfficeScenarioParameters>): OfficeNavObservation {
    const parameters = scenario.parameters;
    this.walker = createWalker(PLAN, {
      levelId: parameters.levelId,
      position: { x: parameters.startX, z: parameters.startZ },
      speed: 2,
      fixedStep: FIXED_STEP,
      maxCatchUpSteps: 1,
    });
    this.previousGoalDistanceM = this.goalDistance();
    this.blockedStreak = 0;
    return this.observation();
  }

  protected normalizeAction(action: OfficeNavAction): OfficeNavAction {
    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 length > 1 ? { x: x / length, z: z / length } : { x, z };
  }

  protected advanceSimulation(
    action: OfficeNavAction,
  ): SimulationTransition<OfficeNavObservation, OfficeNavReward> {
    const walker = this.requireWalker();
    const before = walker.state();
    const after = walker.tick(FIXED_STEP, action);
    const movedX = after.position.x - before.position.x;
    const movedZ = after.position.z - before.position.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 < demand * 2 * FIXED_STEP * 0.2;
    this.blockedStreak = blocked ? this.blockedStreak + 1 : 0;
    const distance = this.goalDistance();
    const progress = this.previousGoalDistanceM - distance;
    this.previousGoalDistanceM = distance;
    const success = distance <= SUCCESS_RADIUS_M;
    const failure = this.blockedStreak >= 8;
    return {
      observation: this.observation(),
      rewardComponents: {
        progress: progress * 0.22,
        success: success ? 2.5 : 0,
        time: -0.01,
        control: -0.001 * demand ** 2,
        collision: blocked ? -0.08 : 0,
        safety: failure ? -1.5 : 0,
      },
      terminated: success || failure,
      terminalReason: success ? "goal" : failure ? "collision-stall" : undefined,
    };
  }

  protected simulationSnapshot(): OfficeSimulationSnapshot {
    return {
      walker: this.requireWalker().state(),
      previousGoalDistanceM: this.previousGoalDistanceM,
      blockedStreak: this.blockedStreak,
    };
  }

  protected restoreSimulation(snapshot: OfficeSimulationSnapshot): OfficeNavObservation {
    this.resetSimulation(this.currentScenario());
    this.requireWalker().restore(snapshot.walker);
    this.previousGoalDistanceM = snapshot.previousGoalDistanceM;
    this.blockedStreak = snapshot.blockedStreak;
    return this.observation();
  }

  private observation(): OfficeNavObservation {
    const state = this.requireWalker().state();
    const goal = this.currentScenario().parameters;
    return {
      levelId: state.levelId,
      x: state.position.x,
      z: state.position.z,
      goalX: goal.goalX,
      goalZ: goal.goalZ,
      deltaX: goal.goalX - state.position.x,
      deltaZ: goal.goalZ - state.position.z,
      distanceToGoalM: this.goalDistance(),
      travelledM: state.distance,
      blockedStreak: this.blockedStreak,
    };
  }

  private goalDistance(): number {
    const state = this.requireWalker().state();
    const goal = this.currentScenario().parameters;
    const dx = goal.goalX - state.position.x;
    const dz = goal.goalZ - state.position.z;
    return Math.sqrt(dx * dx + dz * dz);
  }

  private requireWalker(): WalkerController {
    if (!this.walker) throw new Error("office environment has not been reset");
    return this.walker;
  }
}
