/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#pragma once

#include <iostream>
#include <sstream>
#include <streambuf>
#include <string>
#include <fstream>
#include <fcntl.h>

#ifdef _WIN32
#include <windows.h>
#include <io.h>
#include <stdio.h>
#include <stdlib.h>
#define dup _dup
#define dup2 _dup2
#define fd_close _close
#define unlink _unlink
#define STDERR_FD _fileno(stderr)
#define OPEN_FLAGS (_O_WRONLY | _O_CREAT | _O_TRUNC)
#define OPEN_MODE (_S_IREAD | _S_IWRITE)
#define open _open
#else
#include <unistd.h>
#define fd_close close
#define STDERR_FD STDERR_FILENO
#define OPEN_FLAGS (O_WRONLY | O_CREAT | O_TRUNC)
#define OPEN_MODE 0644
#endif

// Class to capture all stderr output (HIP logging uses stderr)
class OutCapture {
private:
    std::stringstream captured_stream_;
    std::streambuf* cerr_backup_;
    int stderr_backup_;
    std::string temp_file_;

    static std::string getTempFilePath() {
#ifdef _WIN32
        char temp_path[MAX_PATH];
        if (GetTempPathA(MAX_PATH, temp_path)) {
            return std::string(temp_path) + "hip_stderr_capture.txt";
        }
        // Fallback to current directory
        return "hip_stderr_capture.txt";
#else
        return "/tmp/hip_stderr_capture.txt";
#endif
    }

public:
    OutCapture() : temp_file_(getTempFilePath()) {
        // Backup original cerr stream buffer (HIP logging uses stderr)
        cerr_backup_ = std::cerr.rdbuf();

        // Backup original stderr file descriptor
        stderr_backup_ = dup(STDERR_FD);
    }

    void startCapture() {
        // Clear any previous content
        captured_stream_.str("");
        captured_stream_.clear();

        // Redirect std::cerr to our stringstream
        std::cerr.rdbuf(captured_stream_.rdbuf());

        // Redirect stderr file descriptor to temp file (for fprintf to stderr)
        int temp_fd = open(temp_file_.c_str(), OPEN_FLAGS, OPEN_MODE);
        if (temp_fd != -1) {
            dup2(temp_fd, STDERR_FD);
            fd_close(temp_fd);
        }
    }

    std::string stopCapture() {
        // Restore original cerr stream
        std::cerr.rdbuf(cerr_backup_);

        // Restore original stderr file descriptor
        dup2(stderr_backup_, STDERR_FD);

        // Read from temp file (captures fprintf(stderr) output from HIP logging)
        std::ifstream temp_file(temp_file_);
        std::string file_content;
        if (temp_file.is_open()) {
            std::string line;
            while (std::getline(temp_file, line)) {
                file_content += line + "\n";
            }
            temp_file.close();
        }

        // Combine both captures: C++ streams and file descriptor output
        std::string stream_content = captured_stream_.str();
        std::string total_output = stream_content + file_content;

        // Clean up temp file
        unlink(temp_file_.c_str());

        return total_output;
    }

    ~OutCapture() {
        // Ensure everything is restored
        std::cerr.rdbuf(cerr_backup_);
        dup2(stderr_backup_, STDERR_FD);
        fd_close(stderr_backup_);
        unlink(temp_file_.c_str());
    }
};