#include "MatrixInputWidget.h"
#include "SlateOptMacros.h"
#include "Widgets/Layout/SUniformGridPanel.h"
#include "Widgets/Text/STextBlock.h"
#include "Widgets/Input/SButton.h"
#include "Widgets/Layout/SBox.h"
#include "Widgets/Layout/SScrollBox.h"
#include "Styling/AppStyle.h"

BEGIN_SLATE_FUNCTION_BUILD_OPTIMIZATION

void SMatrixInputWidget::Construct(const FArguments& InArgs)
{
    Matrix = InArgs._Matrix;
    OnMatrixChanged = InArgs._OnMatrixChanged;
    Rows = InArgs._Rows;
    Columns = InArgs._Columns;

    // Create a grid panel for the matrix elements
    GridPanel = SNew(SGridPanel);

    // Create numeric entry boxes for each matrix element
    NumericEntryBoxes.SetNum(Rows * Columns); // Allocate numeric entry boxes based on rows and columns

    // Add labels for rows and columns
    for (int32 Col = 0; Col < Columns; ++Col)
    {
        FString ColLabel;
        switch (Col)
        {
        case 0: ColLabel = "X"; break;
        case 1: ColLabel = "Y"; break;
        case 2: ColLabel = "Z"; break;
        default: ColLabel = FString::Printf(TEXT("Column %d"), Col);
        }

        GridPanel->AddSlot(Col + 1, 0)
        [
            SNew(STextBlock)
            .Text(FText::FromString(ColLabel))
            .Margin(FMargin(5.0f))
        ];
    }

    for (int32 Row = 0; Row < Rows; ++Row)
    {
        FString RowLabel;
        switch (Row)
        {
        case 0: RowLabel = "Right"; break;
        case 1: RowLabel = "Forward"; break;
        case 2: RowLabel = "Up"; break;
        case 3: RowLabel = "Position"; break;
        default: RowLabel = FString::Printf(TEXT("Row %d"), Row);
        }

        GridPanel->AddSlot(0, Row + 1)
        [
            SNew(STextBlock)
            .Text(FText::FromString(RowLabel))
            .Margin(FMargin(5.0f))
        ];
    }

    // Add numeric entry boxes for each matrix element
    for (int32 Row = 0; Row < Rows; ++Row)
    {
        for (int32 Col = 0; Col < Columns; ++Col)
        {
            TSharedPtr<SNumericEntryBox<float>> NumericEntryBox = CreateMatrixElementWidget(Row, Col);
            NumericEntryBoxes[Row * Columns + Col] = NumericEntryBox;

            GridPanel->AddSlot(Col + 1, Row + 1)
            [
                SNew(SBox)
                .WidthOverride(80.0f)
                .Padding(FMargin(2.0f))
                [
                    NumericEntryBox.ToSharedRef()
                ]
            ];
        }
    }

    // Create buttons for common operations
    TSharedPtr<SUniformGridPanel> ButtonPanel = SNew(SUniformGridPanel)
        .SlotPadding(FMargin(2.0f));

    ButtonPanel->AddSlot(0, 0)
    [
        SNew(SButton)
        .Text(FText::FromString("Reset to Identity"))
        .ToolTipText(FText::FromString("Reset the matrix to identity"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            ResetToIdentity();
            return FReply::Handled();
        }))
    ];

    ButtonPanel->AddSlot(1, 0)
    [
        SNew(SButton)
        .Text(FText::FromString("Set from Rotation"))
        .ToolTipText(FText::FromString("Set the matrix from rotation values (degrees)"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            SetFromEulerAngles(0.0f, 0.0f, 0.0f);
            return FReply::Handled();
        }))
    ];

    ButtonPanel->AddSlot(2, 0)
    [
        SNew(SButton)
        .Text(FText::FromString("Set from Translation"))
        .ToolTipText(FText::FromString("Set the matrix from translation values"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            SetFromTranslation(FVector::ZeroVector);
            return FReply::Handled();
        }))
    ];

    // Add preset cutting buttons
    ButtonPanel->AddSlot(0, 1)
    [
        SNew(SButton)
        .Text(FText::FromString("Horizontal Cut"))
        .ToolTipText(FText::FromString("Set up a horizontal cutting plane"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            // Set horizontal cutting plane
            FMatrix HorizontalCutMatrix = FMatrix::Identity;
            HorizontalCutMatrix.SetOrigin(FVector(0.0f, 0.0f, 0.0f));
            SetMatrix(HorizontalCutMatrix);
            return FReply::Handled();
        }))
    ];

    ButtonPanel->AddSlot(1, 1)
    [
        SNew(SButton)
        .Text(FText::FromString("Vertical Cut"))
        .ToolTipText(FText::FromString("Set up a vertical cutting plane"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            // Set vertical cutting plane
            FMatrix VerticalCutMatrix = FMatrix::Identity;
            VerticalCutMatrix.SetOrigin(FVector(0.0f, 0.0f, 0.0f));
            // Rotate 90 degrees to make it vertical
            FRotator Rotator(0.0f, 0.0f, 90.0f);
            FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
            SetMatrix(VerticalCutMatrix * RotationMatrix);
            return FReply::Handled();
        }))
    ];

    ButtonPanel->AddSlot(2, 1)
    [
        SNew(SButton)
        .Text(FText::FromString("Diagonal Cut"))
        .ToolTipText(FText::FromString("Set up a diagonal cutting plane"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            // Set diagonal cutting plane
            FMatrix DiagonalCutMatrix = FMatrix::Identity;
            DiagonalCutMatrix.SetOrigin(FVector(0.0f, 0.0f, 0.0f));
            // Rotate 45 degrees
            FRotator Rotator(0.0f, 0.0f, 45.0f);
            FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
            SetMatrix(DiagonalCutMatrix * RotationMatrix);
            return FReply::Handled();
        }))
    ];

    // Add apply cut button
    TSharedPtr<SButton> ApplyCutButton = SNew(SButton)
        .Text(FText::FromString("Apply Cut"))
        .ButtonStyle(FAppStyle::Get(), "FlatButton.Success")
        .ContentPadding(FMargin(8.0f, 4.0f))
        .HAlign(HAlign_Center)
        .ToolTipText(FText::FromString("Apply the current cutting plane to the model"))
        .OnClicked(FOnClicked::CreateLambda([this]()
        {
            // Apply cutting
            if (OnMatrixChanged.IsBound())
            {
                OnMatrixChanged.Execute(Matrix);
            }
            return FReply::Handled();
        }));

    // Create rotation controls
    TSharedPtr<SVerticalBox> RotationControls = SNew(SVerticalBox);
    
    // Add X-axis rotation
    RotationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("X Rotation:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-180.0f)
            .MaxValue(180.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply X-axis rotation
                FRotator Rotator(NewValue, 0.0f, 0.0f);
                FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
                FMatrix NewMatrix = Matrix * RotationMatrix;
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Add Y-axis rotation
    RotationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("Y Rotation:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-180.0f)
            .MaxValue(180.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply Y-axis rotation
                FRotator Rotator(0.0f, NewValue, 0.0f);
                FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
                FMatrix NewMatrix = Matrix * RotationMatrix;
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Add Z-axis rotation
    RotationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("Z Rotation:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-180.0f)
            .MaxValue(180.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply Z-axis rotation
                FRotator Rotator(0.0f, 0.0f, NewValue);
                FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
                FMatrix NewMatrix = Matrix * RotationMatrix;
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Create translation controls
    TSharedPtr<SVerticalBox> TranslationControls = SNew(SVerticalBox);
    
    // Add X-axis translation
    TranslationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("X Position:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-1000.0f)
            .MaxValue(1000.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply X-axis translation
                FVector Translation = Matrix.GetOrigin();
                Translation.X = NewValue;
                FMatrix NewMatrix = Matrix;
                NewMatrix.SetOrigin(Translation);
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Add Y-axis translation
    TranslationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("Y Position:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-1000.0f)
            .MaxValue(1000.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply Y-axis translation
                FVector Translation = Matrix.GetOrigin();
                Translation.Y = NewValue;
                FMatrix NewMatrix = Matrix;
                NewMatrix.SetOrigin(Translation);
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Add Z-axis translation
    TranslationControls->AddSlot()
    .AutoHeight()
    .Padding(0.0f, 4.0f)
    [
        SNew(SHorizontalBox)
        + SHorizontalBox::Slot()
        .AutoWidth()
        .Padding(0.0f, 0.0f, 4.0f, 0.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("Z Position:"))
            .MinDesiredWidth(80.0f)
        ]
        + SHorizontalBox::Slot()
        .FillWidth(1.0f)
        [
            SNew(SSpinBox<float>)
            .MinValue(-1000.0f)
            .MaxValue(1000.0f)
            .Delta(1.0f)
            .Value(0.0f)
            .OnValueChanged_Lambda([this](float NewValue)
            {
                // Apply Z-axis translation
                FVector Translation = Matrix.GetOrigin();
                Translation.Z = NewValue;
                FMatrix NewMatrix = Matrix;
                NewMatrix.SetOrigin(Translation);
                SetMatrix(NewMatrix);
            })
        ]
    ];

    // Create main layout
    ChildSlot
    [
        SNew(SVerticalBox)
        + SVerticalBox::Slot()
        .AutoHeight()
        .Padding(0.0f, 0.0f, 0.0f, 8.0f)
        [
            SNew(STextBlock)
            .Text(FText::FromString("Cutting Plane Transformation Matrix"))
            .Font(FAppStyle::GetFontStyle("HeadingFont"))
        ]
        + SVerticalBox::Slot()
        .AutoHeight()
        [
            SNew(SBorder)
            .BorderImage(FAppStyle::GetBrush("ToolPanel.GroupBorder"))
            .Padding(4.0f)
            [
                GridPanel.ToSharedRef()
            ]
        ]
        + SVerticalBox::Slot()
        .AutoHeight()
        .Padding(0.0f, 8.0f)
        [
            SNew(SBorder)
            .BorderImage(FAppStyle::GetBrush("ToolPanel.GroupBorder"))
            .Padding(4.0f)
            [
                SNew(SVerticalBox)
                + SVerticalBox::Slot()
                .AutoHeight()
                [
                    SNew(STextBlock)
                    .Text(FText::FromString("Rotation Controls"))
                    .Font(FAppStyle::GetFontStyle("NormalFontBold"))
                ]
                + SVerticalBox::Slot()
                .AutoHeight()
                .Padding(0.0f, 4.0f)
                [
                    RotationControls.ToSharedRef()
                ]
            ]
        ]
        + SVerticalBox::Slot()
        .AutoHeight()
        .Padding(0.0f, 8.0f)
        [
            SNew(SBorder)
            .BorderImage(FAppStyle::GetBrush("ToolPanel.GroupBorder"))
            .Padding(4.0f)
            [
                SNew(SVerticalBox)
                + SVerticalBox::Slot()
                .AutoHeight()
                [
                    SNew(STextBlock)
                    .Text(FText::FromString("Translation Controls"))
                    .Font(FAppStyle::GetFontStyle("NormalFontBold"))
                ]
                + SVerticalBox::Slot()
                .AutoHeight()
                .Padding(0.0f, 4.0f)
                [
                    TranslationControls.ToSharedRef()
                ]
            ]
        ]
        + SVerticalBox::Slot()
        .AutoHeight()
        .Padding(0.0f, 8.0f)
        [
            SNew(SBorder)
            .BorderImage(FAppStyle::GetBrush("ToolPanel.GroupBorder"))
            .Padding(4.0f)
            [
                SNew(SVerticalBox)
                + SVerticalBox::Slot()
                .AutoHeight()
                [
                    SNew(STextBlock)
                    .Text(FText::FromString("Presets"))
                    .Font(FAppStyle::GetFontStyle("NormalFontBold"))
                ]
                + SVerticalBox::Slot()
                .AutoHeight()
                .Padding(0.0f, 4.0f)
                [
                    ButtonPanel.ToSharedRef()
                ]
            ]
        ]
        + SVerticalBox::Slot()
        .AutoHeight()
        .Padding(0.0f, 8.0f)
        .HAlign(HAlign_Center)
        [
            ApplyCutButton.ToSharedRef()
        ]
    ];

    // Initialize the UI
    UpdateUI();
}

void SMatrixInputWidget::SetMatrix(const FMatrix& InMatrix)
{
    Matrix = InMatrix;
    UpdateUI();

    // Notify listeners of the change
    if (OnMatrixChanged.IsBound())
    {
        OnMatrixChanged.Execute(Matrix);
    }
}

FMatrix SMatrixInputWidget::GetMatrix() const
{
    return Matrix;
}

void SMatrixInputWidget::ResetToIdentity()
{
    SetMatrix(FMatrix::Identity);
}

void SMatrixInputWidget::SetFromEulerAngles(float Roll, float Pitch, float Yaw)
{
    FRotator Rotator(Pitch, Yaw, Roll);
    FMatrix RotationMatrix = FRotationMatrix::Make(Rotator);
    SetMatrix(RotationMatrix);
}

void SMatrixInputWidget::SetFromTranslation(const FVector& Translation)
{
    FMatrix TranslationMatrix = FMatrix::Identity;
    TranslationMatrix.SetOrigin(Translation);
    SetMatrix(TranslationMatrix);
}

void SMatrixInputWidget::SetFromTransform(const FTransform& Transform)
{
    SetMatrix(Transform.ToMatrixWithScale());
}

TSharedPtr<SNumericEntryBox<float>> SMatrixInputWidget::CreateMatrixElementWidget(int32 Row, int32 Col)
{
    return SNew(SNumericEntryBox<float>)
        .Value_Lambda([this, Row, Col]() -> float
        {
            return Matrix.M[Row][Col];
        })
        .OnValueChanged(SNumericEntryBox<float>::FOnValueChanged::CreateSP(this, &SMatrixInputWidget::OnMatrixElementChanged, Row, Col))
        .AllowSpin(true)
        .Delta(0.1f)
        .MinValue(-10000.0f)
        .MaxValue(10000.0f)
        .MinSliderValue(-10.0f)
        .MaxSliderValue(10.0f);
}

void SMatrixInputWidget::OnMatrixElementChanged(float NewValue, int32 Row, int32 Col)
{
    // Update the matrix element
    Matrix.M[Row][Col] = NewValue;

    // Notify listeners of the change
    if (OnMatrixChanged.IsBound())
    {
        OnMatrixChanged.Execute(Matrix);
    }
}

void SMatrixInputWidget::UpdateUI()
{
    // Update each numeric entry box with the current matrix values
    for (int32 Row = 0; Row < Rows; ++Row)
    {
        for (int32 Col = 0; Col < Columns; ++Col)
        {
            int32 Index = Row * Columns + Col;
            if (NumericEntryBoxes.IsValidIndex(Index) && NumericEntryBoxes[Index].IsValid())
            {
                // In UE5.5.4, we need to rebuild the grid panel to update the value
                NumericEntryBoxes[Index] = CreateMatrixElementWidget(Row, Col);
            }
        }
    }
}

END_SLATE_FUNCTION_BUILD_OPTIMIZATION