"""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)


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={"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={"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 = {
        "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