// Copyright 2025 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
struct Entity {
  value : UInt
} derive(Eq, Debug, Hash)

///|
enum EntityLifecycleState {
  Alive
  PendingSpawn
  PendingDespawn
  PendingRespawn
  Destroyed
} derive(Eq, Debug)

///|
priv struct EntityWorldStore {
  mut entity_generator : UInt
  mut alive_count : Int
  mut lifecycle_revision : UInt
  all_entities : @set.Set[Entity]
  lifecycle_states : Map[Entity, EntityLifecycleState]
  pending_state_changes : @set.Set[Entity]
}

///|
let entity_world_stores : Map[UInt, EntityWorldStore] = Map([])

///|
fn entity_store() -> EntityWorldStore {
  let world = @ecs.require_current_world()
  entity_world_stores.get_or_init(world.id(), () => {
    entity_generator: 0U,
    alive_count: 0,
    lifecycle_revision: 0U,
    all_entities: Set([]),
    lifecycle_states: Map([]),
    pending_state_changes: Set([]),
  })
}

///|
pub fn iter_entities() -> Iter[Entity] {
  entity_store().all_entities.iter()
}

///|
pub fn alive_entity_count() -> Int {
  entity_store().alive_count
}

///|
/// Returns the current world's monotonically increasing Entity lifecycle
/// revision.
///
/// Systems that retain a projection of all live Entities can compare this value
/// with their last synchronized revision and avoid scanning an unchanged World.
pub fn lifecycle_revision() -> UInt {
  entity_store().lifecycle_revision
}

///|
pub fn Entity::is_alive(e : Entity) -> Bool {
  match entity_store().lifecycle_states.get(e).unwrap_or(Destroyed) {
    Alive | PendingRespawn => true
    PendingSpawn | PendingDespawn | Destroyed => false
  }
}

///|
pub fn Entity::Entity() -> Entity {
  let store = entity_store()
  let entity = { value: store.entity_generator }
  store.entity_generator += 1U
  store.all_entities.add(entity)
  store.alive_count += 1
  store.lifecycle_revision += 1U
  store.lifecycle_states.set(entity, Alive)
  entity
}

///|
pub fn Entity::id(self : Entity) -> UInt {
  self.value
}

///|
/// Reserves an entity id that becomes alive on the next flush.
/// This matches deferred-spawn semantics used by commands.
pub fn Entity::reserve_spawn() -> Entity {
  let store = entity_store()
  let entity = { value: store.entity_generator }
  store.entity_generator += 1U
  store.lifecycle_states.set(entity, PendingSpawn)
  store.pending_state_changes.add(entity)
  entity
}

///|
pub fn Entity::destroy(e : Entity) -> Unit {
  let store = entity_store()
  match store.lifecycle_states.get(e).unwrap_or(Destroyed) {
    Destroyed | PendingDespawn => ()
    PendingSpawn => {
      store.lifecycle_states.set(e, PendingDespawn)
      store.pending_state_changes.add(e)
    }
    Alive | PendingRespawn => {
      if store.all_entities.contains(e) {
        if store.alive_count > 0 {
          store.alive_count -= 1
        }
      }
      store.lifecycle_revision += 1U
      store.lifecycle_states.set(e, PendingDespawn)
      store.pending_state_changes.add(e)
      if hierarchy_parents().get(e) is Some(p) {
        for c in p.children {
          c.destroy()
        }
      }
    }
  }
}

///|
/// Removes entries whose Entity keys are no longer alive in the current World.
///
/// Package-owned component stores can use `on_remove` to release resources
/// that are not owned by the Entity hierarchy.
pub fn[T] cleanup_dead_entities(
  storage : Map[Entity, T],
  on_remove? : (Entity, T) -> Unit = fn(_, _) { () },
) -> Unit {
  let dead : Array[Entity] = []
  for entity, _ in storage {
    if !entity.is_alive() {
      dead.push(entity)
    }
  }
  for entity in dead {
    match storage.get(entity) {
      Some(value) => {
        storage.remove(entity)
        on_remove(entity, value)
      }
      None => ()
    }
  }
}

///|
pub fn Entity::respawn(e : Entity) -> Unit {
  let store = entity_store()
  match store.lifecycle_states.get(e).unwrap_or(Destroyed) {
    Alive | PendingRespawn => ()
    Destroyed | PendingDespawn | PendingSpawn => {
      if store.lifecycle_states.get(e) is Some(PendingDespawn) &&
        store.all_entities.contains(e) {
        store.alive_count += 1
      }
      store.lifecycle_revision += 1U
      store.lifecycle_states.set(e, PendingRespawn)
      store.pending_state_changes.add(e)
      if hierarchy_parents().get(e) is Some(p) {
        for c in p.children {
          c.respawn()
        }
      }
    }
  }
}

///|
pub fn entity_flush_system(_delta : Double) -> Unit {
  let store = entity_store()
  for entity in store.pending_state_changes.to_array() {
    match store.lifecycle_states.get(entity).unwrap_or(Destroyed) {
      PendingSpawn => {
        if !store.all_entities.contains(entity) {
          store.all_entities.add(entity)
          store.alive_count += 1
          store.lifecycle_revision += 1U
        }
        store.lifecycle_states.set(entity, Alive)
      }
      PendingDespawn => {
        if store.all_entities.contains(entity) {
          store.all_entities.remove(entity)
        }
        store.lifecycle_states.set(entity, Destroyed)
      }
      PendingRespawn => {
        if !store.all_entities.contains(entity) {
          store.all_entities.add(entity)
          store.alive_count += 1
          store.lifecycle_revision += 1U
        }
        store.lifecycle_states.set(entity, Alive)
      }
      Alive | Destroyed => ()
    }
    store.pending_state_changes.remove(entity)
  }
}