import os
import sys
import time
import uuid

from jcclang.core.const import *
from jcclang.core.logger import jcwLogger

_forced_platform = None
captured = {}


def capture_and_clean_platform_args():
    platform_keys = {
        f"--{OUTPUT}",
        f"--{DATASET_INPUT}",
        f"--{CODE_INPUT}",
        f"--{MODEL_INPUT}",
        f"--{CURRENT_PLATFORM}",
        f"--{GATEWAY_ADDR}",
        f"--{USER_ID}",
        f"--{JCS_ADDR}",
        f"--{JCS_AK}",
        f"--{JCS_SK}",
        f"--{JOBHUB_ADDR}",
        f"--{JOBHUB_PORT}",
        f"--{JOBHUB_JOBSET_ID}",
        f"--{JOBHUB_JOB_ID}",
        f"--{JOBHUB_SECRET}"
    }

    # captured = {}
    cleaned_argv = [sys.argv[0]]  # 保留脚本名称
    i = 1

    while i < len(sys.argv):
        arg = sys.argv[i]

        # jcwLogger.debug(f"capture arg: {arg}")

        # 情况1: --key=value
        if any(arg.startswith(k + "=") for k in platform_keys):
            key, value = arg.split("=", 1)
            captured[key] = value
            i += 1
            continue

        # 情况2: --key value（下一个参数是值）
        if arg in platform_keys:
            if i + 1 < len(sys.argv):
                key = arg
                value = sys.argv[i + 1]
                captured[key] = value
                i += 2  # 跳过 key 和 value
                continue
            else:
                # 没有值？按空字符串处理（或跳过）
                captured[arg] = ""
                i += 1
                continue

        # 不是平台参数，保留
        cleaned_argv.append(arg)
        i += 1

    # 替换 sys.argv，移除所有平台参数
    sys.argv = cleaned_argv


def get_args(key: str):
    value = captured.get(f"--{key}")
    # jcwLogger.debug(f"get args: {key} = {value}")
    return value


def default_context():
    return {
        "trace_id": str(uuid.uuid4()),
        "timestamp": int(time.time()),
        "logger": None  # 可注入 get_logger("task-name")
    }


def set_forced_platform(name: str):
    global _forced_platform
    _forced_platform = name


def get_platform():
    if _forced_platform:
        jcwLogger.debug(f"forced platform: {_forced_platform}")
        return _forced_platform
    if get_args(CURRENT_PLATFORM) is not None:
        jcwLogger.debug(f"get platform args: {get_args(CURRENT_PLATFORM)}")
        return get_args(CURRENT_PLATFORM)
    return os.environ.get(CURRENT_PLATFORM, "local")


def get_gateway_addr():
    if get_args(GATEWAY_ADDR) is not None:
        return get_args(GATEWAY_ADDR)
    return os.environ.get(GATEWAY_ADDR, "")


def get_jcs_addr():
    if get_args(JCS_ADDR) is not None:
        return get_args(JCS_ADDR)
    return os.environ.get(JCS_ADDR, "")


def get_jcs_ak():
    if get_args(JCS_AK) is not None:
        return get_args(JCS_AK)
    return os.environ.get(JCS_AK, "")


def get_jcs_sk():
    if get_args(JCS_SK) is not None:
        return get_args(JCS_SK)
    return os.environ.get(JCS_SK, "")


def get_user_id():
    if get_args(USER_ID) is not None:
        return get_args(USER_ID)
    env_val = os.environ.get(USER_ID)
    return int(env_val) if env_val and env_val.isdigit() else 0


def get_jobhub_addr():
    if get_args(JOBHUB_ADDR) is not None:
        return get_args(JOBHUB_ADDR)
    return os.environ.get(JOBHUB_ADDR, "")


def get_jobhub_port():
    port = get_args(JOBHUB_PORT)
    if port is not None:
        return int(port) if port.isdigit() else 0
    env_val = os.environ.get(JOBHUB_PORT)
    return int(env_val) if env_val and env_val.isdigit() else 0


def get_jobhub_jobset_id():
    if get_args(JOBHUB_JOBSET_ID) is not None:
        return get_args(JOBHUB_JOBSET_ID)
    return os.environ.get(JOBHUB_JOBSET_ID, "")


def get_jobhub_job_id():
    if get_args(JOBHUB_JOB_ID) is not None:
        return get_args(JOBHUB_JOB_ID)
    return os.environ.get(JOBHUB_JOB_ID, "")


def get_jobhub_secret():
    secret = get_args(JOBHUB_SECRET)
    if secret is not None:
        return secret.encode("utf-8")
    env_secret = os.environ.get(JOBHUB_SECRET, "")
    return env_secret.encode("utf-8")


def get_data_output():
    if get_args(OUTPUT) is not None:
        return get_args(OUTPUT)
    return os.environ.get(OUTPUT, "")


def get_dataset_input():
    if get_args(DATASET_INPUT) is not None:
        return get_args(DATASET_INPUT)
    return os.environ.get(DATASET_INPUT, "")


def get_code_input():
    if get_args(CODE_INPUT) is not None:
        return get_args(CODE_INPUT)
    return os.environ.get(CODE_INPUT, "")


def get_model_input():
    if get_args(MODEL_INPUT) is not None:
        return get_args(MODEL_INPUT)
    return os.environ.get(MODEL_INPUT, "")
