Imports System.Drawing.Drawing2D
Imports System.IO
Imports System.Text.Json
Imports System.Linq

Public Class RoomView
    Inherits Form

    Private Const INDENT As Integer = 20
    Private Const SQCOUNTX As Integer = 15
    Private Const SQCOUNTY As Integer = 11
    Private Const WALL_THICKNESS As Integer = 10

    Friend wallsHatchStyle As HatchStyle = HatchStyle.Cross
    Friend wallsColor1 As Color = Color.Gray
    Friend wallsColor2 As Color = Color.DarkGray
    Friend floorsHatchStyle As HatchStyle = HatchStyle.SmallCheckerBoard
    Friend floorsColor1 As Color = Color.LightGray
    Friend floorsColor2 As Color = Color.White

    Private SQSizeX As Integer
    Private SQSizeY As Integer

    Private isRightMouseDown As Boolean = False
    Private isLeftMouseDown As Boolean = False
    Private isMouseInsideGrid As Boolean = False

    Private highlightedX As Integer
    Private highlightedY As Integer
    Private highlightedWallX As Boolean = False
    Private highlightedWallY As Boolean = False
    Private highlightColor As Color = Color.Black
    Private highlightTimer As Timer
    Private isRestoringState As Boolean = False

    Friend verticalWalls(SQCOUNTX, SQCOUNTY - 1) As Boolean
    Friend horizontalWalls(SQCOUNTX - 1, SQCOUNTY) As Boolean

    Private rand As New Random()

    Private premadesFolder As String = "D:\01 - Grags\Visual Studio\Enigma\Data\Premades"
    Private savesFolder As String = "D:\01 - Grags\Visual Studio\Enigma\Data\Saves"

    Private undoStack As New Stack(Of RoomState)
    Private redoStack As New Stack(Of RoomState)

    Private currentFilePath As String = ""

    Friend placedItems(SQCOUNTX - 1, SQCOUNTY - 1) As PlacedItem

    Private isActionInProgress As Boolean = False

    Public Sub New()
        InitializeComponent()
        InitializeWalls()
        highlightTimer = New Timer() With {.Interval = 500}
        AddHandler highlightTimer.Tick, AddressOf HighlightTimer_Tick
        highlightTimer.Start()
        AddColorPictureBoxHandlers(Me)
        AddHandler picBlock1.MouseDown, AddressOf PicBlock_MouseDown
        AddHandler picBaddy1.MouseDown, AddressOf PicBlock_MouseDown
        PicGame.AllowDrop = True
        AddHandler PicGame.DragEnter, AddressOf PicGame_DragEnter
        AddHandler PicGame.DragOver, AddressOf PicGame_DragOver
        AddHandler PicGame.DragDrop, AddressOf PicGame_DragDrop
        If Not isRestoringState Then SaveStateForUndo()
    End Sub

    Private Sub InitializeWalls()
        For x As Integer = 0 To SQCOUNTX
            For y As Integer = 0 To SQCOUNTY - 1
                verticalWalls(x, y) = True
            Next
        Next
        For x As Integer = 0 To SQCOUNTX - 1
            For y As Integer = 0 To SQCOUNTY
                horizontalWalls(x, y) = True
            Next
        Next
        For y As Integer = 0 To SQCOUNTY - 1
            verticalWalls(0, y) = True
            verticalWalls(SQCOUNTX, y) = True
        Next
        For x As Integer = 0 To SQCOUNTX - 1
            horizontalWalls(x, 0) = True
            horizontalWalls(x, SQCOUNTY) = True
        Next
    End Sub

    Private Sub RoomView_Load(sender As Object, e As EventArgs) Handles MyBase.Load
        LoadPremades()
    End Sub

    Private Sub LoadPremades()
        If Not Directory.Exists(premadesFolder) Then Directory.CreateDirectory(premadesFolder)
        lstPremades.Items.Clear()
        For Each file As String In Directory.GetFiles(premadesFolder, "*.Premade")
            lstPremades.Items.Add(Path.GetFileNameWithoutExtension(file))
        Next
    End Sub

    Private Sub lstPremades_SelectedIndexChanged(sender As Object, e As EventArgs) Handles lstPremades.SelectedIndexChanged
        If lstPremades.SelectedIndex >= 0 Then
            If Not isRestoringState Then SaveStateForUndo()
            Dim selectedPremade As String = lstPremades.SelectedItem.ToString()
            LoadSettingsFromFile(Path.Combine(premadesFolder, selectedPremade & ".Premade"))
            PicGame.Invalidate()
        End If
    End Sub

    Private Sub LoadSettingsFromFile(filePath As String)
        Try
            If Not File.Exists(filePath) Then
                MessageBox.Show("File not found: " & filePath)
                Return
            End If
            Dim lines As String() = File.ReadAllLines(filePath)
            Dim section As String = ""
            For Each line As String In lines
                line = line.Trim()
                If String.IsNullOrEmpty(line) OrElse line.StartsWith(";") Then Continue For
                If line.StartsWith("[") AndAlso line.EndsWith("]") Then
                    section = line.Substring(1, line.Length - 2)
                Else
                    Dim parts() As String = line.Split(New Char() {":"c}, 2)
                    If parts.Length = 2 Then
                        Dim key As String = parts(0).Trim()
                        Dim value As String = parts(1).Trim()
                        Select Case section
                            Case "Walls"
                                LoadProperty(key, value, True)
                            Case "Floors"
                                LoadProperty(key, value, False)
                        End Select
                    End If
                End If
            Next
        Catch ex As Exception
            MessageBox.Show("Error loading settings: " & ex.Message)
        End Try
    End Sub

    Private Sub LoadProperty(key As String, value As String, isWall As Boolean)
        Select Case key
            Case "HatchStyle"
                Dim hatchStyle As HatchStyle
                If [Enum].TryParse(value, hatchStyle) Then
                    If isWall Then wallsHatchStyle = hatchStyle Else floorsHatchStyle = hatchStyle
                End If
            Case "Color1"
                Dim color As Color = ParseColor(value)
                If isWall Then wallsColor1 = color Else floorsColor1 = color
            Case "Color2"
                Dim color As Color = ParseColor(value)
                If isWall Then wallsColor2 = color Else floorsColor2 = color
        End Select
    End Sub

    Private Function ParseColor(colorString As String) As Color
        Dim parts() As String = colorString.Split(","c)
        If parts.Length = 4 Then
            Dim a, r, g, b As Integer
            If Integer.TryParse(parts(0), a) AndAlso Integer.TryParse(parts(1), r) AndAlso
               Integer.TryParse(parts(2), g) AndAlso Integer.TryParse(parts(3), b) Then
                Return Color.FromArgb(a, r, g, b)
            End If
        End If
        Return Color.Black
    End Function

    Private Sub btnSavePremade_Click(sender As Object, e As EventArgs) Handles btnSavePremade.Click
        Dim saveDialog As New SaveFileDialog() With {
            .InitialDirectory = premadesFolder,
            .Filter = "Premade Files (*.Premade)|*.Premade",
            .DefaultExt = "Premade",
            .AddExtension = True,
            .Title = "Save Premade"
        }
        If saveDialog.ShowDialog() = DialogResult.OK Then
            SaveSettingsToFile(saveDialog.FileName)
            LoadPremades()
        End If
    End Sub

    Private Sub SaveSettingsToFile(filePath As String)
        Try
            Using sw As New StreamWriter(filePath)
                sw.WriteLine("[Walls]")
                sw.WriteLine("HatchStyle: " & wallsHatchStyle.ToString())
                sw.WriteLine("Color1: " & ColorToString(wallsColor1))
                sw.WriteLine("Color2: " & ColorToString(wallsColor2))
                sw.WriteLine()
                sw.WriteLine("[Floors]")
                sw.WriteLine("HatchStyle: " & floorsHatchStyle.ToString())
                sw.WriteLine("Color1: " & ColorToString(floorsColor1))
                sw.WriteLine("Color2: " & ColorToString(floorsColor2))
            End Using
        Catch ex As Exception
            MessageBox.Show("Error saving settings: " & ex.Message)
        End Try
    End Sub

    Private Function ColorToString(color As Color) As String
        Return $"{color.A},{color.R},{color.G},{color.B}"
    End Function

    Private Sub mnuNew_Click(sender As Object, e As EventArgs) Handles mnuNew.Click
        If Not isRestoringState Then SaveStateForUndo()
        wallsHatchStyle = HatchStyle.Cross
        wallsColor1 = Color.Gray
        wallsColor2 = Color.DarkGray
        floorsHatchStyle = HatchStyle.SmallCheckerBoard
        floorsColor1 = Color.LightGray
        floorsColor2 = Color.White
        InitializeWalls()
        placedItems = New PlacedItem(SQCOUNTX - 1, SQCOUNTY - 1) {}
        currentFilePath = ""
        PicGame.Invalidate()
    End Sub

    Private Sub mnuSave_Click(sender As Object, e As EventArgs) Handles mnuSave.Click
        If String.IsNullOrEmpty(currentFilePath) Then
            mnuSaveAs_Click(sender, e)
        Else
            SaveRoomToFile(currentFilePath)
        End If
    End Sub

    Private Sub mnuLoad_Click(sender As Object, e As EventArgs) Handles mnuLoad.Click
        Dim openDialog As New OpenFileDialog() With {
            .InitialDirectory = savesFolder,
            .Filter = "Room Files (*.room)|*.room",
            .DefaultExt = "room",
            .AddExtension = True,
            .Title = "Load Room"
        }
        If openDialog.ShowDialog() = DialogResult.OK Then
            If Not isRestoringState Then SaveStateForUndo()
            LoadRoomFromFile(openDialog.FileName)
            currentFilePath = openDialog.FileName
            PicGame.Invalidate()
        End If
    End Sub

    Private Sub mnuSaveAs_Click(sender As Object, e As EventArgs) Handles mnuSaveAs.Click
        Dim saveDialog As New SaveFileDialog() With {
            .InitialDirectory = savesFolder,
            .Filter = "Room Files (*.room)|*.room",
            .DefaultExt = "room",
            .AddExtension = True,
            .Title = "Save Room As"
        }
        If saveDialog.ShowDialog() = DialogResult.OK Then
            SaveRoomToFile(saveDialog.FileName)
            currentFilePath = saveDialog.FileName
        End If
    End Sub

    Private Sub SaveRoomToFile(filePath As String)
        Try
            Dim roomState As New RoomState(Me)
            Dim options As New JsonSerializerOptions With {
                .WriteIndented = True,
                .IncludeFields = True
            }
            File.WriteAllText(filePath, JsonSerializer.Serialize(roomState, options))
        Catch ex As Exception
            MessageBox.Show("Error saving room: " & ex.Message)
        End Try
    End Sub

    Private Sub LoadRoomFromFile(filePath As String)
        Try
            If Not File.Exists(filePath) Then
                MessageBox.Show("File not found: " & filePath)
                Return
            End If
            Dim jsonString As String = File.ReadAllText(filePath)
            Dim options As New JsonSerializerOptions With {.IncludeFields = True}
            Dim roomState As RoomState = JsonSerializer.Deserialize(Of RoomState)(jsonString, options)
            roomState.RestoreState(Me)
            PicGame.Invalidate()
        Catch ex As Exception
            MessageBox.Show("Error loading room: " & ex.Message)
        End Try
    End Sub

    Private Sub btnUndo_Click(sender As Object, e As EventArgs) Handles btnUndo.Click
        If undoStack.Count > 1 Then
            SaveStateForRedo()
            undoStack.Pop()
            isRestoringState = True
            Dim previousState As RoomState = undoStack.Peek()
            previousState.RestoreState(Me)
            isRestoringState = False
            PicGame.Invalidate()
        End If
    End Sub

    Private Sub btnRedo_Click(sender As Object, e As EventArgs) Handles btnRedo.Click
        If redoStack.Count > 0 Then
            SaveStateForUndo()
            isRestoringState = True
            Dim redoState As RoomState = redoStack.Pop()
            redoState.RestoreState(Me)
            isRestoringState = False
            PicGame.Invalidate()
        End If
    End Sub

    Private Sub SaveStateForUndo()
        If isRestoringState Then Return
        undoStack.Push(New RoomState(Me))
        If undoStack.Count > 50 Then
            undoStack = New Stack(Of RoomState)(undoStack.Reverse().Take(50))
        End If
    End Sub

    Private Sub SaveStateForRedo()
        If isRestoringState Then Return
        redoStack.Push(New RoomState(Me))
    End Sub

    Private Sub AddColorPictureBoxHandlers(parent As Control)
        For Each ctrl As Control In parent.Controls
            If TypeOf ctrl Is PictureBox AndAlso ctrl.Name.StartsWith("col") Then
                AddHandler ctrl.Click, AddressOf ColorPictureBox_Click
            End If
            If ctrl.HasChildren Then
                AddColorPictureBoxHandlers(ctrl)
            End If
        Next
    End Sub

    Private Sub ColorPictureBox_Click(sender As Object, e As EventArgs)
        If Not isRestoringState Then SaveStateForUndo()
        Dim pb As PictureBox = CType(sender, PictureBox)
        Dim selectedColor As Color = pb.BackColor
        If rdoFloors.Checked Then
            If pb.Name.StartsWith("col") Then
                Dim numStr As String = pb.Name.Substring(3)
                Dim num As Integer
                If Integer.TryParse(numStr, num) Then
                    If num >= 1 AndAlso num <= 48 Then
                        floorsColor1 = selectedColor
                    ElseIf num >= 51 AndAlso num <= 98 Then
                        floorsColor2 = selectedColor
                    End If
                End If
            End If
        ElseIf rdoWalls.Checked Then
            If pb.Name.StartsWith("col") Then
                Dim numStr As String = pb.Name.Substring(3)
                Dim num As Integer
                If Integer.TryParse(numStr, num) Then
                    If num >= 1 AndAlso num <= 48 Then
                        wallsColor1 = selectedColor
                    ElseIf num >= 51 AndAlso num <= 98 Then
                        wallsColor2 = selectedColor
                    End If
                End If
            End If
        End If
        PicGame.Invalidate()
    End Sub

    Private Sub lstHatchstyles_SelectedIndexChanged(sender As Object, e As EventArgs) Handles lstHatchstyles.SelectedIndexChanged
        If Not isRestoringState Then SaveStateForUndo()
        Dim selectedHatchStyleName As String = lstHatchstyles.SelectedItem.ToString()
        Dim selectedHatchStyle As HatchStyle = CType([Enum].Parse(GetType(HatchStyle), selectedHatchStyleName), HatchStyle)
        If rdoFloors.Checked Then
            floorsHatchStyle = selectedHatchStyle
        ElseIf rdoWalls.Checked Then
            wallsHatchStyle = selectedHatchStyle
        End If
        PicGame.Invalidate()
    End Sub

    Private Sub btnRandom_Click(sender As Object, e As EventArgs) Handles btnRandom.Click
        If Not isRestoringState Then SaveStateForUndo()
        floorsColor1 = Color.FromArgb(rand.Next(256), rand.Next(256), rand.Next(256))
        floorsColor2 = Color.FromArgb(rand.Next(256), rand.Next(256), rand.Next(256))
        wallsColor1 = Color.FromArgb(rand.Next(256), rand.Next(256), rand.Next(256))
        wallsColor2 = Color.FromArgb(rand.Next(256), rand.Next(256), rand.Next(256))
        Dim hatchStyles As Array = [Enum].GetValues(GetType(HatchStyle))
        floorsHatchStyle = CType(hatchStyles.GetValue(rand.Next(hatchStyles.Length)), HatchStyle)
        wallsHatchStyle = CType(hatchStyles.GetValue(rand.Next(hatchStyles.Length)), HatchStyle)
        PicGame.Invalidate()
    End Sub

    Private Sub HighlightTimer_Tick(sender As Object, e As EventArgs)
        highlightColor = If(highlightColor = Color.Black, Color.White, Color.Black)
        PicGame.Invalidate()
    End Sub

    Private Sub PicGame_Paint(sender As Object, e As PaintEventArgs) Handles PicGame.Paint
        CalculateSquareSizes(PicGame.ClientSize.Width, PicGame.ClientSize.Height)
        e.Graphics.SmoothingMode = SmoothingMode.None
        DrawGrid(e.Graphics)
        DrawHighlight(e.Graphics)
        DrawItems(e.Graphics)
    End Sub

    Private Sub PicGame_MouseMove(sender As Object, e As MouseEventArgs) Handles PicGame.MouseMove
        HandleMouseMove(e)
    End Sub

    Private Sub PicGame_MouseDown(sender As Object, e As MouseEventArgs) Handles PicGame.MouseDown
        If e.Button = MouseButtons.Right Then
            isRightMouseDown = True
        ElseIf e.Button = MouseButtons.Left Then
            isLeftMouseDown = True
        End If
        If (isLeftMouseDown OrElse isRightMouseDown) AndAlso Not isActionInProgress Then
            isActionInProgress = True
            If Not isRestoringState Then SaveStateForUndo()
        End If
    End Sub

    Private Sub PicGame_MouseUp(sender As Object, e As MouseEventArgs) Handles PicGame.MouseUp
        If e.Button = MouseButtons.Right Then
            isRightMouseDown = False
        ElseIf e.Button = MouseButtons.Left Then
            isLeftMouseDown = False
        End If
        If isActionInProgress Then
            isActionInProgress = False
            redoStack.Clear()
        End If
    End Sub

    Private Sub CalculateSquareSizes(clientWidth As Integer, clientHeight As Integer)
        SQSizeX = (clientWidth - 2 * INDENT - (SQCOUNTX + 1) * WALL_THICKNESS) \ SQCOUNTX
        SQSizeY = (clientHeight - 2 * INDENT - (SQCOUNTY + 1) * WALL_THICKNESS) \ SQCOUNTY
    End Sub

    Private Sub DrawGrid(g As Graphics)
        Dim xOrigin As Integer = INDENT
        Dim yOrigin As Integer = INDENT
        For x As Integer = 0 To SQCOUNTX - 1
            For y As Integer = 0 To SQCOUNTY - 1
                Dim floorRect As New Rectangle(xOrigin + x * (SQSizeX + WALL_THICKNESS) + WALL_THICKNESS \ 2,
                                               yOrigin + y * (SQSizeY + WALL_THICKNESS) + WALL_THICKNESS \ 2,
                                               SQSizeX, SQSizeY)
                Using floorBrush As New HatchBrush(floorsHatchStyle, floorsColor1, floorsColor2)
                    g.FillRectangle(floorBrush, floorRect)
                End Using
            Next
        Next
        For x As Integer = 0 To SQCOUNTX
            For y As Integer = 0 To SQCOUNTY
                Dim intersectionRect As New Rectangle(xOrigin + x * (SQSizeX + WALL_THICKNESS) - WALL_THICKNESS \ 2,
                                                      yOrigin + y * (SQSizeY + WALL_THICKNESS) - WALL_THICKNESS \ 2,
                                                      WALL_THICKNESS, WALL_THICKNESS)
                Using wallBrush As New HatchBrush(wallsHatchStyle, wallsColor1, wallsColor2)
                    g.FillRectangle(wallBrush, intersectionRect)
                End Using
            Next
        Next
        For x As Integer = 0 To SQCOUNTX
            For y As Integer = 0 To SQCOUNTY - 1
                Dim wallRect As New Rectangle(xOrigin + x * (SQSizeX + WALL_THICKNESS) - WALL_THICKNESS \ 2,
                                              yOrigin + y * (SQSizeY + WALL_THICKNESS) + WALL_THICKNESS \ 2,
                                              WALL_THICKNESS, SQSizeY)
                If verticalWalls(x, y) Then
                    Using wallBrush As New HatchBrush(wallsHatchStyle, wallsColor1, wallsColor2)
                        g.FillRectangle(wallBrush, wallRect)
                    End Using
                Else
                    Using floorBrush As New HatchBrush(floorsHatchStyle, floorsColor1, floorsColor2)
                        g.FillRectangle(floorBrush, wallRect)
                    End Using
                End If
            Next
        Next
        For x As Integer = 0 To SQCOUNTX - 1
            For y As Integer = 0 To SQCOUNTY
                Dim wallRect As New Rectangle(xOrigin + x * (SQSizeX + WALL_THICKNESS) + WALL_THICKNESS \ 2,
                                              yOrigin + y * (SQSizeY + WALL_THICKNESS) - WALL_THICKNESS \ 2,
                                              SQSizeX, WALL_THICKNESS)
                If horizontalWalls(x, y) Then
                    Using wallBrush As New HatchBrush(wallsHatchStyle, wallsColor1, wallsColor2)
                        g.FillRectangle(wallBrush, wallRect)
                    End Using
                Else
                    Using floorBrush As New HatchBrush(floorsHatchStyle, floorsColor1, floorsColor2)
                        g.FillRectangle(floorBrush, wallRect)
                    End Using
                End If
            Next
        Next
    End Sub

    Private Sub DrawHighlight(g As Graphics)
        If Not isMouseInsideGrid Then Return
        If highlightedWallX Then
            If highlightedX = 0 OrElse highlightedX = SQCOUNTX Then Return
        ElseIf highlightedWallY Then
            If highlightedY = 0 OrElse highlightedY = SQCOUNTY Then Return
        End If
        Using highlightPen As New Pen(highlightColor, 2) With {.DashStyle = DashStyle.Dash}
            If highlightedWallX Then
                Dim x As Integer = INDENT + highlightedX * (SQSizeX + WALL_THICKNESS) - WALL_THICKNESS \ 2
                Dim y As Integer = INDENT + highlightedY * (SQSizeY + WALL_THICKNESS) + WALL_THICKNESS \ 2
                Dim highlightRect As New Rectangle(x, y, WALL_THICKNESS, SQSizeY)
                g.DrawRectangle(highlightPen, highlightRect)
            ElseIf highlightedWallY Then
                Dim x As Integer = INDENT + highlightedX * (SQSizeX + WALL_THICKNESS) + WALL_THICKNESS \ 2
                Dim y As Integer = INDENT + highlightedY * (SQSizeY + WALL_THICKNESS) - WALL_THICKNESS \ 2
                Dim highlightRect As New Rectangle(x, y, SQSizeX, WALL_THICKNESS)
                g.DrawRectangle(highlightPen, highlightRect)
            Else
                Dim x As Integer = INDENT + highlightedX * (SQSizeX + WALL_THICKNESS) + WALL_THICKNESS \ 2
                Dim y As Integer = INDENT + highlightedY * (SQSizeY + WALL_THICKNESS) + WALL_THICKNESS \ 2
                Dim highlightRect As New Rectangle(x, y, SQSizeX, SQSizeY)
                g.DrawRectangle(highlightPen, highlightRect)
            End If
        End Using
    End Sub

    Private Sub HandleMouseMove(e As MouseEventArgs)
        isMouseInsideGrid = False
        Dim mouseX As Integer = e.X - INDENT
        Dim mouseY As Integer = e.Y - INDENT
        If mouseX < 0 OrElse mouseY < 0 Then Exit Sub
        Dim cellSizeX As Integer = SQSizeX + WALL_THICKNESS
        Dim cellSizeY As Integer = SQSizeY + WALL_THICKNESS
        Dim cellX As Integer = mouseX \ cellSizeX
        Dim cellY As Integer = mouseY \ cellSizeY
        If cellX >= SQCOUNTX OrElse cellY >= SQCOUNTY Then Exit Sub
        Dim offsetX As Integer = mouseX Mod cellSizeX
        Dim offsetY As Integer = mouseY Mod cellSizeY
        highlightedWallX = False
        highlightedWallY = False
        If offsetX < WALL_THICKNESS Then
            highlightedWallX = True
            highlightedX = cellX
            highlightedY = cellY
            If highlightedX = 0 OrElse highlightedX = SQCOUNTX Then Exit Sub
            isMouseInsideGrid = True
        ElseIf offsetY < WALL_THICKNESS Then
            highlightedWallY = True
            highlightedX = cellX
            highlightedY = cellY
            If highlightedY = 0 OrElse highlightedY = SQCOUNTY Then Exit Sub
            isMouseInsideGrid = True
        ElseIf offsetX >= WALL_THICKNESS AndAlso offsetY >= WALL_THICKNESS Then
            highlightedX = cellX
            highlightedY = cellY
            isMouseInsideGrid = True
        Else
            Exit Sub
        End If
        If isLeftMouseDown OrElse isRightMouseDown Then
            If highlightedWallX Then
                If highlightedX > 0 AndAlso highlightedX < SQCOUNTX Then
                    If isLeftMouseDown Then
                        verticalWalls(highlightedX, highlightedY) = True
                    ElseIf isRightMouseDown Then
                        verticalWalls(highlightedX, highlightedY) = False
                    End If
                End If
            ElseIf highlightedWallY Then
                If highlightedY > 0 AndAlso highlightedY < SQCOUNTY Then
                    If isLeftMouseDown Then
                        horizontalWalls(highlightedX, highlightedY) = True
                    ElseIf isRightMouseDown Then
                        horizontalWalls(highlightedX, highlightedY) = False
                    End If
                End If
            End If
            PicGame.Invalidate()
        End If
        PicGame.Invalidate()
    End Sub

    Private Sub PicBlock_MouseDown(sender As Object, e As MouseEventArgs)
        Dim pic As PictureBox = CType(sender, PictureBox)
        pic.DoDragDrop(pic.Name, DragDropEffects.Copy)
    End Sub

    Private Sub PicGame_DragEnter(sender As Object, e As DragEventArgs)
        If e.Data.GetDataPresent(DataFormats.StringFormat) Then
            Dim data As String = CStr(e.Data.GetData(DataFormats.StringFormat))
            If data = "picBlock1" OrElse data = "picBaddy1" Then
                e.Effect = DragDropEffects.Copy
            Else
                e.Effect = DragDropEffects.None
            End If
        Else
            e.Effect = DragDropEffects.None
        End If
    End Sub

    Private Sub PicGame_DragOver(sender As Object, e As DragEventArgs)
        e.Effect = DragDropEffects.Copy
    End Sub

    Private Sub PicGame_DragDrop(sender As Object, e As DragEventArgs)
        Dim clientPoint As Point = PicGame.PointToClient(New Point(e.X, e.Y))
        Dim mouseX As Integer = clientPoint.X - INDENT
        Dim mouseY As Integer = clientPoint.Y - INDENT
        If mouseX < 0 OrElse mouseY < 0 Then Return
        Dim cellSizeX As Integer = SQSizeX + WALL_THICKNESS
        Dim cellSizeY As Integer = SQSizeY + WALL_THICKNESS
        Dim cellX As Integer = mouseX \ cellSizeX
        Dim cellY As Integer = mouseY \ cellSizeY
        If cellX >= SQCOUNTX OrElse cellY >= SQCOUNTY Then Return
        Dim offsetX As Integer = mouseX Mod cellSizeX
        Dim offsetY As Integer = mouseY Mod cellSizeY
        If offsetX >= WALL_THICKNESS AndAlso offsetY >= WALL_THICKNESS Then
            If e.Data.GetDataPresent(DataFormats.StringFormat) Then
                Dim imgKey As String = CStr(e.Data.GetData(DataFormats.StringFormat))
                If imgKey = "picBlock1" OrElse imgKey = "picBaddy1" Then
                    If Not isRestoringState Then SaveStateForUndo()
                    PlaceItem(cellX, cellY, imgKey)
                    redoStack.Clear()
                    PicGame.Invalidate()
                End If
            End If
        End If
    End Sub

    Private Sub PlaceItem(x As Integer, y As Integer, imgKey As String)
        placedItems(x, y) = New PlacedItem() With {.ImageKey = imgKey}
    End Sub

    Private Sub DrawItems(g As Graphics)
        For x As Integer = 0 To SQCOUNTX - 1
            For y As Integer = 0 To SQCOUNTY - 1
                If placedItems(x, y) IsNot Nothing Then
                    Dim xPos As Integer = INDENT + x * (SQSizeX + WALL_THICKNESS) + WALL_THICKNESS \ 2
                    Dim yPos As Integer = INDENT + y * (SQSizeY + WALL_THICKNESS) + WALL_THICKNESS \ 2
                    Dim img As Image = GetImageByKey(placedItems(x, y).ImageKey)
                    If img IsNot Nothing Then
                        g.DrawImage(img, New Rectangle(xPos, yPos, SQSizeX, SQSizeY))
                    End If
                End If
            Next
        Next
    End Sub

    Private Function GetImageByKey(imageKey As String) As Image
        Select Case imageKey
            Case "picBlock1"
                Return picBlock1.Image
            Case "picBaddy1"
                Return picBaddy1.Image
            Case Else
                Return Nothing
        End Select
    End Function

End Class

' RoomState and PlacedItem classes remain the same as before
Public Class RoomState
    Public wallsHatchStyle As HatchStyle
    Public wallsColor1 As Color
    Public wallsColor2 As Color
    Public floorsHatchStyle As HatchStyle
    Public floorsColor1 As Color
    Public floorsColor2 As Color

    Public verticalWalls(,) As Boolean
    Public horizontalWalls(,) As Boolean

    Public placedItems(,) As PlacedItemData

    Public Sub New(roomView As RoomView)
        wallsHatchStyle = roomView.wallsHatchStyle
        wallsColor1 = roomView.wallsColor1
        wallsColor2 = roomView.wallsColor2
        floorsHatchStyle = roomView.floorsHatchStyle
        floorsColor1 = roomView.floorsColor1
        floorsColor2 = roomView.floorsColor2
        verticalWalls = CType(roomView.verticalWalls.Clone(), Boolean(,))
        horizontalWalls = CType(roomView.horizontalWalls.Clone(), Boolean(,))
        Dim itemData(roomView.placedItems.GetLength(0) - 1, roomView.placedItems.GetLength(1) - 1) As PlacedItemData
        For x As Integer = 0 To roomView.placedItems.GetLength(0) - 1
            For y As Integer = 0 To roomView.placedItems.GetLength(1) - 1
                If roomView.placedItems(x, y) IsNot Nothing Then
                    itemData(x, y) = New PlacedItemData With {.ImageKey = roomView.placedItems(x, y).ImageKey}
                End If
            Next
        Next
        placedItems = itemData
    End Sub

    Public Sub RestoreState(roomView As RoomView)
        roomView.wallsHatchStyle = wallsHatchStyle
        roomView.wallsColor1 = wallsColor1
        roomView.wallsColor2 = wallsColor2
        roomView.floorsHatchStyle = floorsHatchStyle
        roomView.floorsColor1 = floorsColor1
        roomView.floorsColor2 = floorsColor2
        roomView.verticalWalls = CType(verticalWalls.Clone(), Boolean(,))
        roomView.horizontalWalls = CType(horizontalWalls.Clone(), Boolean(,))
        roomView.placedItems = New PlacedItem(placedItems.GetLength(0) - 1, placedItems.GetLength(1) - 1) {}
        For x As Integer = 0 To placedItems.GetLength(0) - 1
            For y As Integer = 0 To placedItems.GetLength(1) - 1
                If placedItems(x, y) IsNot Nothing Then
                    roomView.placedItems(x, y) = New PlacedItem With {.ImageKey = placedItems(x, y).ImageKey}
                End If
            Next
        Next
    End Sub
End Class

Public Class PlacedItem
    Public Property ImageKey As String
End Class

Public Class PlacedItemData
    Public Property ImageKey As String
End Class

I've added to the desinger 2 more pictureboxes, I have named them picBaddy2 and picLazer. Can you add these 2 to the drag and drop section. picLazer works differently to picBaddy1, picbaddy2 and picBlock1. PicLazer can only be attached to a wall (if the wall is not there, been removed by the user) if the Lazer is dragged to a horizontal wall it must be rotated 90 degrees. On a vertical wall no changes to rotation.