"""任务执行模块 - 在本地执行推理/计算任务"""
import json
import time
import logging

logger = logging.getLogger("kogun.executor")


async def execute_task(task: dict, upstream_config: dict = None) -> dict:
    """
    执行一个计算任务

    Args:
        task: 任务字典 {id, type, model, prompt, ...}
        upstream_config: 上游API配置 {url, key, models}

    Returns:
        {"result": ..., "actual_cost": ..., "duration_seconds": ...}
    """
    task_type = task.get("type", "inference")
    model = task.get("model", "")
    prompt = task.get("prompt", "")

    start_time = time.time()

    if task_type == "inference" or task_type == "chat":
        return await _do_inference(model, prompt, upstream_config)
    elif task_type == "compute":
        return _do_compute(prompt)
    elif task_type == "benchmark":
        return _do_benchmark()
    else:
        return _do_generic(task)


async def _do_inference(model: str, prompt: str, upstream: dict) -> dict:
    """执行推理任务 - 转发到上游API或模拟"""
    import httpx

    duration = 0

    if upstream and upstream.get("url") and upstream.get("key"):
        try:
            # 转发到上游API
            upstream_url = upstream["url"].rstrip("/")
            upstream_key = upstream["key"]
            upstream_model = upstream.get("models", {}).get(model, model)

            headers = {
                "Authorization": f"Bearer {upstream_key}",
                "Content-Type": "application/json",
            }
            payload = {
                "model": upstream_model,
                "messages": [
                    {"role": "user", "content": prompt or "hello"}
                ],
                "max_tokens": 512,
            }

            start = time.time()
            async with httpx.AsyncClient(timeout=120) as client:
                resp = await client.post(
                    f"{upstream_url}/v1/chat/completions",
                    headers=headers, json=payload,
                )
                duration = time.time() - start

            if resp.status_code == 200:
                data = resp.json()
                result_text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                return {
                    "result": result_text,
                    "actual_cost": round(duration * 0.0005, 6),  # 按时间计费
                    "duration_seconds": round(duration, 2),
                }
        except Exception as e:
            logger.warning(f"上游API调用失败: {e}")

    # 模拟执行（超轻量级）
    start = time.time()

    # 模拟计算 - 假装在运行模型
    simulated_time = max(0.5, len(prompt or "") * 0.01)
    simulated_time = min(simulated_time, 10)  # 最多模拟10秒

    # 用CPU做点实际计算来模拟（而不是sleep）
    _burn_cpu(simulated_time * 0.5)
    remaining = simulated_time - 0.5
    if remaining > 0:
        import asyncio
        await asyncio.sleep(remaining)

    duration = time.time() - start

    result_text = f"[模拟推理完成] 模型: {model} | 输入长度: {len(prompt or '')}字符 | 耗时: {duration:.2f}s"

    return {
        "result": result_text,
        "actual_cost": round(duration * 0.0002, 6),  # 模拟执行更便宜
        "duration_seconds": round(duration, 2),
    }


def _do_compute(prompt: str) -> dict:
    """执行通用计算任务"""
    start = time.time()

    params = {}
    try:
        if prompt:
            params = json.loads(prompt) if isinstance(prompt, str) else prompt
    except (json.JSONDecodeError, TypeError):
        params = {"input": str(prompt)[:200]}

    iterations = params.get("iterations", 1000000)
    workload = params.get("workload", "hash")

    result = {}
    if workload == "hash":
        import hashlib
        data = params.get("data", "benchmark_data")
        result["hash"] = hashlib.sha256(data.encode()).hexdigest()

        # 做多次哈希来模拟工作量
        for i in range(min(iterations, 500000)):
            data = hashlib.sha256((data + str(i)).encode()).hexdigest()
        result["final_hash"] = data
        result["iterations"] = min(iterations, 500000)

    elif workload == "pi":
        # 计算Pi的近似值
        n = min(iterations, 100000)
        pi = 0
        for i in range(n):
            pi += ((-1) ** i) / (2 * i + 1)
        result["pi_approx"] = pi * 4
        result["iterations"] = n
    else:
        result["message"] = f"未知工作负载: {workload}"

    duration = time.time() - start

    return {
        "result": json.dumps(result, ensure_ascii=False),
        "actual_cost": round(duration * 0.0001, 6),
        "duration_seconds": round(duration, 2),
    }


def _do_benchmark() -> dict:
    """执行基准测试"""
    import hashlib
    import time

    start = time.time()

    # CPU哈希基准
    hash_ops = 0
    end_time = time.time() + 3  # 跑3秒
    data = "benchmark"
    while time.time() < end_time:
        data = hashlib.sha256(data.encode()).hexdigest()
        hash_ops += 1

    duration = time.time() - start
    hashes_per_sec = hash_ops / duration

    return {
        "result": json.dumps({
            "benchmark_type": "cpu_hash",
            "hashes_per_sec": round(hashes_per_sec, 0),
            "total_hashes": hash_ops,
            "duration_seconds": round(duration, 2),
        }, ensure_ascii=False),
        "actual_cost": round(duration * 0.0001, 6),
        "duration_seconds": round(duration, 2),
    }


def _do_generic(task: dict) -> dict:
    """通用任务处理"""
    import time
    start = time.time()
    # 简单模拟
    simulated_time = min(2, len(str(task.get("prompt", ""))) * 0.01)
    _burn_cpu(simulated_time)

    duration = time.time() - start
    return {
        "result": f"[完成] 任务#{task.get('id')} 类型: {task.get('type', 'unknown')}",
        "actual_cost": round(duration * 0.0001, 6),
        "duration_seconds": round(duration, 2),
    }


def _burn_cpu(seconds: float):
    """用CPU计算来模拟负载，而不是sleep"""
    import time
    import hashlib
    end = time.time() + max(0.1, seconds)
    i = 0
    data = "x"
    while time.time() < end:
        data = hashlib.sha256((data + str(i)).encode()).hexdigest()
        i += 1
        if i > 500000:  # 防止无限循环
            break
