Chessistics/chessistics-engine/Model/BoardState.cs

using Chessistics.Engine.Events;

namespace Chessistics.Engine.Model;

public class BoardState
{
    public int Width { get; private set; }
    public int Height { get; private set; }
    public CellType[,] Grid { get; private set; }
    public Dictionary<Coords, ProductionDef> Productions { get; }
    public Dictionary<Coords, DemandState> Demands { get; }
    public List<PieceState> Pieces { get; }
    public List<PieceState> DestroyedPieces { get; } = new();
    public Dictionary<Coords, int> ProductionBuffers { get; }
    public Dictionary<Coords, TransformerDef> Transformers { get; }
    public Dictionary<Coords, int> TransformerInputBuffers { get; }
    public Dictionary<Coords, int> TransformerOutputBuffers { get; }
    public SimPhase Phase { get; set; }
    public int TurnNumber { get; set; }
    public int NextPieceId { get; set; }
    public Dictionary<PieceKind, int> RemainingStock { get; }

    // Campaign state (null for legacy level mode)
    public CampaignState? Campaign { get; private set; }

    // Tracks all cells ever occupied by a piece (for metrics)
    public HashSet<Coords> OccupiedCells { get; }

    private readonly LevelDef? _levelDef;
    private bool _isApplyingCommand;

    private BoardState(int width, int height)
    {
        Width = width;
        Height = height;

        Grid = new CellType[Width, Height];
        Productions = new Dictionary<Coords, ProductionDef>();
        Demands = new Dictionary<Coords, DemandState>();
        Pieces = new List<PieceState>();
        ProductionBuffers = new Dictionary<Coords, int>();
        Transformers = new Dictionary<Coords, TransformerDef>();
        TransformerInputBuffers = new Dictionary<Coords, int>();
        TransformerOutputBuffers = new Dictionary<Coords, int>();
        RemainingStock = new Dictionary<PieceKind, int>();
        OccupiedCells = new HashSet<Coords>();

        Phase = SimPhase.Paused;
        TurnNumber = 0;
        NextPieceId = 1;

        for (int c = 0; c < Width; c++)
            for (int r = 0; r < Height; r++)
                Grid[c, r] = CellType.Empty;
    }

    private BoardState(LevelDef level) : this(level.Width, level.Height)
    {
        _levelDef = level;
        ApplyLevelDef(level);
    }

    public static BoardState FromLevel(LevelDef level) => new(level);

    public static BoardState FromCampaign(CampaignDef campaignDef)
    {
        var state = new BoardState(campaignDef.InitialWidth, campaignDef.InitialHeight);
        state.Campaign = new CampaignState(campaignDef);
        return state;
    }

    public CellType GetCell(Coords coords) => Grid[coords.Col, coords.Row];

    public bool IsOnBoard(Coords coords) => coords.IsOnBoard(Width, Height);

    /// <summary>
    /// Returns all cells currently occupied by any piece.
    /// In campaign mode (no Edit phase), always uses CurrentCell.
    /// </summary>
    public HashSet<Coords> GetOccupiedCells()
    {
        var occupied = new HashSet<Coords>();
        foreach (var piece in Pieces)
        {
            occupied.Add(piece.CurrentCell);
        }
        return occupied;
    }

    public PieceState? GetPieceById(int pieceId) => Pieces.Find(p => p.Id == pieceId);

    public void ApplyCommand(Action<BoardState, List<IWorldEvent>> apply, List<IWorldEvent> changeList)
    {
        if (_isApplyingCommand)
            throw new InvalidOperationException("Nested command not allowed.");

        _isApplyingCommand = true;
        try
        {
            apply(this, changeList);
        }
        finally
        {
            _isApplyingCommand = false;
        }
    }

    /// <summary>
    /// Expand the board to new dimensions, preserving existing state.
    /// </summary>
    public void Resize(int newWidth, int newHeight)
    {
        if (newWidth < Width || newHeight < Height)
            throw new InvalidOperationException("Cannot shrink the board.");

        if (newWidth == Width && newHeight == Height)
            return;

        var newGrid = new CellType[newWidth, newHeight];
        for (int c = 0; c < newWidth; c++)
            for (int r = 0; r < newHeight; r++)
                newGrid[c, r] = CellType.Empty;

        // Copy existing grid
        for (int c = 0; c < Width; c++)
            for (int r = 0; r < Height; r++)
                newGrid[c, r] = Grid[c, r];

        Grid = newGrid;
        Width = newWidth;
        Height = newHeight;
    }

    /// <summary>
    /// Apply a terrain patch (new cells, productions, demands, walls).
    /// </summary>
    public void ApplyTerrainPatch(TerrainPatch patch, int missionIndex)
    {
        if (patch.NewWidth > Width || patch.NewHeight > Height)
            Resize(patch.NewWidth, patch.NewHeight);

        foreach (var cell in patch.Cells)
        {
            var coords = new Coords(cell.Col, cell.Row);

            // Always clear existing buildings before overwriting
            ClearBuildingAt(coords);
            Grid[cell.Col, cell.Row] = cell.Type;

            switch (cell.Type)
            {
                case CellType.Production when cell.Production != null:
                    Productions[coords] = cell.Production;
                    ProductionBuffers[coords] = 0;
                    break;
                case CellType.Demand when cell.Demand != null:
                    Demands[coords] = new DemandState(cell.Demand, missionIndex);
                    break;
                case CellType.Transformer when cell.Transformer != null:
                    Transformers[coords] = cell.Transformer;
                    TransformerInputBuffers[coords] = 0;
                    TransformerOutputBuffers[coords] = 0;
                    break;
                case CellType.Wall:
                    // Remove pieces whose start or end cell is on this wall
                    RemovePiecesOnCell(coords);
                    break;
            }
        }
    }

    /// <summary>
    /// Add stock to the remaining stock (cumulative for campaigns).
    /// </summary>
    public void AddStock(IReadOnlyList<PieceStock> stock)
    {
        foreach (var s in stock)
            RemainingStock[s.Kind] = RemainingStock.GetValueOrDefault(s.Kind) + s.Count;
    }

    public void ResetFromLevel()
    {
        if (_levelDef == null)
            throw new InvalidOperationException("Cannot reset: no level definition.");

        Pieces.Clear();
        DestroyedPieces.Clear();
        Productions.Clear();
        Demands.Clear();
        ProductionBuffers.Clear();
        Transformers.Clear();
        TransformerInputBuffers.Clear();
        TransformerOutputBuffers.Clear();
        RemainingStock.Clear();
        OccupiedCells.Clear();

        Phase = SimPhase.Paused;
        TurnNumber = 0;
        NextPieceId = 1;

        for (int c = 0; c < Width; c++)
            for (int r = 0; r < Height; r++)
                Grid[c, r] = CellType.Empty;

        ApplyLevelDef(_levelDef);
    }

    /// <summary>
    /// Remove pieces whose StartCell or EndCell is on the given cell (return to stock).
    /// Used when a wall overwrites an occupied cell during terrain patching.
    /// </summary>
    private void RemovePiecesOnCell(Coords coords)
    {
        var toRemove = Pieces
            .Where(p => p.StartCell == coords || p.EndCell == coords)
            .ToList();

        foreach (var piece in toRemove)
        {
            Pieces.Remove(piece);
            RemainingStock[piece.Kind] = RemainingStock.GetValueOrDefault(piece.Kind) + 1;
        }
    }

    /// <summary>
    /// Capture a deep copy of every mutable field (for QuickSave).
    /// Immutable defs and CampaignDef are shared by reference.
    /// </summary>
    public WorldSave CaptureSave()
    {
        var grid = new CellType[Width, Height];
        Array.Copy(Grid, grid, Grid.Length);

        return new WorldSave
        {
            Width = Width,
            Height = Height,
            Grid = grid,
            Phase = Phase,
            TurnNumber = TurnNumber,
            NextPieceId = NextPieceId,
            Productions = new Dictionary<Coords, ProductionDef>(Productions),
            ProductionBuffers = new Dictionary<Coords, int>(ProductionBuffers),
            Demands = Demands.ToDictionary(kv => kv.Key, kv => kv.Value.Clone()),
            Transformers = new Dictionary<Coords, TransformerDef>(Transformers),
            TransformerInputBuffers = new Dictionary<Coords, int>(TransformerInputBuffers),
            TransformerOutputBuffers = new Dictionary<Coords, int>(TransformerOutputBuffers),
            Pieces = Pieces.Select(p => p.Clone()).ToList(),
            DestroyedPieces = DestroyedPieces.Select(p => p.Clone()).ToList(),
            RemainingStock = new Dictionary<PieceKind, int>(RemainingStock),
            OccupiedCells = new HashSet<Coords>(OccupiedCells),
            Campaign = Campaign == null ? null : new CampaignSaveData
            {
                CurrentMissionIndex = Campaign.CurrentMissionIndex,
                CompletedMissions = new List<int>(Campaign.CompletedMissions),
                AvailablePieceKinds = new HashSet<PieceKind>(Campaign.AvailablePieceKinds),
                AvailableLevels = new HashSet<PieceUpgrade>(Campaign.AvailableLevels)
            }
        };
    }

    /// <summary>
    /// Restore every mutable field from a save. Board dimensions can grow
    /// or shrink; the grid is fully replaced.
    /// </summary>
    public void RestoreFromSave(WorldSave save)
    {
        Width = save.Width;
        Height = save.Height;
        Grid = new CellType[Width, Height];
        Array.Copy(save.Grid, Grid, save.Grid.Length);

        Phase = save.Phase;
        TurnNumber = save.TurnNumber;
        NextPieceId = save.NextPieceId;

        Productions.Clear();
        foreach (var kv in save.Productions) Productions[kv.Key] = kv.Value;

        ProductionBuffers.Clear();
        foreach (var kv in save.ProductionBuffers) ProductionBuffers[kv.Key] = kv.Value;

        Demands.Clear();
        foreach (var kv in save.Demands) Demands[kv.Key] = kv.Value.Clone();

        Transformers.Clear();
        foreach (var kv in save.Transformers) Transformers[kv.Key] = kv.Value;

        TransformerInputBuffers.Clear();
        foreach (var kv in save.TransformerInputBuffers) TransformerInputBuffers[kv.Key] = kv.Value;

        TransformerOutputBuffers.Clear();
        foreach (var kv in save.TransformerOutputBuffers) TransformerOutputBuffers[kv.Key] = kv.Value;

        Pieces.Clear();
        foreach (var p in save.Pieces) Pieces.Add(p.Clone());

        DestroyedPieces.Clear();
        foreach (var p in save.DestroyedPieces) DestroyedPieces.Add(p.Clone());

        RemainingStock.Clear();
        foreach (var kv in save.RemainingStock) RemainingStock[kv.Key] = kv.Value;

        OccupiedCells.Clear();
        foreach (var c in save.OccupiedCells) OccupiedCells.Add(c);

        if (save.Campaign != null && Campaign != null)
        {
            Campaign.CurrentMissionIndex = save.Campaign.CurrentMissionIndex;
            Campaign.CompletedMissions.Clear();
            Campaign.CompletedMissions.AddRange(save.Campaign.CompletedMissions);
            Campaign.AvailablePieceKinds.Clear();
            foreach (var k in save.Campaign.AvailablePieceKinds) Campaign.AvailablePieceKinds.Add(k);
            Campaign.AvailableLevels.Clear();
            foreach (var u in save.Campaign.AvailableLevels) Campaign.AvailableLevels.Add(u);
        }
    }

    private void ClearBuildingAt(Coords coords)
    {
        Productions.Remove(coords);
        ProductionBuffers.Remove(coords);
        Demands.Remove(coords);
        Transformers.Remove(coords);
        TransformerInputBuffers.Remove(coords);
        TransformerOutputBuffers.Remove(coords);
    }

    private void ApplyLevelDef(LevelDef level)
    {
        foreach (var wall in level.Walls)
            Grid[wall.Col, wall.Row] = CellType.Wall;

        foreach (var prod in level.Productions)
        {
            Grid[prod.Position.Col, prod.Position.Row] = CellType.Production;
            Productions[prod.Position] = prod;
            ProductionBuffers[prod.Position] = 0;
        }

        foreach (var demand in level.Demands)
        {
            Grid[demand.Position.Col, demand.Position.Row] = CellType.Demand;
            Demands[demand.Position] = new DemandState(demand);
        }

        foreach (var stock in level.Stock)
            RemainingStock[stock.Kind] = stock.Count;
    }
}