import json
import re
import uuid
from copy import deepcopy

import pandas as pd
import requests

from .common import clamp, sample_lognormal


DEFAULT_PROFILES = [
    {
        "task_type": "simple_qa",
        "description": "单节点即可完成的简短问答，少量情况下调用工具。",
        "phase": "answering",
        "tool_probability": 0.10,
        "split_probability": 0.05,
        "mean_subtasks": 1.2,
        "timeout_probability": 0.02,
        "failure_probability": 0.01,
        "think_time_ms_mean": 420,
        "think_time_cv": 0.55,
        "request_size_bytes_mean": 1800,
        "response_size_bytes_mean": 2600,
        "message_size_cv": 0.45,
    },
    {
        "task_type": "tool_research",
        "description": "需要搜索、数据库或工具结果的研究任务。",
        "phase": "evidence_collection",
        "tool_probability": 0.78,
        "split_probability": 0.20,
        "mean_subtasks": 1.8,
        "timeout_probability": 0.07,
        "failure_probability": 0.025,
        "think_time_ms_mean": 900,
        "think_time_cv": 0.75,
        "request_size_bytes_mean": 3200,
        "response_size_bytes_mean": 8500,
        "message_size_cv": 0.70,
    },
    {
        "task_type": "collaborative_analysis",
        "description": "协调多个执行 Agent 并汇总结果的复杂分析任务。",
        "phase": "multi_agent_synthesis",
        "tool_probability": 0.48,
        "split_probability": 0.82,
        "mean_subtasks": 3.4,
        "timeout_probability": 0.09,
        "failure_probability": 0.035,
        "think_time_ms_mean": 1450,
        "think_time_cv": 0.85,
        "request_size_bytes_mean": 5200,
        "response_size_bytes_mean": 11800,
        "message_size_cv": 0.80,
    },
]


PROFILE_KEYS = {
    "task_type", "description", "phase", "tool_probability", "split_probability",
    "mean_subtasks", "timeout_probability", "failure_probability",
    "think_time_ms_mean", "think_time_cv", "request_size_bytes_mean",
    "response_size_bytes_mean", "message_size_cv",
}


def _extract_json(text):
    text = text.strip()
    fenced = re.search(r"```(?:json)?\s*(.*?)```", text, flags=re.S | re.I)
    if fenced:
        text = fenced.group(1).strip()
    start = min([i for i in (text.find("["), text.find("{")) if i >= 0] or [0])
    text = text[start:]
    if text.startswith("{"):
        obj = json.loads(text)
        return obj.get("profiles", obj)
    return json.loads(text)


def _validate_profiles(profiles, requested_types):
    if not isinstance(profiles, list):
        raise ValueError("LLM output must be a JSON list")
    validated = []
    by_type = {p["task_type"]: p for p in DEFAULT_PROFILES}
    for raw in profiles:
        if not isinstance(raw, dict) or "task_type" not in raw:
            continue
        base = deepcopy(by_type.get(raw["task_type"], DEFAULT_PROFILES[0]))
        for key in PROFILE_KEYS:
            if key in raw:
                base[key] = raw[key]
        for key in ("tool_probability", "split_probability", "timeout_probability", "failure_probability"):
            base[key] = clamp(float(base[key]), 0.0, 0.95)
        base["mean_subtasks"] = clamp(float(base["mean_subtasks"]), 1.0, 8.0)
        for key in ("think_time_ms_mean", "request_size_bytes_mean", "response_size_bytes_mean"):
            base[key] = max(float(base[key]), 1.0)
        for key in ("think_time_cv", "message_size_cv"):
            base[key] = clamp(float(base[key]), 0.05, 2.5)
        validated.append(base)
    found = {p["task_type"] for p in validated}
    for task_type in requested_types:
        if task_type not in found:
            validated.append(deepcopy(by_type.get(task_type, DEFAULT_PROFILES[0])))
    return [p for p in validated if p["task_type"] in requested_types]


def generate_profiles(mode, task_types, llm_config, allow_fallback=True):
    if mode == "local":
        return _validate_profiles(DEFAULT_PROFILES, task_types), {"generator": "local_builtin"}

    import os
    api_key = os.environ.get("OPENAI_API_KEY")
    base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1").rstrip("/")
    model = os.environ.get("OPENAI_MODEL", "")
    if not api_key or not model:
        if allow_fallback:
            return _validate_profiles(DEFAULT_PROFILES, task_types), {
                "generator": "local_fallback", "reason": "OPENAI_API_KEY or OPENAI_MODEL missing"
            }
        raise RuntimeError("LLM mode requires OPENAI_API_KEY and OPENAI_MODEL")

    prompt = f"""你是多智能体网络仿真实验的数据设计专家。请为以下任务类型生成行为画像：
{json.dumps(task_types, ensure_ascii=False)}

仅输出 JSON 数组，每个对象必须包含这些字段：
task_type, description, phase, tool_probability, split_probability, mean_subtasks,
timeout_probability, failure_probability, think_time_ms_mean, think_time_cv,
request_size_bytes_mean, response_size_bytes_mean, message_size_cv。

约束：概率在0到0.95之间；mean_subtasks在1到8之间；时间单位毫秒；消息大小单位字节；
不同任务类型应体现从简单问答、工具密集到多Agent协作的明显差异。画像将用于科学仿真，数值应合理且可解释。"""
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "Return valid JSON only. Do not include markdown commentary."},
            {"role": "user", "content": prompt},
        ],
        "temperature": float(llm_config.get("temperature", 0.5)),
        "max_tokens": int(llm_config.get("max_tokens", 3500)),
    }
    try:
        response = requests.post(
            base_url + "/chat/completions",
            headers={"Authorization": "Bearer " + api_key, "Content-Type": "application/json"},
            json=payload,
            timeout=float(llm_config.get("timeout_seconds", 90)),
        )
        response.raise_for_status()
        body = response.json()
        content = body["choices"][0]["message"]["content"]
        profiles = _validate_profiles(_extract_json(content), task_types)
        return profiles, {"generator": "llm", "model": model, "base_url": base_url}
    except Exception as exc:
        if not allow_fallback:
            raise
        return _validate_profiles(DEFAULT_PROFILES, task_types), {
            "generator": "local_fallback", "reason": str(exc), "requested_model": model
        }


def _event(task_id, msg_id, parent_id, timestamp, task_type, phase, source, destination,
           state_before, state_after, duration, message_type, message_size, queue_length,
           success=True, retry_count=0, is_external=False):
    return {
        "timestamp": round(float(timestamp), 6),
        "task_id": task_id,
        "message_id": msg_id,
        "parent_message_id": parent_id or "",
        "task_type": task_type,
        "task_phase": phase,
        "source": source,
        "destination": destination,
        "state_before": state_before,
        "state_after": state_after,
        "state_duration_ms": round(float(duration), 3),
        "message_type": message_type,
        "message_size_bytes": int(max(0, round(message_size))),
        "queue_length": int(max(0, queue_length)),
        "success": bool(success),
        "retry_count": int(retry_count),
        "is_external": bool(is_external),
    }


def generate_event_log(profiles, tasks_per_type, rng, network):
    events = []
    agent_count = int(network["agent_count"])
    tool_count = int(network["tool_count"])
    coordinator = network.get("coordinator", "agent_0")
    clock = 0.0
    for profile in profiles:
        for index in range(tasks_per_type):
            task_id = "%s_%05d" % (profile["task_type"], index)
            clock += float(rng.exponential(0.7))
            t = clock
            external_size = sample_lognormal(rng, profile["request_size_bytes_mean"] * 0.75,
                                              profile["message_size_cv"])
            root_id = uuid.uuid4().hex[:16]
            think = sample_lognormal(rng, profile["think_time_ms_mean"], profile["think_time_cv"])
            queue = int(rng.poisson(0.5 + 2.0 * profile["split_probability"]))
            events.append(_event(task_id, root_id, "", t, profile["task_type"], profile["phase"],
                                 "external", coordinator, "Idle", "Think", think,
                                 "external_task", external_size, queue, is_external=True))
            t += think / 1000.0
            parent = root_id
            split = rng.random() < profile["split_probability"]
            subtask_count = max(1, int(rng.poisson(max(profile["mean_subtasks"] - 1.0, 0.01)) + 1)) if split else 1
            subtask_count = min(subtask_count, 8)
            result_sizes = []
            all_success = True
            for sub_index in range(subtask_count):
                worker = "agent_%d" % (1 + ((index + sub_index) % max(1, agent_count - 1)))
                req_id = uuid.uuid4().hex[:16]
                req_size = sample_lognormal(rng, profile["request_size_bytes_mean"], profile["message_size_cv"])
                state = "Split" if split else "CallAgent"
                events.append(_event(task_id, req_id, parent, t, profile["task_type"], profile["phase"],
                                     coordinator, worker, "Think", state, 5.0, "agent_request",
                                     req_size, queue))
                worker_think = sample_lognormal(rng, profile["think_time_ms_mean"], profile["think_time_cv"])
                t += (8.0 + worker_think) / 1000.0
                tool_used = rng.random() < profile["tool_probability"]
                retries = 0
                success = True
                tool_response_size = 0.0
                if tool_used:
                    tool = "tool_%d" % ((index + sub_index) % max(1, tool_count))
                    tool_req_parent = req_id
                    while True:
                        tool_req_id = uuid.uuid4().hex[:16]
                        tool_req_size = sample_lognormal(rng, profile["request_size_bytes_mean"] * 0.35,
                                                         profile["message_size_cv"])
                        events.append(_event(task_id, tool_req_id, tool_req_parent, t,
                                             profile["task_type"], profile["phase"], worker, tool,
                                             "Think" if retries == 0 else "Retry", "CallTool", 8.0,
                                             "tool_request", tool_req_size, queue, retry_count=retries))
                        t += 0.008
                        timed_out = rng.random() < profile["timeout_probability"]
                        failed = rng.random() < profile["failure_probability"]
                        if not timed_out and not failed:
                            tool_response_size = sample_lognormal(
                                rng, profile["response_size_bytes_mean"] * 0.72,
                                profile["message_size_cv"])
                            response_id = uuid.uuid4().hex[:16]
                            tool_ms = sample_lognormal(rng, profile["think_time_ms_mean"] * 0.65,
                                                       profile["think_time_cv"])
                            t += tool_ms / 1000.0
                            events.append(_event(task_id, response_id, tool_req_id, t,
                                                 profile["task_type"], profile["phase"], tool, worker,
                                                 "CallTool", "Think", tool_ms, "tool_response",
                                                 tool_response_size, queue, retry_count=retries))
                            break
                        if retries >= 2:
                            success = False
                            fail_id = uuid.uuid4().hex[:16]
                            events.append(_event(task_id, fail_id, tool_req_id, t + 2.2,
                                                 profile["task_type"], profile["phase"], worker, coordinator,
                                                 "Wait", "Failed", 2200.0, "error", 600, queue,
                                                 success=False, retry_count=retries))
                            t += 2.2
                            break
                        retries += 1
                        t += 2.2 + 0.18 * retries
                result_id = uuid.uuid4().hex[:16]
                result_size = sample_lognormal(rng, profile["response_size_bytes_mean"],
                                                profile["message_size_cv"]) + tool_response_size * 0.15
                events.append(_event(task_id, result_id, req_id, t, profile["task_type"], profile["phase"],
                                     worker, coordinator, "Think", "Send" if success else "Failed", 5.0,
                                     "agent_response" if success else "error", result_size if success else 600,
                                     queue, success=success, retry_count=retries))
                result_sizes.append(result_size if success else 600)
                all_success = all_success and success
                t += 0.005
            final_id = uuid.uuid4().hex[:16]
            final_size = max(800.0, sum(result_sizes) * 0.38)
            final_think = sample_lognormal(rng, profile["think_time_ms_mean"] * 0.55,
                                           profile["think_time_cv"])
            t += final_think / 1000.0
            events.append(_event(task_id, final_id, parent, t, profile["task_type"], profile["phase"],
                                 coordinator, "external", "Think", "Send" if all_success else "Failed",
                                 final_think, "final_response" if all_success else "error", final_size,
                                 queue, success=all_success))
    return pd.DataFrame(events).sort_values(["timestamp", "task_id"]).reset_index(drop=True)


def aggregate_task_truth(events):
    rows = []
    for task_id, group in events.groupby("task_id"):
        group = group.sort_values("timestamp")
        external = group[group["is_external"]]
        internal = group[~group["is_external"]]
        input_bytes = int(external["message_size_bytes"].sum())
        rows.append({
            "task_id": task_id,
            "task_type": group["task_type"].iloc[0],
            "task_phase": group["task_phase"].iloc[0],
            "latency_ms": max(0.0, (group["timestamp"].max() - group["timestamp"].min()) * 1000.0),
            "internal_bytes": int(internal["message_size_bytes"].sum()),
            "input_bytes": input_bytes,
            "amplification": float(internal["message_size_bytes"].sum()) / max(input_bytes, 1),
            "message_count": int(len(internal)),
            "tool_calls": int((group["message_type"] == "tool_request").sum()),
            "retries": int(group["retry_count"].max()),
            "success": bool(group["success"].all()),
        })
    return pd.DataFrame(rows)

