#include "vtkViewport.h"
#include <vtkPropAssembly.h>
#include <vtkAssembly.h>
#include <vtkPropCollection.h>
#include <vtkProp3D.h>
#include <vtkProp3DCollection.h>
#include <vtkCamera.h>
#include <algorithm>
#include <vtkInteractorStyleTrackballCamera.h>
#include <vtkObjectFactory.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkNew.h>
#include <vtkTextProperty.h>

#include <vtkPolyDataMapper.h>

#include <vtkProperty.h>
#include <vtkCallbackCommand.h>
#include <vtkMath.h>
#include <cmath>
#include <iostream>
#include "vtkHandlerWidget.h"

namespace uLib {
namespace Vtk {

Viewport::Viewport()
    : m_Renderer(vtkSmartPointer<vtkRenderer>::New())
    , m_Annotation(vtkSmartPointer<vtkCornerAnnotation>::New())
    , m_Marker(vtkSmartPointer<vtkOrientationMarkerWidget>::New())
    , m_CameraWidget(nullptr)
    , m_Colors(vtkSmartPointer<vtkNamedColors>::New())
    , m_GridAxis(Y)
{
}

Viewport::~Viewport()
{
    if (m_Renderer) {
        m_Renderer->RemoveAllViewProps();
    }
}

void Viewport::SetupPipeline(vtkRenderWindowInteractor* iren)
{
    if (!iren) return;

    // Trackball-camera interaction style
    vtkNew<vtkInteractorStyleTrackballCamera> style;
    iren->SetInteractorStyle(style);

    // Corner annotation
    m_Annotation->GetTextProperty()->SetColor(1, 1, 1);
    m_Annotation->GetTextProperty()->SetFontFamilyToArial();
    m_Annotation->GetTextProperty()->SetOpacity(0.5);
    m_Annotation->SetMaximumFontSize(10);
    m_Annotation->SetText(0, "uLib VTK viewer.");
    m_Renderer->AddViewProp(m_Annotation);

    // right corner annotation
    m_Annotation->SetText(1, "Grid: -");

    // Orientation axes marker (bottom-left corner)
    vtkNew<vtkAxesActor> axes;
    m_Marker->SetInteractor(iren);
    m_Marker->SetOrientationMarker(axes);
    m_Marker->SetViewport(0.0, 0.0, 0.2, 0.2);
    m_Marker->SetEnabled(true);
    m_Marker->InteractiveOff();

    // Grid Plane centered at (0,0,0)
    m_GridSource = vtkSmartPointer<vtkPlaneSource>::New();
    m_GridActor = vtkSmartPointer<vtkActor>::New();
    vtkNew<vtkPolyDataMapper> gridMapper;
    gridMapper->SetInputConnection(m_GridSource->GetOutputPort());

    m_GridActor->SetMapper(gridMapper);
    m_GridActor->GetProperty()->SetRepresentationToWireframe();
    m_GridActor->GetProperty()->SetColor(0.4, 0.4, 0.4);
    m_GridActor->GetProperty()->SetLighting(0);
    m_GridActor->GetProperty()->SetOpacity(0.5);
    m_GridActor->PickableOff();
    m_Renderer->AddActor(m_GridActor);
    
    // Global Origin Axes
    m_OriginAxes = vtkSmartPointer<vtkAxes>::New();
    m_OriginAxes->SetScaleFactor(1.0); // will be updated
    
    vtkNew<vtkPolyDataMapper> axesMapper;
    axesMapper->SetInputConnection(m_OriginAxes->GetOutputPort());

    m_OriginAxesActor = vtkSmartPointer<vtkActor>::New();
    m_OriginAxesActor->SetMapper(axesMapper);
    m_OriginAxesActor->PickableOff();
    m_Renderer->AddActor(m_OriginAxesActor);

    UpdateGrid();

    // Observe interactor to update grid during interaction
    vtkNew<vtkCallbackCommand> interactionCallback;
    interactionCallback->SetClientData(this);
    interactionCallback->SetCallback([](vtkObject*, unsigned long, void* clientdata, void*){
        static_cast<Viewport*>(clientdata)->UpdateGrid();
    });
    iren->AddObserver(vtkCommand::InteractionEvent, interactionCallback);
    m_Renderer->GetActiveCamera()->AddObserver(vtkCommand::ModifiedEvent, interactionCallback);


    // Camera-orientation widget (VTK >= 9)
#if VTK_MAJOR_VERSION >= 9
    m_CameraWidget = vtkSmartPointer<vtkCameraOrientationWidget>::New();
    m_CameraWidget->SetParentRenderer(m_Renderer);
    m_CameraWidget->SetInteractor(iren);
    m_CameraWidget->On();
#endif
    m_Renderer->SetBackground(0.15, 0.15, 0.15);
    m_Renderer->ResetCamera();

    // Setup layering for overimposed rendering
    if (iren->GetRenderWindow()) {
        iren->GetRenderWindow()->SetNumberOfLayers(2);
        m_Renderer->SetLayer(0);
    }

    // Setup Handler Widget
    m_HandlerWidget = vtkSmartPointer<vtkHandlerWidget>::New();
    m_HandlerWidget->SetInteractor(iren);
    m_HandlerWidget->SetCurrentRenderer(m_Renderer);
    if (m_HandlerWidget->GetOverlayRenderer()) {
        m_HandlerWidget->GetOverlayRenderer()->SetLayer(1);
    }

    // Observe InteractionEvent to update the selected puppet when the widget moves it
    vtkNew<vtkCallbackCommand> widgetInteractionCallback;
    widgetInteractionCallback->SetClientData(this);
    widgetInteractionCallback->SetCallback([](vtkObject*, unsigned long, void* clientdata, void*){
        auto* self = static_cast<Viewport*>(clientdata);
        for (auto* p : self->m_Puppets) {
            if (p->IsSelected()) {
                p->Update();
            }
        }
    });
    m_HandlerWidget->AddObserver(vtkCommand::InteractionEvent, widgetInteractionCallback);

    // Picking for selection
    m_Picker = vtkSmartPointer<vtkCellPicker>::New();
    vtkNew<vtkCallbackCommand> clickCallback;
    clickCallback->SetClientData(this);
    clickCallback->SetCallback([](vtkObject* caller, unsigned long, void* clientdata, void*){
        auto* iren = static_cast<vtkRenderWindowInteractor*>(caller);
        auto* self = static_cast<Viewport*>(clientdata);
        
        int* pos = iren->GetEventPosition();
        self->m_Picker->Pick(pos[0], pos[1], 0, self->m_Renderer);
        vtkProp* picked = self->m_Picker->GetViewProp();
        
        Puppet* target = nullptr;
        if (picked) {
            for (auto* p : self->m_Puppets) {
                if (p->GetProp() == picked) {
                    target = p;
                    break;
                }
                auto* propAssembly = vtkPropAssembly::SafeDownCast(p->GetProp());
                auto* actorAssembly = vtkAssembly::SafeDownCast(p->GetProp());
                vtkPropCollection* parts = nullptr;
                if (propAssembly) parts = propAssembly->GetParts();
                else if (actorAssembly) parts = actorAssembly->GetParts();

                if (parts) {
                    bool found = false;
                    parts->InitTraversal();
                    for (int i=0; i<parts->GetNumberOfItems(); ++i) {
                        if (parts->GetNextProp() == picked) {
                            found = true;
                            break;
                        }
                    }
                    if (found) {
                        target = p;
                        break;
                    }
                }
            }
        }
        self->SelectPuppet(target);
    });
    iren->AddObserver(vtkCommand::LeftButtonPressEvent, clickCallback);

    // Keyboard events for widget coordinate frame
    m_KeyCallback = vtkSmartPointer<vtkCallbackCommand>::New();
    m_KeyCallback->SetClientData(this);
    m_KeyCallback->SetCallback([](vtkObject* caller, unsigned long event, void* clientdata, void*){
        auto* iren = static_cast<vtkRenderWindowInteractor*>(caller);
        auto* self = static_cast<Viewport*>(clientdata);
        
        std::string key = iren->GetKeySym();
        bool handled = false;

        if (self->m_HandlerWidget && self->m_HandlerWidget->GetEnabled()) {
            if (key == "l") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetReferenceFrame(vtkHandlerWidget::LOCAL);
                    std::cout << "Widget Frame: LOCAL" << std::endl;
                }
                handled = true;
            }
            else if (key == "g") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetReferenceFrame(vtkHandlerWidget::GLOBAL);
                    std::cout << "Widget Frame: GLOBAL" << std::endl;
                }
                handled = true;
            }
            else if (key == "c") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetReferenceFrame(vtkHandlerWidget::CENTER);
                    std::cout << "Widget Frame: CENTER" << std::endl;
                }
                handled = true;
            }
            else if (key == "k") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetReferenceFrame(vtkHandlerWidget::CENTER_LOCAL);
                    std::cout << "Widget Frame: CENTER_LOCAL" << std::endl;
                }
                handled = true;
            }
            else if (key == "1") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetTranslationEnabled(!self->m_HandlerWidget->GetTranslationEnabled());
                }
                handled = true;
            }
            else if (key == "2") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetRotationEnabled(!self->m_HandlerWidget->GetRotationEnabled());
                }
                handled = true;
            }
            else if (key == "3") {
                if (event == vtkCommand::KeyPressEvent) {
                    self->m_HandlerWidget->SetScalingEnabled(!self->m_HandlerWidget->GetScalingEnabled());
                }
                handled = true;
            }
        }

        if (key == "f") {
            if (event == vtkCommand::KeyPressEvent) {
                self->ZoomSelected();
            }
            handled = true;
        }

        if (handled) {
            self->m_KeyCallback->SetAbortFlag(1);
            iren->Render();
        }
    });
    iren->AddObserver(vtkCommand::KeyPressEvent, m_KeyCallback, 1.0);
    iren->AddObserver(vtkCommand::CharEvent, m_KeyCallback, 1.0);
}

void Viewport::Reset()
{
    ZoomAuto();
    Render();
}

void Viewport::ZoomAuto()
{
    if (m_Renderer) {
        m_Renderer->ResetCameraClippingRange();
        m_Renderer->ResetCamera();
    }
}

void Viewport::ZoomSelected()
{
    if (!m_Renderer) return;

    Puppet* selected = nullptr;
    for (auto* p : m_Puppets) {
        if (p->IsSelected()) {
            selected = p;
            break;
        }
    }
    if (!selected) return;

    vtkProp* prop = selected->GetProp();
    if (!prop) return;

    double* b = prop->GetBounds();
    if (!b) return;
    
    double bounds[6];
    std::copy(b, b + 6, bounds);
    
    if (bounds[0] > bounds[1]) return; // Invalid bounds

    // Expand bounds by a factor from center (e.g. 2.0 to have some margin)
    double center[3] = {(bounds[0] + bounds[1]) / 2.0, (bounds[2] + bounds[3]) / 2.0, (bounds[4] + bounds[5]) / 2.0};
    double h_ext[3] =  {(bounds[1] - bounds[0]) / 2.0, (bounds[3] - bounds[2]) / 2.0, (bounds[5] - bounds[4]) / 2.0};
    
    // Ensure a minimum size to avoid camera issues with flat/point objects
    double max_h = std::max({h_ext[0], h_ext[1], h_ext[2], 0.1});
    
    double newBounds[6];
    for (int i=0; i<3; ++i) {
        double current_h = std::max(h_ext[i], max_h * 0.1); 
        newBounds[2*i]   = center[i] - 2.5 * current_h;
        newBounds[2*i+1] = center[i] + 2.5 * current_h;
    }
    
    m_Renderer->ResetCamera(newBounds);
    m_Renderer->ResetCameraClippingRange();
    this->Render();
}

void Viewport::AddPuppet(Puppet& prop)
{
    m_Puppets.push_back(&prop);
    prop.ConnectRenderer(m_Renderer);
    Render();
}

void Viewport::RemovePuppet(Puppet& prop)
{
    if (prop.IsSelected()) SelectPuppet(nullptr);
    auto it = std::find(m_Puppets.begin(), m_Puppets.end(), &prop);
    if (it != m_Puppets.end()) m_Puppets.erase(it);
    prop.DisconnectRenderer(m_Renderer);
    Render();
}

void Viewport::SelectPuppet(Puppet* prop)
{
    for (auto* p : m_Puppets) {
        p->SetSelected(p == prop);
    }
    
    if (m_HandlerWidget) {
        if (prop) {
            vtkProp3D* prop3d = vtkProp3D::SafeDownCast(prop->GetProp());
            if (prop3d) {
                m_HandlerWidget->SetProp3D(prop3d);
                m_HandlerWidget->SetEnabled(1);
                m_HandlerWidget->PlaceWidget(prop3d->GetBounds());
            }
        } else {
            m_HandlerWidget->SetEnabled(0);
            m_HandlerWidget->SetProp3D(nullptr);
        }
    }
    
    Render();
}

void Viewport::SetGridVisible(bool visible)
{
    if (m_GridActor) {
        m_GridActor->SetVisibility(visible);
        Render();
    }
}

bool Viewport::GetGridVisible() const
{
    if (m_GridActor) {
        return m_GridActor->GetVisibility() != 0;
    }
    return false;
}

void Viewport::SetGridAxis(Axis axis)
{
    m_GridAxis = axis;
    UpdateGrid();
    Render();
}

void Viewport::addProp(vtkProp* prop)
{
    if (m_Renderer) {
        m_Renderer->AddActor(prop);
        Render();
    }
}

void Viewport::RemoveProp(vtkProp* prop)
{
    if (m_Renderer) {
        m_Renderer->RemoveViewProp(prop);
        Render();
    }
}

void Viewport::UpdateGrid()
{
    if (!m_Renderer || !m_GridSource) return;
    if (m_GridActor && !m_GridActor->GetVisibility()) return;

    vtkCamera* camera = m_Renderer->GetActiveCamera();
    if (!camera) return;

    // Determine the "scale" of the view (how many units are visible vertically)
    double viewHeight;
    if (camera->GetParallelProjection()) {
        viewHeight = 2.0 * camera->GetParallelScale();
    } else {
        double distance = camera->GetDistance();
        // ViewAngle is height angle in degrees
        double angleRad = camera->GetViewAngle() * vtkMath::Pi() / 180.0;
        viewHeight = 2.0 * distance * std::tan(angleRad / 2.0);
    }

    if (viewHeight <= 0) viewHeight = 1.0;

    // We want roughly 5-15 grid divisions visible.
    // Spacing should be a power of 10 (1mm, 1cm, 10cm, 1m, 10m...)
    double log10Spacing = std::floor(std::log10(viewHeight / 5.0));
    double spacing = std::pow(10.0, log10Spacing);

    // Get current focal point to center the grid near what we're looking at
    double focalPoint[3];
    camera->GetFocalPoint(focalPoint);

    // Indices for the two dimensions of the grid plane
    int idxH, idxV, idxN;
    if (m_GridAxis == X) { idxH = 1; idxV = 2; idxN = 0; }
    else if (m_GridAxis == Y) { idxH = 0; idxV = 2; idxN = 1; }
    else { idxH = 0; idxV = 1; idxN = 2; }

    // Align center to spacing
    double centerH = std::round(focalPoint[idxH] / spacing) * spacing;
    double centerV = std::round(focalPoint[idxV] / spacing) * spacing;
    double centerN = 0.0; // Grid plane typically passes through the origin

    // Number of lines
    int numLines = 20;
    double halfSize = (numLines / 2.0) * spacing;

    double minH = centerH - halfSize;
    double maxH = centerH + halfSize;
    double minV = centerV - halfSize;
    double maxV = centerV + halfSize;

    // Update Plane Source mapping axes to origin/point1/point2
    double origin[3] = {0,0,0}, p1[3] = {0,0,0}, p2[3] = {0,0,0};
    origin[idxH] = minH; origin[idxV] = minV; origin[idxN] = centerN;
    p1[idxH]     = maxH; p1[idxV]     = minV; p1[idxN]     = centerN;
    p2[idxH]     = minH; p2[idxV]     = maxV; p2[idxN]     = centerN;

    m_GridSource->SetOrigin(origin);
    m_GridSource->SetPoint1(p1);
    m_GridSource->SetPoint2(p2);
    m_GridSource->SetXResolution(numLines);
    m_GridSource->SetYResolution(numLines);
    m_GridSource->Update();

    if (m_OriginAxes) {
        m_OriginAxes->SetScaleFactor(spacing);
    }

    // Update annotation for grid size
    char gridLabel[32];
    if (spacing >= 1000.0) {
        sprintf(gridLabel, "Grid: %.0f m", spacing / 1000.0);
    } else if (spacing >= 10.0) {
        sprintf(gridLabel, "Grid: %.0f cm", spacing / 10.0);
    } else {
        sprintf(gridLabel, "Grid: %.0f mm", spacing);
    }
    m_Annotation->SetText(1, gridLabel);
}




} // namespace Vtk
} // namespace uLib