export type Vec2 = { x: number; y: number };

export enum WarehouseAction {
  Up = 0,
  Down = 1,
  Left = 2,
  Right = 3,
  Pickup = 4,
  Drop = 5,
}

export const WAREHOUSE_ACTIONS: Array<{ action: WarehouseAction; label: string; key: string }> = [
  { action: WarehouseAction.Up, label: "Up", key: "↑" },
  { action: WarehouseAction.Down, label: "Down", key: "↓" },
  { action: WarehouseAction.Left, label: "Left", key: "←" },
  { action: WarehouseAction.Right, label: "Right", key: "→" },
  { action: WarehouseAction.Pickup, label: "Pickup", key: "P" },
  { action: WarehouseAction.Drop, label: "Drop", key: "D" },
];

export type Shelf = {
  id: string;
  pos: Vec2;
  items: number;
};

export type WarehouseRewards = {
  stepCost: number;
  invalidActionPenalty: number;
  pickupReward: number;
  dropoffRewardPerItem: number;
};

export type WarehouseConfig = {
  width: number;
  height: number;
  start: Vec2;
  dropoff: Vec2;
  shelves: Shelf[];
  capacity: number;
  maxSteps: number;
  rewards: WarehouseRewards;
};

export type WarehouseState = {
  robot: Vec2;
  carrying: number;
  delivered: number;
  steps: number;
  shelvesItems: number[];
  totalItems: number;
};

export type StepInfo = {
  event:
    | "move"
    | "pickup"
    | "drop"
    | "noop"
    | "invalid"
    | "done"
    | "truncated";
  picked: number;
  dropped: number;
  deliveredTotal: number;
};

export type StepResult = {
  state: WarehouseState;
  reward: number;
  done: boolean;
  truncated: boolean;
  info: StepInfo;
};

export function defaultWarehouseConfig(): WarehouseConfig {
  return {
    width: 10,
    height: 8,
    start: { x: 1, y: 1 },
    dropoff: { x: 8, y: 6 },
    shelves: [
      { id: "s1", pos: { x: 2, y: 3 }, items: 2 },
      { id: "s2", pos: { x: 6, y: 2 }, items: 2 },
      { id: "s3", pos: { x: 7, y: 5 }, items: 1 },
    ],
    capacity: 2,
    maxSteps: 200,
    rewards: {
      stepCost: -0.1,
      invalidActionPenalty: -0.5,
      pickupReward: 0.2,
      dropoffRewardPerItem: 5,
    },
  };
}

function clampInt(v: number, min: number, max: number): number {
  if (!Number.isFinite(v)) return min;
  return Math.max(min, Math.min(max, Math.trunc(v)));
}

export function normalizeWarehouseConfig(config: WarehouseConfig): WarehouseConfig {
  const width = clampInt(config.width, 3, 30);
  const height = clampInt(config.height, 3, 30);
  const capacity = clampInt(config.capacity, 0, 20);
  const maxSteps = clampInt(config.maxSteps, 1, 5000);

  const start = {
    x: clampInt(config.start.x, 0, width - 1),
    y: clampInt(config.start.y, 0, height - 1),
  };
  const dropoff = {
    x: clampInt(config.dropoff.x, 0, width - 1),
    y: clampInt(config.dropoff.y, 0, height - 1),
  };

  const shelves: Shelf[] = [];
  const seen = new Set<string>();
  for (const shelf of config.shelves) {
    const x = clampInt(shelf.pos.x, 0, width - 1);
    const y = clampInt(shelf.pos.y, 0, height - 1);
    const key = `${x},${y}`;
    if (key === `${start.x},${start.y}`) continue;
    if (key === `${dropoff.x},${dropoff.y}`) continue;
    if (seen.has(key)) continue;
    seen.add(key);
    shelves.push({
      id: shelf.id || `s${shelves.length + 1}`,
      pos: { x, y },
      items: clampInt(shelf.items, 0, 50),
    });
  }

  return {
    width,
    height,
    start,
    dropoff,
    shelves,
    capacity,
    maxSteps,
    rewards: {
      stepCost: Number.isFinite(config.rewards.stepCost) ? config.rewards.stepCost : -0.1,
      invalidActionPenalty: Number.isFinite(config.rewards.invalidActionPenalty)
        ? config.rewards.invalidActionPenalty
        : -0.5,
      pickupReward: Number.isFinite(config.rewards.pickupReward) ? config.rewards.pickupReward : 0.2,
      dropoffRewardPerItem: Number.isFinite(config.rewards.dropoffRewardPerItem)
        ? config.rewards.dropoffRewardPerItem
        : 5,
    },
  };
}

export function initialWarehouseState(config: WarehouseConfig): WarehouseState {
  const shelvesItems = config.shelves.map((s) => clampInt(s.items, 0, 50));
  const totalItems = shelvesItems.reduce((a, b) => a + b, 0);
  return {
    robot: { ...config.start },
    carrying: 0,
    delivered: 0,
    steps: 0,
    shelvesItems,
    totalItems,
  };
}

export function isTerminal(state: WarehouseState): boolean {
  return state.delivered >= state.totalItems;
}

export function stateKey(state: WarehouseState): string {
  // Key is specific to the current config (shelf ordering), which is fine for a tabular demo.
  return `${state.robot.x},${state.robot.y}|c${state.carrying}|d${state.delivered}|s${state.shelvesItems.join(",")}`;
}

export function stepWarehouse(configRaw: WarehouseConfig, state: WarehouseState, action: WarehouseAction): StepResult {
  const config = normalizeWarehouseConfig(configRaw);

  if (isTerminal(state)) {
    return {
      state,
      reward: 0,
      done: true,
      truncated: false,
      info: { event: "done", picked: 0, dropped: 0, deliveredTotal: state.delivered },
    };
  }

  if (state.steps >= config.maxSteps) {
    return {
      state,
      reward: 0,
      done: false,
      truncated: true,
      info: { event: "truncated", picked: 0, dropped: 0, deliveredTotal: state.delivered },
    };
  }

  let reward = config.rewards.stepCost;
  let nextRobot = state.robot;
  let carrying = state.carrying;
  let delivered = state.delivered;
  let shelvesItems = state.shelvesItems;
  let event: StepInfo["event"] = "noop";
  let picked = 0;
  let dropped = 0;

  const applyInvalid = () => {
    reward += config.rewards.invalidActionPenalty;
    event = "invalid";
  };

  if (action === WarehouseAction.Up) {
    if (state.robot.y > 0) {
      nextRobot = { x: state.robot.x, y: state.robot.y - 1 };
      event = "move";
    } else {
      applyInvalid();
    }
  } else if (action === WarehouseAction.Down) {
    if (state.robot.y < config.height - 1) {
      nextRobot = { x: state.robot.x, y: state.robot.y + 1 };
      event = "move";
    } else {
      applyInvalid();
    }
  } else if (action === WarehouseAction.Left) {
    if (state.robot.x > 0) {
      nextRobot = { x: state.robot.x - 1, y: state.robot.y };
      event = "move";
    } else {
      applyInvalid();
    }
  } else if (action === WarehouseAction.Right) {
    if (state.robot.x < config.width - 1) {
      nextRobot = { x: state.robot.x + 1, y: state.robot.y };
      event = "move";
    } else {
      applyInvalid();
    }
  } else if (action === WarehouseAction.Pickup) {
    const shelfIndex = config.shelves.findIndex(
      (s) => s.pos.x === state.robot.x && s.pos.y === state.robot.y
    );
    if (shelfIndex < 0) {
      applyInvalid();
    } else if (carrying >= config.capacity) {
      applyInvalid();
    } else if ((state.shelvesItems[shelfIndex] ?? 0) <= 0) {
      applyInvalid();
    } else {
      const nextShelves = state.shelvesItems.slice();
      nextShelves[shelfIndex] = clampInt(nextShelves[shelfIndex] - 1, 0, 50);
      shelvesItems = nextShelves;
      carrying += 1;
      picked = 1;
      reward += config.rewards.pickupReward;
      event = "pickup";
    }
  } else if (action === WarehouseAction.Drop) {
    const atDropoff = state.robot.x === config.dropoff.x && state.robot.y === config.dropoff.y;
    if (!atDropoff) {
      applyInvalid();
    } else if (carrying <= 0) {
      applyInvalid();
    } else {
      dropped = carrying;
      delivered += carrying;
      carrying = 0;
      reward += dropped * config.rewards.dropoffRewardPerItem;
      event = "drop";
    }
  }

  const nextState: WarehouseState = {
    robot: nextRobot,
    carrying,
    delivered,
    steps: state.steps + 1,
    shelvesItems,
    totalItems: state.totalItems,
  };

  const done = isTerminal(nextState);
  const truncated = nextState.steps >= config.maxSteps && !done;

  return {
    state: nextState,
    reward,
    done,
    truncated,
    info: {
      event: truncated ? "truncated" : done ? "done" : event,
      picked,
      dropped,
      deliveredTotal: delivered,
    },
  };
}