#include "BarIntersectionVisualizer.h"
#include "PointCloudLoader.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <algorithm>
#include <cstdlib>
#include <cmath>

#ifdef _WIN32
#include <windows.h>
#else
#include <dirent.h>
#include <sys/stat.h>
#endif

void PrintUsage(const char* programName) {
    std::cout << "Usage: " << programName << " <point_cloud_file_or_directory> [options]" << std::endl;
    std::cout << "\nOptions:" << std::endl;
    std::cout << "  --interactive         Show interactive windows (default)" << std::endl;
    std::cout << "  --batch               Save screenshots instead of interactive" << std::endl;
    std::cout << "  --batch-dir           Process all .txt files in the specified directory" << std::endl;
    std::cout << "  --output-dir <dir>    Output directory for screenshots (default: ./visualization_output)" << std::endl;
    std::cout << "\nExample:" << std::endl;
    std::cout << "  " << programName << " data/test1/sample.txt --interactive" << std::endl;
    std::cout << "  " << programName << " data/test1 --batch-dir --output-dir ./results" << std::endl;
}

// Helper function to collect all .txt files from a directory
std::vector<std::string> CollectTxtFiles(const std::string& dirPath) {
    std::vector<std::string> txtFiles;
#ifdef _WIN32
    WIN32_FIND_DATAA findData;
    std::string searchPath = dirPath + "\\*.txt";
    HANDLE hFind = FindFirstFileA(searchPath.c_str(), &findData);
    if (hFind == INVALID_HANDLE_VALUE) {
        std::cerr << "Error: Cannot open directory: " << dirPath << std::endl;
        return txtFiles;
    }
    do {
        if (!(findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) {
            txtFiles.push_back(dirPath + "\\" + findData.cFileName);
        }
    } while (FindNextFileA(hFind, &findData) != 0);
    FindClose(hFind);
#else
    DIR* dir = opendir(dirPath.c_str());
    if (!dir) {
        std::cerr << "Error: Cannot open directory: " << dirPath << std::endl;
        return txtFiles;
    }
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string filename = entry->d_name;
        if (filename.length() > 4 && filename.substr(filename.length() - 4) == ".txt") {
            txtFiles.push_back(dirPath + "/" + filename);
        }
    }
    closedir(dir);
#endif
    std::sort(txtFiles.begin(), txtFiles.end());
    return txtFiles;
}

bool CreateDirectoryRecursive(const std::string& path) {
#ifdef _WIN32
    std::string normalizedPath = path;
    for (size_t i = 0; i < normalizedPath.length(); i++) {
        if (normalizedPath[i] == '/') normalizedPath[i] = '\\';
    }
    size_t pos = 0;
    do {
        pos = normalizedPath.find_first_of("\\/", pos + 1);
        std::string subPath = normalizedPath.substr(0, pos);
        if (!subPath.empty() && subPath.back() != ':') {
            CreateDirectoryA(subPath.c_str(), NULL);
        }
    } while (pos != std::string::npos);
    return true;
#else
    std::string cmd = "mkdir -p " + path;
    return system(cmd.c_str()) == 0;
#endif
}

std::string GetFileBaseName(const std::string& filePath) {
    size_t lastSlash = filePath.find_last_of("/\\");
    size_t lastDot = filePath.find_last_of('.');
    std::string filename;
    if (lastSlash != std::string::npos) {
        filename = filePath.substr(lastSlash + 1);
    } else {
        filename = filePath;
    }
    if (lastDot != std::string::npos && lastDot > lastSlash) {
        return filename.substr(0, lastDot - (lastSlash != std::string::npos ? lastSlash + 1 : 0));
    }
    return filename;
}

// ========== Debug Callbacks ==========

struct CallbackUserData {
    const SVzNLPointXYZ* points;
    int rows;
    int cols;
    BarIntersectionVisualizer* visualizer;
    bool enableVisualization;
};

void OnPlaneSegmented(
    const SVzNLPointXYZ* planePoints, int planePointCount,
    const SVzNL3DPointF* normal,
    const SVzNLPointXYZ* allPoints, int totalCount,
    void* userData
) {
    CallbackUserData* data = static_cast<CallbackUserData*>(userData);
    if (!data || !data->enableVisualization || !data->visualizer) return;

    std::cout << "[DEBUG] Plane segmented: " << planePointCount << " points" << std::endl;
    std::cout << "  Normal: (" << normal->x << ", " << normal->y << ", " << normal->z << ")" << std::endl;

    data->visualizer->VisualizePlaneSegmentation(
        planePoints, planePointCount, normal,
        allPoints, totalCount,
        "Step 1: Plane Segmentation"
    );
}

void OnPointsAligned(
    const SVzNLPointXYZ* alignedPoints, int count,
    void* userData
) {
    CallbackUserData* data = static_cast<CallbackUserData*>(userData);
    if (!data) return;

    std::cout << "[DEBUG] Points aligned: " << count << " points" << std::endl;

    if (!data->enableVisualization || !data->visualizer) return;

    data->visualizer->VisualizeAlignedPoints(
        alignedPoints, count,
        "Step 1b: Aligned Points"
    );
}

void OnPlaneFiltered(
    const SVzNLPointXYZ* filteredPoints, int count,
    void* userData
) {
    CallbackUserData* data = static_cast<CallbackUserData*>(userData);
    if (!data || !data->enableVisualization || !data->visualizer) return;

    std::cout << "[DEBUG] Plane filtered, remaining: " << count << " bar points" << std::endl;

    data->visualizer->VisualizeFilteredPoints(
        filteredPoints, count,
        "Step 2: Filtered Points (Bars Only)"
    );
}

void OnClustersDetected(
    const SGrowthCluster* clusters, int clusterCount,
    void* userData
) {
    CallbackUserData* data = static_cast<CallbackUserData*>(userData);
    if (!data) return;

    std::cout << "[DEBUG] Clusters detected: " << clusterCount << std::endl;
    for (int i = 0; i < clusterCount; i++) {
        std::cout << "  Cluster #" << i
                  << " pts=" << clusters[i].pointCount
                  << " centroid=(" << clusters[i].centroid.x
                  << "," << clusters[i].centroid.y
                  << "," << clusters[i].centroid.z << ")"
                  << " bbox=[(" << clusters[i].minBound.x << "," << clusters[i].minBound.y << "," << clusters[i].minBound.z << ")"
                  << " - (" << clusters[i].maxBound.x << "," << clusters[i].maxBound.y << "," << clusters[i].maxBound.z << ")]"
                  << std::endl;
    }
}

int ProcessSingleFile(
    const std::string& inputFile,
    bool interactive,
    const std::string& outputDir
) {
    std::cout << "\n" << std::string(70, '=') << std::endl;
    std::cout << "Processing: " << inputFile << std::endl;
    std::cout << std::string(70, '=') << std::endl;

    // Load point cloud
    SVzNLPointXYZ* points = nullptr;
    int rows = 0, cols = 0, errCode = 0;
    int ret = LoadTxtFile(inputFile.c_str(), &points, &rows, &cols, &errCode);
    if (ret != 0) {
        std::cerr << "Error loading point cloud: " << errCode << std::endl;
        return 1;
    }

    std::cout << "Loaded " << rows << " x " << cols << " = "
              << (rows * cols) << " points" << std::endl;

    // Create visualizer
    BarIntersectionVisualizer visualizer;
    visualizer.SetInteractive(interactive);
    if (!interactive) {
        std::string fileBaseName = GetFileBaseName(inputFile);
        std::string fileOutputDir = outputDir + "/" + fileBaseName;
        CreateDirectoryRecursive(fileOutputDir);
        visualizer.SetOutputDirectory(fileOutputDir);
    }

    // Setup callbacks
    CallbackUserData callbackData;
    callbackData.points = points;
    callbackData.rows = rows;
    callbackData.cols = cols;
    callbackData.visualizer = &visualizer;
    callbackData.enableVisualization = true;

    SBarIntersectionDebugCallbacks debugCallbacks;
    debugCallbacks.onPlaneSegmented = nullptr;// OnPlaneSegmented;
    debugCallbacks.onPointsAligned = nullptr;// OnPointsAligned;
    debugCallbacks.onPlaneFiltered = nullptr;// OnPlaneFiltered;
    debugCallbacks.onClustersDetected = nullptr;// OnClustersDetected;
    debugCallbacks.userData = &callbackData;

    // Set parameters
    SPlaneSegmentationParams planeParams;
    SGrowthParams growthParams;

    // Run detection
    SBarIntersectionResult result;
    ret = DetectBarIntersections(
        points, rows, cols,
        planeParams, growthParams,
        &result, &debugCallbacks
    );

    if (ret != 0) {
        std::cerr << "Detection failed with error: " << ret << std::endl;
        delete[] points;
        return 1;
    }

    // Print results
    PrintDetectionResults(result, true);

    // Final visualization
    visualizer.VisualizeResult(points, rows, cols, result, "Final Result");

    // Best cluster visualization (filtered aligned points)
    visualizer.VisualizeBestCluster(result, "Best Cluster (min Z, high pts)");

    // Cleanup
    FreeBarIntersectionResult(&result);
    delete[] points;

    std::cout << "Completed: " << inputFile << std::endl;
    return 0;
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        PrintUsage(argv[0]);
        return 1;
    }

    std::string firstArg = argv[1];
    if (firstArg == "--help" || firstArg == "-h") {
        PrintUsage(argv[0]);
        return 0;
    }

    std::string inputPath = argv[1];
    bool interactive = true;
    bool batchDir = false;
    std::string outputDir = "./visualization_output";

    for (int i = 2; i < argc; i++) {
        std::string arg = argv[i];
        if (arg == "--interactive") {
            interactive = true;
        } else if (arg == "--batch") {
            interactive = false;
        } else if (arg == "--batch-dir") {
            interactive = false;
            batchDir = true;
        } else if (arg == "--output-dir" && i + 1 < argc) {
            outputDir = argv[++i];
        } else {
            std::cerr << "Unknown option: " << arg << std::endl;
            PrintUsage(argv[0]);
            return 1;
        }
    }

    std::cout << "=== Bar Intersection Detection Visualization Demo ===" << std::endl;

    const char* ver = GetAlgorithmVersion();
    const char* name = GetAlgorithmName();
    std::cout << "Algorithm: " << (name ? name : "Unknown")
              << " v" << (ver ? ver : "Unknown") << std::endl;

    if (batchDir) {
        std::cout << "Mode: Batch Directory Processing" << std::endl;
        std::vector<std::string> txtFiles = CollectTxtFiles(inputPath);
        if (txtFiles.empty()) {
            std::cerr << "No .txt files found in: " << inputPath << std::endl;
            return 1;
        }
        std::cout << "Found " << txtFiles.size() << " file(s)" << std::endl;

        int successCount = 0, failCount = 0;
        for (const auto& file : txtFiles) {
            if (ProcessSingleFile(file, interactive, outputDir) == 0) {
                successCount++;
            } else {
                failCount++;
            }
        }

        std::cout << "\nBatch Summary: " << successCount << " success, "
                  << failCount << " failed" << std::endl;
    } else {
        std::cout << "Mode: " << (interactive ? "Interactive" : "Batch") << std::endl;
        ProcessSingleFile(inputPath, interactive, outputDir);
    }

    std::cout << "\nProgram completed successfully!" << std::endl;
    return 0;
}