import os
import allure
import pytest
import subprocess
from common import *

g_test_suites = {"AllGather": {}, "AllReduce": {}, "AllToAll": {}, 
                 "AllToAllv": {}, "Broadcast": {}, "Gather": {}, 
                 "GroupCall": {}, "NonBlocking": {}, "Performance": {}, 
                 "ReduceScatter": {}, "Reduce": {}, "Scatter": {}, 
                 "SendRecv": {}, "Standalone": {}}

def run_case(case_cfg):
    if not os.path.exists(g_log_path):
        try:
            result = subprocess.run(["mkdir", "-p", g_log_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            # 打印标准输出和标准错误
            print(result.stdout.decode())
            print(result.stderr.decode())
        except subprocess.CalledProcessError as e:
            print(f"An error occurred while trying to create directory: {e}")
        except subprocess.TimeoutExpired as e:
            print(f"The command timed out: {e}")
    run_path="{}/{}".format(g_main_path, case_cfg["run_dir"])
    log_path = "{}/{}.txt".format(g_log_path, case_cfg["case"])
    run_cmd = "source /opt/{}/env.sh; cd {}; ./{} {} 1> {} 2>&1".format(g_dtk_version, run_path, case_cfg["run_cmd"], 
                                                                    case_cfg["run_args"], log_path)
    print(run_cmd)
    process_handle = subprocess.Popen(run_cmd, shell=True, preexec_fn=os.setpgrp)
    result = tear_down(process_handle, case_cfg["timeout"])

    return analysis_result(log_path) and result, log_path

gen_test_suites_from_csv("{}/rccl_test_cfg.csv".format(g_cur_path), g_test_suites)

@allure.feature("rccl")
@allure.story("AllGather")
@pytest.mark.AllGather
@pytest.mark.parametrize("test_case", g_test_suites["AllGather"].keys())
def test_rccl_AllGather(test_case):
    result, log_path = run_case(g_test_suites["hysmi_test"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("AllGather")
@pytest.mark.AllGather
@pytest.mark.parametrize("test_case", g_test_suites["AllGather"].keys())
def test_rccl_AllGather(test_case):
    result, log_path = run_case(g_test_suites["AllGather"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("AllReduce")
@pytest.mark.AllReduce
@pytest.mark.parametrize("test_case", g_test_suites["AllReduce"].keys())
def test_rccl_AllReduce(test_case):
    result, log_path = run_case(g_test_suites["AllReduce"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("AllToAll")
@pytest.mark.AllToAll
@pytest.mark.parametrize("test_case", g_test_suites["AllToAll"].keys())
def test_rccl_AllToAll(test_case):
    result, log_path = run_case(g_test_suites["AllToAll"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("AllToAllv")
@pytest.mark.AllToAllv
@pytest.mark.parametrize("test_case", g_test_suites["AllToAllv"].keys())
def test_rccl_AllToAllv(test_case):
    result, log_path = run_case(g_test_suites["AllToAllv"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("Broadcast")
@pytest.mark.Broadcast
@pytest.mark.parametrize("test_case", g_test_suites["Broadcast"].keys())
def test_rccl_Broadcast(test_case):
    result, log_path = run_case(g_test_suites["Broadcast"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("Gather")
@pytest.mark.Gather
@pytest.mark.parametrize("test_case", g_test_suites["Gather"].keys())
def test_rccl_Gather(test_case):
    result, log_path = run_case(g_test_suites["Gather"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("GroupCall")
@pytest.mark.GroupCall
@pytest.mark.parametrize("test_case", g_test_suites["GroupCall"].keys())
def test_rccl_GroupCall(test_case):
    result, log_path = run_case(g_test_suites["GroupCall"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("NonBlocking")
@pytest.mark.NonBlocking
@pytest.mark.parametrize("test_case", g_test_suites["NonBlocking"].keys())
def test_rccl_NonBlocking(test_case):
    result, log_path = run_case(g_test_suites["NonBlocking"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("Performance")
@pytest.mark.Performance
@pytest.mark.parametrize("test_case", g_test_suites["Performance"].keys())
def test_rccl_Performance(test_case):
    result, log_path = run_case(g_test_suites["Performance"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"

@allure.feature("rccl")
@allure.story("ReduceScatter")
@pytest.mark.ReduceScatter
@pytest.mark.parametrize("test_case", g_test_suites["ReduceScatter"].keys())
def test_rccl_ReduceScatter(test_case):
    result, log_path = run_case(g_test_suites["ReduceScatter"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("Reduce")
@pytest.mark.Reduce
@pytest.mark.parametrize("test_case", g_test_suites["Reduce"].keys())
def test_rccl_Reduce(test_case):
    result, log_path = run_case(g_test_suites["Reduce"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("Scatter")
@pytest.mark.Scatter
@pytest.mark.parametrize("test_case", g_test_suites["Scatter"].keys())
def test_rccl_Scatter(test_case):
    result, log_path = run_case(g_test_suites["Scatter"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("SendRecv")
@pytest.mark.SendRecv
@pytest.mark.parametrize("test_case", g_test_suites["SendRecv"].keys())
def test_rccl_SendRecv(test_case):
    result, log_path = run_case(g_test_suites["SendRecv"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"


@allure.feature("rccl")
@allure.story("Standalone")
@pytest.mark.Standalone
@pytest.mark.parametrize("test_case", g_test_suites["Standalone"].keys())
def test_rccl_Standalone(test_case):
    result, log_path = run_case(g_test_suites["Standalone"][test_case])
    allure.attach.file(log_path, name="test_log", attachment_type=allure.attachment_type.TEXT)
    assert result, "run fail!"