#include "DismembermentGraph/SDismembermentPreviewViewport.h"
#include "DismembermentGraph/DismembermentPreviewManager.h"
#include "EditorViewportClient.h"
#include "SEditorViewport.h"
#include "PreviewScene.h"
#include "Components/SkeletalMeshComponent.h"
#include "Engine/SkeletalMesh.h"
#include "Engine/StaticMesh.h"
#include "AdvancedPreviewScene.h"
#include "AssetViewerSettings.h"
#include "EditorModeManager.h"
#include "EngineUtils.h"
#include "GameFramework/Actor.h"
#include "Engine/World.h"

// Preview actor class
class ADismembermentPreviewActor : public AActor
{
public:
    ADismembermentPreviewActor(const FObjectInitializer& ObjectInitializer)
        : AActor(ObjectInitializer)
    {
        // Create a skeletal mesh component
        SkeletalMeshComponent = CreateDefaultSubobject<USkeletalMeshComponent>(TEXT("SkeletalMeshComponent"));
        RootComponent = SkeletalMeshComponent;
    }

    // Set the skeletal mesh
    void SetSkeletalMesh(USkeletalMesh* InSkeletalMesh)
    {
        if (SkeletalMeshComponent)
        {
            SkeletalMeshComponent->SetSkeletalMesh(InSkeletalMesh);
        }
    }

    // Get the skeletal mesh component
    USkeletalMeshComponent* GetSkeletalMeshComponent() const
    {
        return SkeletalMeshComponent;
    }

private:
    // The skeletal mesh component
    UPROPERTY()
    TObjectPtr<USkeletalMeshComponent> SkeletalMeshComponent;
};

//////////////////////////////////////////////////////////////////////////
// SDismembermentPreviewViewport

void SDismembermentPreviewViewport::Construct(const FArguments& InArgs)
{
    // Create the preview scene
    PreviewScene = MakeShareable(new FAdvancedPreviewScene(FPreviewScene::ConstructionValues()));

    // Set up the preview scene
    UAssetViewerSettings* Settings = UAssetViewerSettings::Get();
    const int32 ProfileIndex = Settings->Profiles.IsValidIndex(0) ? 0 : INDEX_NONE;
    if (ProfileIndex != INDEX_NONE)
    {
        Settings->Profiles[ProfileIndex].bRotateLightingRig = false;
        PreviewScene->SetLightDirection(FRotator(Settings->Profiles[ProfileIndex].LightingRigRotation, 0, 0));
        // The following methods may not exist in UE5.5.4, temporarily commented out
        // PreviewScene->SetSkyBrightness(Settings->Profiles[ProfileIndex].EnvironmentCubeMapBrightness);
        // PreviewScene->SetLightColor(Settings->Profiles[ProfileIndex].LightColor);
        // PreviewScene->SetLightIntensity(Settings->Profiles[ProfileIndex].LightBrightness);
        // PreviewScene->SetEnvironmentCubeMap(Settings->Profiles[ProfileIndex].EnvironmentCubeMap);
    }

    // Create the preview actor
    CreatePreviewActor();

    // Call parent constructor
    SEditorViewport::Construct(SEditorViewport::FArguments());
}

SDismembermentPreviewViewport::~SDismembermentPreviewViewport()
{
    // Clean up the preview scene
    if (PreviewScene.IsValid())
    {
        // Cannot remove Actor directly, need to remove the Actor's root component
        if (PreviewActor && PreviewActor->GetRootComponent())
        {
            PreviewScene->RemoveComponent(PreviewActor->GetRootComponent());
        }
        PreviewScene.Reset();
    }

    // Clean up the preview actor
    if (PreviewActor)
    {
        if (PreviewActor->GetWorld())
        {
            PreviewActor->Destroy();
        }
        PreviewActor = nullptr;
    }

    // Clean up the viewport client
    ViewportClient.Reset();
}

void SDismembermentPreviewViewport::SetPreviewManager(UDismembermentPreviewManager* InPreviewManager)
{
    PreviewManager = InPreviewManager;

    // Update the viewport client
    if (ViewportClient.IsValid())
    {
        ViewportClient->SetPreviewManager(PreviewManager);
    }

    // Update the preview actor
    if (PreviewManager && PreviewActor)
    {
        PreviewManager->SetTargetActor(PreviewActor);
    }
}

void SDismembermentPreviewViewport::SetPreviewSkeletalMesh(USkeletalMesh* InSkeletalMesh)
{
    // Set the skeletal mesh on the preview actor
    if (PreviewActor)
    {
        ADismembermentPreviewActor* PreviewActorCasted = Cast<ADismembermentPreviewActor>(PreviewActor);
        if (PreviewActorCasted)
        {
            PreviewActorCasted->SetSkeletalMesh(InSkeletalMesh);
        }
    }

    // Update the preview actor
    UpdatePreviewActor();
}

AActor* SDismembermentPreviewViewport::GetPreviewActor() const
{
    return PreviewActor;
}

void SDismembermentPreviewViewport::RefreshViewport()
{
    // Invalidate the viewport
    if (ViewportClient.IsValid())
    {
        ViewportClient->Invalidate();
    }
}

TSharedRef<FEditorViewportClient> SDismembermentPreviewViewport::MakeEditorViewportClient()
{
    // Create the viewport client
    ViewportClient = MakeShareable(new SDismembermentPreviewViewportClient(PreviewScene.Get(), SharedThis(this)));
    
    // Set the preview manager
    ViewportClient->SetPreviewManager(PreviewManager);

    // Set up the viewport client
    ViewportClient->SetViewLocation(FVector(0.0f, 0.0f, 200.0f));
    ViewportClient->SetViewRotation(FRotator(-20.0f, 0.0f, 0.0f));
    ViewportClient->SetViewLocationForOrbiting(FVector(0.0f, 0.0f, 0.0f));
    ViewportClient->bSetListenerPosition = false;
    ViewportClient->EngineShowFlags.SetPostProcessing(true);
    ViewportClient->EngineShowFlags.SetLighting(true);
    ViewportClient->EngineShowFlags.SetIndirectLightingCache(true);
    ViewportClient->EngineShowFlags.SetSeparateTranslucency(true);
    ViewportClient->EngineShowFlags.SetTemporalAA(true);
    ViewportClient->EngineShowFlags.SetGrid(false);
    ViewportClient->EngineShowFlags.SetAtmosphere(true);
    ViewportClient->EngineShowFlags.SetSkeletalMeshes(true);
    ViewportClient->EngineShowFlags.SetDecals(true);
    ViewportClient->EngineShowFlags.SetParticles(true);
    ViewportClient->EngineShowFlags.SetVolumetricFog(true);
    ViewportClient->EngineShowFlags.SetDynamicShadows(true);
    ViewportClient->EngineShowFlags.SetSkyLighting(true);
    ViewportClient->EngineShowFlags.SetAmbientOcclusion(true);
    ViewportClient->EngineShowFlags.SetScreenSpaceReflections(true);
    ViewportClient->EngineShowFlags.SetAntiAliasing(true);
    ViewportClient->EngineShowFlags.SetMotionBlur(false);
    ViewportClient->EngineShowFlags.SetBounds(false);
    ViewportClient->EngineShowFlags.SetCollision(false);
    ViewportClient->EngineShowFlags.SetBSP(false);
    ViewportClient->EngineShowFlags.SetFog(true);
    ViewportClient->EngineShowFlags.SetStaticMeshes(true);
    ViewportClient->EngineShowFlags.SetLandscape(true);
    ViewportClient->EngineShowFlags.SetTranslucency(true);
    ViewportClient->EngineShowFlags.SetInstancedFoliage(true);
    ViewportClient->EngineShowFlags.SetInstancedGrass(true);
    ViewportClient->EngineShowFlags.SetInstancedStaticMeshes(true);
    ViewportClient->EngineShowFlags.SetInstancedFoliage(true);
    ViewportClient->EngineShowFlags.SetInstancedGrass(true);
    ViewportClient->EngineShowFlags.SetInstancedStaticMeshes(true);
    ViewportClient->EngineShowFlags.SetSplines(true);
    ViewportClient->EngineShowFlags.SetSelectionOutline(true);
    ViewportClient->EngineShowFlags.SetMeshEdges(false);
    ViewportClient->EngineShowFlags.SetVertexColors(false);
    ViewportClient->EngineShowFlags.SetLightComplexity(false);
    ViewportClient->EngineShowFlags.SetShaderComplexity(false);
    // The following methods may not exist in UE5.5.4, temporarily commented out
    // ViewportClient->EngineShowFlags.SetStaticMeshLODColoration(false);
    ViewportClient->EngineShowFlags.SetLightMapDensity(false);
    ViewportClient->EngineShowFlags.SetLightInfluences(false);
    // ViewportClient->EngineShowFlags.SetOptimizeVizibility(false);
    ViewportClient->EngineShowFlags.SetTextRender(true);
    ViewportClient->EngineShowFlags.SetTestImage(false);
    ViewportClient->EngineShowFlags.SetVisualizeLightCulling(false);
    ViewportClient->EngineShowFlags.SetPrecomputedVisibility(true);
    ViewportClient->EngineShowFlags.SetPrecomputedVisibilityCells(false);
    ViewportClient->EngineShowFlags.SetVolumes(false);
    ViewportClient->EngineShowFlags.SetGame(false);
    ViewportClient->EngineShowFlags.SetBSPTriangles(false);
    ViewportClient->EngineShowFlags.SetHitProxies(false);
    ViewportClient->EngineShowFlags.SetGBufferHints(false);
    ViewportClient->EngineShowFlags.SetVisualizeShadingModels(false);
    // The following methods may not exist in UE5.5.4, temporarily commented out
    // ViewportClient->EngineShowFlags.SetVisualizeAdaptiveDOF(false);
    ViewportClient->EngineShowFlags.SetVisualizeSSR(false);
    ViewportClient->EngineShowFlags.SetVisualizeSSS(false);
    ViewportClient->EngineShowFlags.SetVolumetricLightmap(true);
    ViewportClient->EngineShowFlags.SetVisualizeOutOfBoundsPixels(false);
    ViewportClient->EngineShowFlags.SetHighResScreenshotMask(false);
    ViewportClient->EngineShowFlags.SetHMDDistortion(false);
    ViewportClient->EngineShowFlags.SetStereoRendering(false);
    ViewportClient->EngineShowFlags.SetTonemapper(true);
    ViewportClient->EngineShowFlags.SetLumenReflections(true);
    ViewportClient->EngineShowFlags.SetLumenGlobalIllumination(true);
    // ViewportClient->EngineShowFlags.SetVirtualShadowMaps(true);
    // ViewportClient->EngineShowFlags.SetNanite(true);

    return ViewportClient.ToSharedRef();
}

void SDismembermentPreviewViewport::OnFocusViewportToSelection()
{
    // Focus the viewport on the preview actor
    if (ViewportClient.IsValid() && PreviewActor)
    {
        ViewportClient->FocusViewportOnBox(PreviewActor->GetComponentsBoundingBox());
    }
}

bool SDismembermentPreviewViewport::IsVisible() const
{
    return SEditorViewport::IsVisible();
}

void SDismembermentPreviewViewport::CreatePreviewActor()
{
    // Clean up any existing preview actor
    if (PreviewActor)
    {
        if (PreviewActor->GetWorld())
        {
            PreviewActor->Destroy();
        }
        PreviewActor = nullptr;
    }

    // Create a new preview actor
    if (PreviewScene.IsValid())
    {
        UWorld* World = PreviewScene->GetWorld();
        if (World)
        {
            PreviewActor = World->SpawnActor<ADismembermentPreviewActor>();
            if (PreviewActor)
            {
                // Add the actor to the preview scene
                PreviewScene->AddComponent(PreviewActor->GetRootComponent(), FTransform::Identity);

                // Update the preview manager
                if (PreviewManager)
                {
                    PreviewManager->SetTargetActor(PreviewActor);
                }
            }
        }
    }
}

void SDismembermentPreviewViewport::UpdatePreviewActor()
{
    // Update the preview manager
    if (PreviewManager && PreviewActor)
    {
        PreviewManager->SetTargetActor(PreviewActor);
    }

    // Refresh the viewport
    RefreshViewport();
}

//////////////////////////////////////////////////////////////////////////
// SDismembermentPreviewViewportClient

SDismembermentPreviewViewportClient::SDismembermentPreviewViewportClient(FPreviewScene* InPreviewScene, const TWeakPtr<SDismembermentPreviewViewport>& InViewportWidget)
    : FEditorViewportClient(nullptr, InPreviewScene, InViewportWidget)
    , ViewportWidget(InViewportWidget)
    , PreviewManager(nullptr)
{
    // Set up the viewport client
    SetRealtime(true);
    
    // In UE5.5.4, these variables may no longer be members of FEditorViewportClient
    // We only keep the necessary settings, others are commented out
    bSetListenerPosition = false;
    bShouldCheckHitProxy = true;
    
    // The following settings may no longer be supported in UE5.5.4, temporarily commented out
    /*
    bShowGrid = false;
    bDisableInput = false;
    bAllowMatineePreview = false;
    bUsingOrbitCamera = true;
    bForceInitialFocus = true;
    bShowBounds = false;
    bShowBoundsActors = false;
    bShowFloor = true;
    bShowBoxes = false;
    bShowWireframe = false;
    bShowCollision = false;
    bShowSockets = true;
    bDrawAxes = false;
    bShowNormals = false;
    bShowTangents = false;
    bShowBinormals = false;
    bShowConstraints = false;
    bShowCameras = false;
    bShowLightRadius = false;
    bShowLightVolumes = false;
    bShowLightInfluences = false;
    bShowLightFrustums = false;
    bShowShadowFrustums = false;
    bShowLightingOnlyOverlap = false;
    bShowLightingVisualization = false;
    bShowLightingStats = false;
    bShowShadowDensity = false;
    bShowPhysicalMaterialMasks = false;
    bShowSpriteSockets = false;
    bShowParticleSystemComponents = true;
    bShowParticleSystems = true;
    bShowLOD = false;
    bShowHUD = false;
    bShowDebugInfo = false;
    bDrawPreviewShadowsInGame = false;
    bEnableDirectLightMap = true;
    bEnableIndirectLightMap = true;
    // All bEnableColorize related variables are commented out
    */
    // Continue to comment out more bEnableColorize related variables
    /*
    bEnableColorizeVolumetricLightmapIndirectIntensityGrayscale = false;
    bEnableColorizeVolumetricLightmapIndirectIntensityColor = false;
    bEnableColorizeVolumetricLightmapIndirectIntensityGrayscale = false;
    bEnableColorizeVolumetricLightmapIndirectIntensityColor = false;
    bEnableColorizeVolumetricLightmapIndirectIntensity = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusion = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusionGrayscale = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusionColor = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusionGrayscale = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusionColor = false;
    bEnableColorizeVolumetricLightmapAmbientOcclusion = false;
    */
    // Continue to comment out all remaining bEnableColorize related variables
    /*
    bEnableColorizeVolumetricLightmapSHBand0 = false;
    bEnableColorizeVolumetricLightmapSHBand0Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand0Color = false;
    bEnableColorizeVolumetricLightmapSHBand0Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand0Color = false;
    bEnableColorizeVolumetricLightmapSHBand1 = false;
    bEnableColorizeVolumetricLightmapSHBand1Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand1Color = false;
    bEnableColorizeVolumetricLightmapSHBand1Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand1Color = false;
    bEnableColorizeVolumetricLightmapSHBand1 = false;
    bEnableColorizeVolumetricLightmapSHBand2 = false;
    bEnableColorizeVolumetricLightmapSHBand2Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand2Color = false;
    bEnableColorizeVolumetricLightmapSHBand2Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand2Color = false;
    bEnableColorizeVolumetricLightmapSHBand2 = false;
    bEnableColorizeVolumetricLightmapSHBand3 = false;
    bEnableColorizeVolumetricLightmapSHBand3Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand3Color = false;
    bEnableColorizeVolumetricLightmapSHBand3Grayscale = false;
    bEnableColorizeVolumetricLightmapSHBand3Color = false;
    bEnableColorizeVolumetricLightmapSHBand3 = false;
    */
}

SDismembermentPreviewViewportClient::~SDismembermentPreviewViewportClient()
{
    // Clean up
    PreviewManager = nullptr;
}

void SDismembermentPreviewViewportClient::Tick(float DeltaSeconds)
{
    // Call parent tick
    FEditorViewportClient::Tick(DeltaSeconds);

    // Tick the preview manager
    if (PreviewManager)
    {
        PreviewManager->Tick(DeltaSeconds);
    }
}

void SDismembermentPreviewViewportClient::Draw(const FSceneView* View, FPrimitiveDrawInterface* PDI)
{
    // Call parent draw
    FEditorViewportClient::Draw(View, PDI);

    // Draw preview elements
    if (PreviewManager)
    {
        // TODO: Draw preview elements
    }
}

void SDismembermentPreviewViewportClient::DrawCanvas(FViewport& InViewport, FSceneView& View, FCanvas& Canvas)
{
    // Call parent draw canvas
    FEditorViewportClient::DrawCanvas(InViewport, View, Canvas);

    // Draw preview elements
    if (PreviewManager)
    {
        // TODO: Draw preview elements
    }
}

void SDismembermentPreviewViewportClient::SetPreviewManager(UDismembermentPreviewManager* InPreviewManager)
{
    PreviewManager = InPreviewManager;
}