"""A2A 网关接入客户端：Agent 服务启动后自动注册 + 心跳保活 + 注销。

用法（在 api_server.py 的 main() 中）:
    from gateway_client import register_agent, start_heartbeat, unregister_agent

    if register_agent(dev_app.name, f"http://127.0.0.1:{PORT}"):
        start_heartbeat(dev_app.name)

环境变量:
    GATEWAY_URL     网关地址，默认 http://127.0.0.1:8000
    GATEWAY_AUTH    网关认证密码，需与网关 GATEWAY_AUTH 一致
    AGENT_TAGS      能力标签，逗号分隔，默认 code
    AGENT_MAX_CONCURRENT  最大并发数，默认 1
    HEARTBEAT_INTERVAL    心跳间隔秒，默认 10
"""
import logging
import os
import threading
import time

import httpx

logger = logging.getLogger(__name__)

GATEWAY_URL = os.getenv("GATEWAY_URL", "http://127.0.0.1:8000").rstrip("/")
GATEWAY_AUTH = os.getenv("GATEWAY_AUTH", "dev-gateway-auth")
HEARTBEAT_INTERVAL = int(os.getenv("HEARTBEAT_INTERVAL", "10"))

_client = httpx.Client(timeout=5.0)

# 取消/停止指令下发：网关复用 CLI 的 /api/cli/events 通道（event=task_stop）
# Agent 通过同一个 cli_session_id 订阅，收到匹配自身正在执行 request_id 的 task_stop 时置停止标志。
_stop_registry: dict[str, threading.Event] = {}


def mark_stop_requested(request_id: str) -> None:
    """标记某任务需要停止（收到 task_stop 后调用）。"""
    ev = _stop_registry.get(request_id)
    if ev is None:
        ev = threading.Event()
        _stop_registry[request_id] = ev
    ev.set()


def clear_stop_requested(request_id: str) -> None:
    """任务开始执行前清除停止标志。"""
    ev = _stop_registry.get(request_id)
    if ev is not None:
        ev.clear()


def is_stop_requested(request_id: str) -> bool:
    """判断某任务是否已被要求停止（供执行循环轮询检查）。"""
    ev = _stop_registry.get(request_id)
    return ev is not None and ev.is_set()


def wait_stop(request_id: str, timeout: float = 0.5) -> bool:
    """等待停止标志，返回 True 表示已收到停止请求。执行循环可用它做可中断 sleep。"""
    ev = _stop_registry.get(request_id)
    if ev is None:
        ev = threading.Event()
        _stop_registry[request_id] = ev
    return ev.wait(timeout)


def _sse_subscribe_poll(cli_session_id: str) -> None:
    """后台线程：订阅网关 /api/cli/events 通道，识别 task_stop 指令并标记停止。"""
    url = f"{GATEWAY_URL}/api/cli/events?cli_session_id={cli_session_id}&auth={GATEWAY_AUTH}"
    while True:
        try:
            with _client.stream("GET", url, timeout=None) as resp:
                if resp.status_code != 200:
                    logger.warning("sse subscribe failed status=%s", resp.status_code)
                    time.sleep(HEARTBEAT_INTERVAL)
                    continue
                for line in resp.iter_lines():
                    if not line or not line.startswith("data:"):
                        continue
                    try:
                        import json
                        evt = json.loads(line[len("data:"):].strip())
                    except Exception:
                        continue
                    if evt.get("event") == "task_stop":
                        rid = evt.get("request_id")
                        if rid:
                            mark_stop_requested(rid)
                            logger.info("stop requested received request=%s", rid)
        except Exception as e:
            logger.warning("sse subscribe loop error err=%s", e)
        time.sleep(HEARTBEAT_INTERVAL)


def start_stop_listener(cli_session_id: str) -> threading.Thread:
    """启动 SSE 停止指令订阅线程（daemon）。"""
    t = threading.Thread(
        target=_sse_subscribe_poll,
        args=(cli_session_id,),
        name=f"gateway-sse-{cli_session_id[:8]}",
        daemon=True,
    )
    t.start()
    logger.info("stop listener started session=%s", cli_session_id)
    return t


def _parse_tags() -> list[str]:
    raw = os.getenv("AGENT_TAGS", "code")
    return [t.strip() for t in raw.split(",") if t.strip()]


def register_agent(agent_id: str, endpoint: str, agent_tags: list[str] | None = None,
                   max_concurrent: int | None = None) -> bool:
    """向网关注册本 Agent。成功返回 True，失败（网关未启动/认证失败）返回 False。"""
    tags = agent_tags if agent_tags is not None else _parse_tags()
    max_conc = max_concurrent or int(os.getenv("AGENT_MAX_CONCURRENT", "1"))
    body = {
        "auth": GATEWAY_AUTH,
        "agent_id": agent_id,
        "endpoint": endpoint,
        "agent_tags": tags,
        "max_concurrent": max_conc,
        "current_load": 0,
    }
    try:
        resp = _client.post(f"{GATEWAY_URL}/api/agent/register", json=body)
        if resp.status_code == 200:
            logger.info("registered to gateway agent=%s endpoint=%s tags=%s", agent_id, endpoint, tags)
            return True
        logger.error("register failed agent=%s status=%s body=%s", agent_id, resp.status_code, resp.text)
    except httpx.HTTPError as e:
        logger.error("register network error agent=%s err=%s", agent_id, e)
    return False


def _heartbeat_once(agent_id: str, current_load: int) -> bool:
    try:
        resp = _client.post(
            f"{GATEWAY_URL}/api/agent/heartbeat",
            json={"auth": GATEWAY_AUTH, "agent_id": agent_id, "current_load": current_load},
        )
        if resp.status_code == 200:
            return True
        logger.warning("heartbeat failed status=%s body=%s", resp.status_code, resp.text)
    except httpx.HTTPError as e:
        logger.warning("heartbeat network error err=%s", e)
    return False


def _heartbeat_loop(agent_id: str) -> None:
    while True:
        try:
            load = get_current_load()
        except Exception:
            load = 0
        _heartbeat_once(agent_id, load)
        time.sleep(HEARTBEAT_INTERVAL)


def start_heartbeat(agent_id: str) -> threading.Thread:
    """启动后台心跳线程（daemon，随进程退出）。"""
    t = threading.Thread(target=_heartbeat_loop, args=(agent_id,), name=f"gateway-hb-{agent_id}", daemon=True)
    t.start()
    logger.info("heartbeat started agent=%s interval=%ss", agent_id, HEARTBEAT_INTERVAL)
    return t


def unregister_agent(agent_id: str) -> bool:
    """向网关注销本 Agent。"""
    try:
        resp = _client.post(f"{GATEWAY_URL}/api/agent/unregister", json={"auth": GATEWAY_AUTH, "agent_id": agent_id})
        if resp.status_code in (200, 404):
            logger.info("unregistered from gateway agent=%s", agent_id)
            return True
        logger.warning("unregister failed status=%s body=%s", resp.status_code, resp.text)
    except httpx.HTTPError as e:
        logger.error("unregister network error agent=%s err=%s", agent_id, e)
    return False


def report_result(request_id: str, agent_id: str, status: str = "success",
                  progress: int = 100, result: dict | None = None,
                  error_info: str | None = None) -> bool:
    """任务执行完成后，将结果回传给网关（POST /api/agent/result）。

    Args:
        request_id: 网关分配的任务 ID
        agent_id: 本 Agent ID
        status: success / failed
        progress: 进度百分比
        result: 结果字典（可选）
        error_info: 错误信息（失败时必填）
    """
    body = {
        "auth": GATEWAY_AUTH,
        "request_id": request_id,
        "agent_id": agent_id,
        "status": status,
        "progress": progress,
        "result": result,
        "error_info": error_info,
    }
    try:
        resp = _client.post(f"{GATEWAY_URL}/api/agent/result", json=body)
        if resp.status_code == 200:
            logger.info("result reported request=%s status=%s", request_id, status)
            return True
        logger.error("report result failed request=%s status=%s body=%s", request_id, resp.status_code, resp.text)
    except httpx.HTTPError as e:
        logger.error("report result network error request=%s err=%s", request_id, e)
    return False


def get_current_load() -> int:
    """Agent 当前负载。子类/调用方可覆写以报告真实并发数。"""
    return 0