#pragma once
#include <iostream>
#include <string>
#include <filesystem>
#include <TopoDS_Shape.hxx>
#include <string>

class GeometryIO {
public:
    // 加载几何模型：支持 .brep, .step/.stp, .iges/.igs
    static TopoDS_Shape Load(const std::string& filename);

    // 保存几何模型：支持 .brep, .step/.stp, .iges/.igs
    static bool Save(const TopoDS_Shape& shape, const std::string& filename);
};

#endif // GEOMETRY_IO_H

#include "GeometryIO.h"
#include <BRepTools.hxx>
#include <BRep_Builder.hxx>
#include <STEPControl_Reader.hxx>
#include <STEPControl_Writer.hxx>
#include <IGESControl_Reader.hxx>
#include <IGESControl_Writer.hxx>
#include <Interface_Static.hxx>
#include <filesystem>
#include <iostream>

TopoDS_Shape GeometryIO::Load(const std::string& filename) {
    std::string ext = std::filesystem::path(filename).extension().string();

    if (ext == ".brep") {
        TopoDS_Shape shape;
        BRep_Builder builder;
        if (!BRepTools::Read(shape, filename.c_str(), builder)) {
            std::cerr << "Failed to load BREP file: " << filename << std::endl;
        }
        return shape;
    }
    else if (ext == ".step" || ext == ".stp") {
        STEPControl_Reader reader;
        IFSelect_ReturnStatus stat = reader.ReadFile(filename.c_str());
        if (stat == IFSelect_RetDone) {
            reader.TransferRoots();
            return reader.OneShape();
        }
        else {
            std::cerr << "Failed to read STEP file: " << filename << std::endl;
        }
    }
    else if (ext == ".iges" || ext == ".igs") {
        IGESControl_Reader reader;
        IFSelect_ReturnStatus stat = reader.ReadFile(filename.c_str());
        if (stat == IFSelect_RetDone) {
            reader.TransferRoots();
            return reader.OneShape();
        }
        else {
            std::cerr << "Failed to read IGES file: " << filename << std::endl;
        }
    }

    std::cerr << "Unsupported file format: " << ext << std::endl;
    return TopoDS_Shape();
}

bool GeometryIO::Save(const TopoDS_Shape& shape, const std::string& filename) {
    std::string ext = std::filesystem::path(filename).extension().string();

    if (ext == ".brep") {
        return BRepTools::Write(shape, filename.c_str());
    }
    else if (ext == ".step" || ext == ".stp") {
        STEPControl_Writer writer;
        writer.Transfer(shape, STEPControl_AsIs);
        IFSelect_ReturnStatus stat = writer.Write(filename.c_str());
        return stat == IFSelect_RetDone;
    }
    else if (ext == ".iges" || ext == ".igs") {
        IGESControl_Writer writer;
        writer.AddShape(shape);
        writer.ComputeModel();
        return writer.Write(filename.c_str());
    }

    std::cerr << "Unsupported file format for saving: " << ext << std::endl;
    return false;
}

// ShapeViewer.hpp
#pragma once

#include <TopoDS_Shape.hxx>
#include <vtkSmartPointer.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkActor.h>
#include <vtkPolyDataMapper.h>
#include <vtkInteractorStyleTrackballCamera.h>

class ShapeViewer {
public:
    ShapeViewer();
    void SetShape(const TopoDS_Shape& shape);
    void Show();

private:
    vtkSmartPointer<vtkRenderer> renderer;
    vtkSmartPointer<vtkRenderWindow> renderWindow;
    vtkSmartPointer<vtkRenderWindowInteractor> interactor;
    vtkSmartPointer<vtkInteractorStyleTrackballCamera> interactorStyle;
    vtkSmartPointer<vtkActor> actor;
    vtkSmartPointer<vtkPolyDataMapper> mapper;
};

// ShapeViewer.cpp

#include "ShapeViewer.h"
#include <IVtkTools_ShapeDataSource.hxx>
#include <IVtkOCC_Shape.hxx>

ShapeViewer::ShapeViewer() {
    renderer = vtkSmartPointer<vtkRenderer>::New();
    renderWindow = vtkSmartPointer<vtkRenderWindow>::New();
    interactor = vtkSmartPointer<vtkRenderWindowInteractor>::New();
    interactorStyle = vtkSmartPointer<vtkInteractorStyleTrackballCamera>::New();
    actor = vtkSmartPointer<vtkActor>::New();
    mapper = vtkSmartPointer<vtkPolyDataMapper>::New();

    renderWindow->AddRenderer(renderer);
    interactor->SetRenderWindow(renderWindow);
    interactor->SetInteractorStyle(interactorStyle);
}

void ShapeViewer::SetShape(const TopoDS_Shape& shape) {
    vtkNew<IVtkTools_ShapeDataSource> occSource;
    occSource->SetShape(new IVtkOCC_Shape(shape));
    mapper->SetInputConnection(occSource->GetOutputPort());
    actor->SetMapper(mapper);
    renderer->AddActor(actor);
}

void ShapeViewer::Show() {
    renderWindow->Render();
    interactor->Start();
}

// 加载并显示保存的 BoxShape
#include "ShapeViewer.h"
#include "ShapeBase.h"
#include "BoxShape.h"
#include "CylinderShape.h"
#include "ConeShape.h"
#include "CustomShape.h"

#include <memory>
#include <vtkAutoInit.h>
#include <iostream>

VTK_MODULE_INIT(vtkRenderingOpenGL2);
VTK_MODULE_INIT(vtkInteractionStyle);

int main() {
    // 创建并保存 BoxShape
    std::unique_ptr<ShapeBase> boxShape = std::make_unique<BoxShape>(1.0, 2.0, 3.0);
    if (boxShape->Save("box.brep")) {
        std::cout << "Box shape saved successfully." << std::endl;
    }
    else {
        std::cerr << "Failed to save Box shape." << std::endl;
    }

    // 创建并保存 CylinderShape
    std::unique_ptr<ShapeBase> cylinderShape = std::make_unique<CylinderShape>(1.0, 2.0);
    if (cylinderShape->Save("cylinder.brep")) {
        std::cout << "Cylinder shape saved successfully." << std::endl;
    }
    else {
        std::cerr << "Failed to save Cylinder shape." << std::endl;
    }

    // 创建并保存 ConeShape
    std::unique_ptr<ShapeBase> coneShape = std::make_unique<ConeShape>(1.0, 0.5, 2.0);
    if (coneShape->Save("cone.brep")) {
        std::cout << "Cone shape saved successfully." << std::endl;
    }
    else {
        std::cerr << "Failed to save Cone shape." << std::endl;
    }

    // 创建并保存 CustomShape
    std::unique_ptr<ShapeBase> customShape = std::make_unique<CustomShape>();
    if (customShape->Save("custom_shape.brep")) {
        std::cout << "Custom shape saved successfully." << std::endl;
    }
    else {
        std::cerr << "Failed to save Custom shape." << std::endl;
    }

    // 加载并显示保存的 BoxShape
    std::unique_ptr<ShapeBase> loadedShape = std::make_unique<BoxShape>();
    if (loadedShape->Load("box.brep")) {
        std::cout << "Box shape loaded successfully." << std::endl;
        ShapeViewer viewer;
        viewer.SetShape(loadedShape->GetShape());
        viewer.Show();
    }
    else {
        std::cerr << "Failed to load Box shape." << std::endl;
    }

    return 0;
}