// 角色分配 + role list 展开。RNG 注入，便于测试。

export type RoleDict = Record<string, number>
export type RoleList = string[]
export type Preferences = Record<number, string[]> // player_id (1-based) -> 偏好角色列表

export interface Rng {
  // [0, 1)
  next(): number
}

export const defaultRng: Rng = { next: () => Math.random() }

/** 把 {狼人: 2, 村民: 1} 展开成 ['狼人', '狼人', '村民']。已是 list 则原样返回。 */
export function expandRolesToList(roles: RoleDict | RoleList): RoleList {
  if (Array.isArray(roles)) return [...roles]
  const out: RoleList = []
  for (const [role, count] of Object.entries(roles)) {
    for (let i = 0; i < count; i++) out.push(role)
  }
  return out
}

/** 缩回 dict，方便存历史。 */
export function rolesListToDict(roles: RoleList): RoleDict {
  const out: RoleDict = {}
  for (const r of roles) out[r] = (out[r] || 0) + 1
  return out
}

// Fisher-Yates shuffle，用注入的 rng 保证可测。
export function shuffle<T>(arr: T[], rng: Rng = defaultRng): T[] {
  const a = [...arr]
  for (let i = a.length - 1; i > 0; i--) {
    const j = Math.floor(rng.next() * (i + 1))
    ;[a[i], a[j]] = [a[j], a[i]]
  }
  return a
}

function weightedPick(weighted: string[], rng: Rng): string {
  // weighted 是已经按权重展开的 list，直接随机选
  const idx = Math.floor(rng.next() * weighted.length)
  return weighted[idx]
}

/**
 * 按偏好分配角色。
 *
 * 算法（从 partiverse python 同步）：
 *  - 先随机化玩家顺序（保证公平 — 谁都可能先选）
 *  - 维护剩余角色池
 *  - 对每个玩家：偏好角色 4x 权重，其他 1x，加权随机抽一个
 *  - 抽完从池里扣除
 *
 * 返回 list：index 0 = player 1 的角色。
 */
export function assignRolesWithPreferences(
  rolesList: RoleList,
  preferences: Preferences,
  rng: Rng = defaultRng,
): RoleList {
  const numPlayers = rolesList.length
  if (numPlayers === 0) return []

  // 池：role -> 剩余 count
  const pool: Record<string, number> = {}
  for (const r of rolesList) pool[r] = (pool[r] || 0) + 1

  const playerOrder = shuffle(
    Array.from({ length: numPlayers }, (_, i) => i + 1),
    rng,
  )

  const assignments: Record<number, string> = {}

  for (const playerId of playerOrder) {
    const preferred = new Set(preferences[playerId] || [])

    const weighted: string[] = []
    for (const [role, count] of Object.entries(pool)) {
      if (count <= 0) continue
      const w = preferred.has(role) ? count * 4 : count
      for (let i = 0; i < w; i++) weighted.push(role)
    }

    if (weighted.length === 0) {
      // 不该发生，兜底
      assignments[playerId] = rolesList[playerId - 1]
      continue
    }

    const picked = weightedPick(weighted, rng)
    assignments[playerId] = picked
    pool[picked] -= 1
    if (pool[picked] === 0) delete pool[picked]
  }

  const out: RoleList = []
  for (let i = 1; i <= numPlayers; i++) out.push(assignments[i])
  return out
}
