use bevy::{ecs::system::SystemState, prelude::*};
use serde::{Deserialize, Serialize};

use crate::{
    actions::Action,
    animals::{Animal, AnimalAssets, AnimalKind, Tame, spawn_animal},
    building::{Blueprints, BuildKind},
    colonist::{Colonist, ColonistAssets, NAMES, spawn_colonist},
    designations::Designations,
    economy::ColonyGoals,
    family::{Affinity, CHILD_SCALE, Child, Family, Homes, Partner},
    food::Hunger,
    items::{Carrying, Item, ItemAssets, spawn_stack},
    jobs::{Tool, WorkPriorities},
    mood::{Age, LastMeal, Losses, Mood, Sick, Unhappy},
    needs::Rest,
    population::{ColonyLog, Starvation},
    rng::SimRng,
    sim::{SimPos, SimTick, sim_pos},
    skills::Skills,
    structures::{Fuel, Structure, StructureAssets, StructureKind, spawn_structure},
    threats::Health,
    world::{MapLayers, TilePos, WorldMap},
    worldgen::{WorldGenerated, WorldSeed, generate},
    zones::{Zones, ZonesSave},
};

pub struct SavePlugin;

impl Plugin for SavePlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Update, save_game)
            .add_systems(PreUpdate, start_load)
            .add_systems(PostUpdate, finish_load);
    }
}

const SAVE_VERSION: u32 = 1;

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct ColonistSave {
    name: String,
    pos: Vec2,
    tile: UVec2,
    rest: f32,
    hunger: f32,
    mood: Option<f32>,
    born: Option<i64>,
    family: Option<u32>,
    partner: Option<usize>,
    child_born: Option<u64>,
    carrying: Option<Item>,
    priorities: WorkPriorities,
    tool_wear: Option<u32>,
    sick_until: Option<u64>,
    last_meal: Option<(bool, u64)>,
    unhappy: u32,
    starvation: u32,
    #[serde(default)]
    skills: Skills,
    #[serde(default = "full_health")]
    health: f32,
}

fn full_health() -> f32 {
    1.0
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct AnimalSave {
    kind: AnimalKind,
    pos: Vec2,
    tile: UVec2,
    tame: bool,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct SaveGame {
    version: u32,
    seed: u32,
    tick: u64,
    rng: SimRng,
    map: MapLayers,
    zones: ZonesSave,
    blueprints: Vec<(UVec2, BuildKind)>,
    homes: Homes,
    designations: Designations,
    goals: ColonyGoals,
    log: ColonyLog,
    losses: Losses,
    colonists: Vec<ColonistSave>,
    items: Vec<(UVec2, Item)>,
    structures: Vec<(UVec2, StructureKind, Option<u32>)>,
    animals: Vec<AnimalSave>,
}

#[allow(clippy::type_complexity)]
pub fn snapshot(world: &mut World) -> SaveGame {
    let mut colonist_query = world.query::<(
        Entity,
        &Colonist,
        &SimPos,
        &TilePos,
        &Rest,
        &Hunger,
        &Carrying,
        &WorkPriorities,
        (
            Option<&Mood>,
            Option<&Age>,
            Option<&Family>,
            Option<&Partner>,
            Option<&Child>,
            Option<&Tool>,
            Option<&Sick>,
            Option<&LastMeal>,
            Option<&Unhappy>,
            Option<&Starvation>,
            Option<&Skills>,
            Option<&Health>,
        ),
    )>();
    let rows: Vec<_> = colonist_query.iter(world).collect();
    let order: Vec<Entity> = rows.iter().map(|r| r.0).collect();
    let colonists = rows
        .iter()
        .map(
            |&(_, colonist, pos, tile, rest, hunger, carrying, priorities, extra)| {
                let (
                    mood,
                    age,
                    family,
                    partner,
                    child,
                    tool,
                    sick,
                    meal,
                    unhappy,
                    starvation,
                    skills,
                    health,
                ) = extra;
                ColonistSave {
                    name: colonist.name.clone(),
                    pos: pos.0,
                    tile: tile.0,
                    rest: rest.0,
                    hunger: hunger.0,
                    mood: mood.map(|m| m.0),
                    born: age.map(|a| a.born),
                    family: family.map(|f| f.0),
                    partner: partner.and_then(|p| order.iter().position(|&e| e == p.0)),
                    child_born: child.map(|c| c.born),
                    carrying: carrying.0,
                    priorities: priorities.clone(),
                    tool_wear: tool.map(|t| t.wear),
                    sick_until: sick.map(|s| s.until),
                    last_meal: meal.map(|m| (m.cooked, m.tick)),
                    unhappy: unhappy.map_or(0, |u| u.0),
                    starvation: starvation.map_or(0, |s| s.0),
                    skills: skills.cloned().unwrap_or_default(),
                    health: health.map_or(1.0, |h| h.0),
                }
            },
        )
        .collect();
    let items = world
        .query::<(&Item, &TilePos)>()
        .iter(world)
        .map(|(item, pos)| (pos.0, *item))
        .collect();
    let structures = world
        .query::<(&Structure, &TilePos, Option<&Fuel>)>()
        .iter(world)
        .map(|(s, pos, fuel)| (pos.0, s.kind, fuel.map(|f| f.0)))
        .collect();
    let animals = world
        .query::<(&Animal, &SimPos, &TilePos, Has<Tame>)>()
        .iter(world)
        .map(|(animal, pos, tile, tame)| AnimalSave {
            kind: animal.kind,
            pos: pos.0,
            tile: tile.0,
            tame,
        })
        .collect();
    let mut blueprints: Vec<(UVec2, BuildKind)> = world
        .resource::<Blueprints>()
        .0
        .iter()
        .map(|(t, k)| (*t, *k))
        .collect();
    blueprints.sort_by_key(|(t, _)| (t.y, t.x));
    SaveGame {
        version: SAVE_VERSION,
        seed: world.resource::<WorldSeed>().0,
        tick: world.resource::<SimTick>().0,
        rng: world.resource::<SimRng>().clone(),
        map: world.resource::<WorldMap>().layers(),
        zones: world.resource::<Zones>().save(),
        blueprints,
        homes: world.resource::<Homes>().clone(),
        designations: Designations(world.resource::<Designations>().0.clone()),
        goals: world.resource::<ColonyGoals>().clone(),
        log: ColonyLog(world.resource::<ColonyLog>().0.clone()),
        losses: Losses(world.resource::<Losses>().0.clone()),
        colonists,
        items,
        structures,
        animals,
    }
}

pub fn restore_map(save: &SaveGame, size: UVec2) -> Result<WorldMap, String> {
    if save.version != SAVE_VERSION {
        return Err(format!("unsupported save version {}", save.version));
    }
    let mut map = generate(size, save.seed);
    map.restore(&save.map)?;
    Ok(map)
}

#[derive(Resource)]
struct PendingLoad(SaveGame);

fn despawn_all<T: Component>(world: &mut World) {
    let entities: Vec<Entity> = world
        .query_filtered::<Entity, With<T>>()
        .iter(world)
        .collect();
    for entity in entities {
        world.despawn(entity);
    }
}

type SpawnParams<'w, 's> = (
    Commands<'w, 's>,
    Res<'w, ColonistAssets>,
    Res<'w, ItemAssets>,
    Res<'w, StructureAssets>,
    Res<'w, AnimalAssets>,
    ResMut<'w, Assets<ColorMaterial>>,
);

pub fn apply(world: &mut World, save: SaveGame) {
    despawn_all::<Colonist>(world);
    despawn_all::<Animal>(world);
    despawn_all::<Item>(world);
    despawn_all::<Structure>(world);
    world.insert_resource(SimTick(save.tick));
    world.insert_resource(save.rng);
    world.insert_resource(Zones::load(save.zones));
    world.insert_resource(Blueprints(save.blueprints.into_iter().collect()));
    world.insert_resource(save.homes);
    world.insert_resource(save.designations);
    world.insert_resource(save.goals);
    world.insert_resource(save.log);
    world.insert_resource(save.losses);
    world.insert_resource(Affinity::default());
    let mut state: SystemState<SpawnParams> = SystemState::new(world);
    let (
        mut commands,
        colonist_assets,
        item_assets,
        structure_assets,
        animal_assets,
        mut materials,
    ) = state.get_mut(world).expect("assets exist");
    for (tile, item) in save.items {
        spawn_stack(&mut commands, &item_assets, tile, item);
    }
    for (tile, kind, fuel) in save.structures {
        let entity = spawn_structure(&mut commands, &structure_assets, tile, kind);
        if let Some(fuel) = fuel {
            commands.entity(entity).insert(Fuel(fuel));
        }
    }
    for animal in save.animals {
        let entity = spawn_animal(&mut commands, &animal_assets, animal.kind, animal.tile);
        let mut entity = commands.entity(entity);
        entity.insert(sim_pos(animal.pos));
        if animal.tame {
            entity.insert(Tame);
        }
    }
    let entities: Vec<Entity> = save
        .colonists
        .iter()
        .map(|c| {
            let index = NAMES.iter().position(|n| *n == c.name).unwrap_or(0);
            let entity = spawn_colonist(
                &mut commands,
                &colonist_assets,
                &item_assets,
                &mut materials,
                index,
                c.tile,
            );
            let mut e = commands.entity(entity);
            e.insert((
                Colonist {
                    name: c.name.clone(),
                },
                Name::new(c.name.clone()),
                sim_pos(c.pos),
                Rest(c.rest),
                Hunger(c.hunger),
                Carrying(c.carrying),
                c.priorities.clone(),
                Unhappy(c.unhappy),
                Starvation(c.starvation),
                c.skills.clone(),
                Health(c.health),
            ));
            if let Some(mood) = c.mood {
                e.insert(Mood(mood));
            }
            if let Some(born) = c.born {
                e.insert(Age { born });
            }
            if let Some(family) = c.family {
                e.insert(Family(family));
            }
            if let Some(born) = c.child_born {
                e.insert((
                    Child { born },
                    Transform::from_translation(c.pos.extend(1.0))
                        .with_scale(Vec3::splat(CHILD_SCALE)),
                ));
            }
            if let Some(wear) = c.tool_wear {
                e.insert(Tool { wear });
            }
            if let Some(until) = c.sick_until {
                e.insert(Sick { until });
            }
            if let Some((cooked, tick)) = c.last_meal {
                e.insert(LastMeal { cooked, tick });
            }
            entity
        })
        .collect();
    for (c, &entity) in save.colonists.iter().zip(&entities) {
        if let Some(&partner) = c.partner.and_then(|i| entities.get(i)) {
            commands.entity(entity).insert(Partner(partner));
        }
    }
    state.apply(world);
}

const SAVE_KEY: &str = "simulator-save";

#[cfg(target_arch = "wasm32")]
fn storage() -> Result<web_sys::Storage, String> {
    web_sys::window()
        .and_then(|w| w.local_storage().ok().flatten())
        .ok_or_else(|| "no local storage".to_string())
}

#[cfg(target_arch = "wasm32")]
fn write_save(data: &str) -> Result<(), String> {
    storage()?
        .set_item(SAVE_KEY, data)
        .map_err(|_| "storage is full".to_string())
}

#[cfg(target_arch = "wasm32")]
fn read_save() -> Result<String, String> {
    storage()?
        .get_item(SAVE_KEY)
        .ok()
        .flatten()
        .ok_or_else(|| "no saved game".to_string())
}

#[cfg(not(target_arch = "wasm32"))]
fn save_path() -> std::path::PathBuf {
    std::path::PathBuf::from(format!("{SAVE_KEY}.json"))
}

#[cfg(not(target_arch = "wasm32"))]
fn write_save(data: &str) -> Result<(), String> {
    std::fs::write(save_path(), data).map_err(|e| e.to_string())
}

#[cfg(not(target_arch = "wasm32"))]
fn read_save() -> Result<String, String> {
    std::fs::read_to_string(save_path()).map_err(|_| "no saved game".to_string())
}

fn save_game(world: &mut World, mut reader: Local<Option<SystemState<MessageReader<Action>>>>) {
    let state = reader.get_or_insert_with(|| SystemState::new(world));
    let requested = state
        .get_mut(world)
        .is_ok_and(|mut actions| actions.read().any(|a| *a == Action::Save));
    if !requested {
        return;
    }
    let save = snapshot(world);
    let result = serde_json::to_string(&save)
        .map_err(|e| e.to_string())
        .and_then(|data| write_save(&data));
    let tick = *world.resource::<SimTick>();
    let message = match result {
        Ok(()) => "Game saved".to_string(),
        Err(e) => format!("Save failed: {e}"),
    };
    world.resource_mut::<ColonyLog>().push(tick, message);
}

fn start_load(world: &mut World, mut reader: Local<Option<SystemState<MessageReader<Action>>>>) {
    let state = reader.get_or_insert_with(|| SystemState::new(world));
    let requested = state
        .get_mut(world)
        .is_ok_and(|mut actions| actions.read().any(|a| *a == Action::Load));
    if !requested {
        return;
    }
    let size = world.resource::<WorldMap>().size;
    let loaded = read_save().and_then(|data| {
        let save: SaveGame = serde_json::from_str(&data).map_err(|e| e.to_string())?;
        let map = restore_map(&save, size)?;
        Ok((save, map))
    });
    match loaded {
        Ok((save, map)) => {
            world.insert_resource(WorldSeed(save.seed));
            world.insert_resource(map);
            world.write_message(WorldGenerated);
            world.insert_resource(PendingLoad(save));
        }
        Err(e) => {
            let tick = *world.resource::<SimTick>();
            world
                .resource_mut::<ColonyLog>()
                .push(tick, format!("Load failed: {e}"));
        }
    }
}

fn finish_load(world: &mut World) {
    let Some(PendingLoad(save)) = world.remove_resource::<PendingLoad>() else {
        return;
    };
    let tick = SimTick(save.tick);
    apply(world, save);
    world
        .resource_mut::<ColonyLog>()
        .push(tick, "Game loaded".to_string());
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{items::ItemKind, world::Plant};

    fn app() -> App {
        let mut app = App::new();
        app.insert_resource(WorldSeed(7))
            .insert_resource(generate(UVec2::splat(64), 7))
            .insert_resource(SimTick(12345))
            .insert_resource(SimRng::new(3))
            .init_resource::<Zones>()
            .init_resource::<Blueprints>()
            .init_resource::<Homes>()
            .init_resource::<Designations>()
            .init_resource::<ColonyGoals>()
            .init_resource::<ColonyLog>()
            .init_resource::<Losses>()
            .init_resource::<ColonistAssets>()
            .init_resource::<ItemAssets>()
            .init_resource::<StructureAssets>()
            .init_resource::<AnimalAssets>()
            .init_resource::<Assets<ColorMaterial>>();
        app
    }

    // Jobs, paths and reservations are not saved: loaded colonists start idle and choose again,
    // which every job already handles after an interruption. Everything else round-trips.
    #[test]
    fn a_saved_colony_loads_back_the_same() {
        let mut app = app();
        let world = app.world_mut();
        {
            let mut map = world.resource_mut::<WorldMap>();
            map.set_plant(UVec2::new(3, 3), Some(Plant::Crop));
            map.set_growth(UVec2::new(3, 3), 0.42);
            map.set_wall(UVec2::new(4, 4), true);
            map.set_roof(UVec2::new(5, 5), true);
            map.set_wear(UVec2::new(6, 6), 12.5);
        }
        world.resource_mut::<Blueprints>().0.insert(
            UVec2::new(9, 9),
            BuildKind::Structure(StructureKind::Statue),
        );
        let ada = world
            .spawn((
                Colonist { name: "Ada".into() },
                sim_pos(Vec2::new(40.0, 40.0)),
                TilePos(UVec2::new(2, 2)),
                Hunger(0.4),
                Family(1),
                Tool { wear: 99 },
                Carrying(Some(Item::new(ItemKind::Wood, 7))),
            ))
            .id();
        let bram = world
            .spawn((
                Colonist {
                    name: "Bram".into(),
                },
                Family(1),
                Partner(ada),
            ))
            .id();
        world.entity_mut(ada).insert(Partner(bram));
        world.spawn((Item::new(ItemKind::Meal, 3), TilePos(UVec2::new(1, 1))));
        world.spawn((
            Structure {
                kind: StructureKind::Campfire,
            },
            TilePos(UVec2::new(8, 8)),
            Fuel(500),
        ));
        let saved = snapshot(world);
        let json = serde_json::to_string(&saved).unwrap();
        assert!(json.len() < 100_000, "{}", json.len());

        let mut fresh = self::app();
        let world = fresh.world_mut();
        world.spawn(Colonist {
            name: "Stray".into(),
        });
        let loaded: SaveGame = serde_json::from_str(&json).unwrap();
        let map = restore_map(&loaded, UVec2::splat(64)).unwrap();
        world.insert_resource(map);
        apply(world, loaded);

        let map = world.resource::<WorldMap>();
        assert_eq!(map.growth(IVec2::new(3, 3)), 0.42);
        assert!(map.is_wall(IVec2::new(4, 4)) && map.is_roofed(IVec2::new(5, 5)));
        assert_eq!(map.wear(IVec2::new(6, 6)), 12.5);
        assert_eq!(map.worn_tiles(), vec![UVec2::new(6, 6)]);
        assert_eq!(world.resource::<SimTick>().0, 12345);
        assert_eq!(world.resource::<Blueprints>().0.len(), 1);
        let mut colonists = world.query::<(Entity, &Colonist, Option<&Partner>, &Hunger)>();
        let rows: Vec<_> = colonists.iter(world).collect();
        assert_eq!(rows.len(), 2);
        let find = |name: &str| rows.iter().find(|r| r.1.name == name).unwrap();
        let (ada, bram) = (find("Ada"), find("Bram"));
        assert_eq!(ada.2.map(|p| p.0), Some(bram.0));
        assert_eq!(ada.3.0, 0.4);
        assert_eq!(world.get::<Tool>(ada.0), Some(&Tool { wear: 99 }));
        assert_eq!(world.get::<Carrying>(ada.0).unwrap().0.unwrap().count, 7);
        let fuel: Vec<u32> = world.query::<&Fuel>().iter(world).map(|f| f.0).collect();
        assert_eq!(fuel, vec![500]);
        assert_eq!(world.query::<&Item>().iter(world).count(), 1);
    }
}