"""
agentpaas xiaoyi — 华为小艺 OpenClaw A2A 桥接模块

直接实现小艺 A2A WebSocket 协议，无需 OpenClaw 框架。
将小艺消息转发到 agent67，回复自动发回小艺。

协议: A2A (Agent-to-Agent) over WebSocket
认证: AK/SK + HMAC-SHA256
消息: JSON-RPC 2.0

流程:
    小艺用户发消息 → 华为云 WebSocket → agent67 处理 → WebSocket 回复 → 小艺

用法:
    agentpaas xiaoyi start <agent_id> --ak <AK> --sk <SK> --agent-id <XIAOYI_AGENT_ID>
    agentpaas xiaoyi status
    agentpaas xiaoyi stop

参考: @ynhcj/xiaoyi OpenClaw 插件 (MIT License)
"""
from __future__ import annotations

import asyncio
import base64
import hashlib
import hmac
import json
import os
import signal
import ssl
import sys
import time
from pathlib import Path

# ── 配置 ────────────────────────────────────────────────────

CREDENTIALS_DIR = Path.home() / ".agentpaas" / "xiaoyi"
CREDENTIALS_FILE = CREDENTIALS_DIR / "credentials.json"
PID_FILE = CREDENTIALS_DIR / "daemon.pid"

# 华为小艺 A2A WebSocket 端点 (双服务器)
DEFAULT_WS_URL_1 = "wss://hag.cloud.huawei.com/openclaw/v1/ws/link"
DEFAULT_WS_URL_2 = "wss://116.63.174.231/openclaw/v1/ws/link"

HEARTBEAT_INTERVAL = 30  # 秒
RECONNECT_DELAY = 5      # 秒
MAX_RECONNECT_ATTEMPTS = 50


def _save_credentials(data):
    CREDENTIALS_DIR.mkdir(parents=True, exist_ok=True)
    CREDENTIALS_DIR.chmod(0o700)
    CREDENTIALS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False))
    CREDENTIALS_FILE.chmod(0o600)


def _load_credentials():
    if CREDENTIALS_FILE.exists():
        return json.loads(CREDENTIALS_FILE.read_text())
    return None


# ── 认证 (HMAC-SHA256) ────────────────────────────────────

def generate_signature(sk: str, timestamp: str) -> str:
    """生成 HMAC-SHA256 签名: Base64(HMAC-SHA256(sk, timestamp))"""
    mac = hmac.new(sk.encode("utf-8"), timestamp.encode("utf-8"), hashlib.sha256)
    return base64.b64encode(mac.digest()).decode("utf-8")


def generate_auth_headers(ak: str, sk: str, agent_id: str) -> dict:
    """生成 WebSocket 认证 headers"""
    ts = str(int(time.time() * 1000))
    sig = generate_signature(sk, ts)
    return {
        "x-access-key": ak,
        "x-sign": sig,
        "x-ts": ts,
        "x-agent-id": agent_id,
    }


# ── A2A 消息构造 ──────────────────────────────────────────
# 小艺要求外层包装: {msgType, agentId, sessionId, taskId, msgDetail: "<JSON-RPC>"}

def _wrap_outbound(agent_id: str, session_id: str, task_id: str,
                   jsonrpc_payload: dict) -> str:
    """将 JSON-RPC 响应包装为小艺 outbound 格式"""
    outbound = {
        "msgType": "agent_response",
        "agentId": agent_id,
        "sessionId": session_id,
        "taskId": task_id,
        "msgDetail": json.dumps(jsonrpc_payload),
    }
    return json.dumps(outbound)


def build_artifact_response(agent_id: str, session_id: str,
                            task_id: str, message_id: str,
                            text: str, final: bool = True) -> str:
    """构造 A2A artifact-update 响应 (文本回复) — 保留旧签名供兼容"""
    import uuid
    jsonrpc = {
        "jsonrpc": "2.0",
        "id": message_id,
        "result": {
            "taskId": task_id,
            "kind": "artifact-update",
            "append": False,
            "lastChunk": True,
            "final": final,
            "artifact": {
                "artifactId": str(uuid.uuid4()),
                "parts": [
                    {"kind": "text", "text": text}
                ],
            },
        },
    }
    return _wrap_outbound(agent_id, session_id, task_id, jsonrpc)


def build_text_chunk(agent_id: str, session_id: str, task_id: str,
                     message_id: str, artifact_id: str, text: str) -> str:
    """内容 artifact: append=true, final=false, lastChunk=false"""
    jsonrpc = {
        "jsonrpc": "2.0",
        "id": message_id,
        "result": {
            "taskId": task_id,
            "kind": "artifact-update",
            "append": True,
            "lastChunk": False,
            "final": False,
            "artifact": {
                "artifactId": artifact_id,
                "parts": [{"kind": "text", "text": text}],
            },
        },
    }
    return _wrap_outbound(agent_id, session_id, task_id, jsonrpc)


def build_final_marker(agent_id: str, session_id: str, task_id: str,
                       message_id: str, artifact_id: str) -> str:
    """空结束标记: final=true + text="" (小艺协议要求)"""
    jsonrpc = {
        "jsonrpc": "2.0",
        "id": message_id,
        "result": {
            "taskId": task_id,
            "kind": "artifact-update",
            "append": True,
            "lastChunk": True,
            "final": True,
            "artifact": {
                "artifactId": artifact_id,
                "parts": [{"kind": "text", "text": ""}],
            },
        },
    }
    return _wrap_outbound(agent_id, session_id, task_id, jsonrpc)


def build_reasoning_chunk(agent_id: str, session_id: str, task_id: str,
                          message_id: str, text: str) -> str:
    """keepalive: reasoningText artifact (每 8s 发一次，防小艺超时)"""
    jsonrpc = {
        "jsonrpc": "2.0",
        "id": message_id,
        "result": {
            "taskId": task_id,
            "kind": "artifact-update",
            "append": True,
            "lastChunk": True,
            "final": False,
            "artifact": {
                "artifactId": "reason_{}".format(int(time.time() * 1000)),
                "parts": [{"kind": "reasoningText",
                           "reasoningText": text + "\n"}],
            },
        },
    }
    return _wrap_outbound(agent_id, session_id, task_id, jsonrpc)


def build_status_response(agent_id: str, session_id: str,
                          task_id: str, message_id: str,
                          state: str, text: str = "") -> str:
    """构造 A2A status-update 响应"""
    jsonrpc = {
        "jsonrpc": "2.0",
        "id": message_id,
        "result": {
            "taskId": task_id,
            "kind": "status-update",
            "final": state in ("completed", "failed", "canceled"),
            "status": {
                "state": state,
                "message": {
                    "role": "agent",
                    "parts": [{"kind": "text", "text": text}] if text else [],
                },
            },
        },
    }
    return _wrap_outbound(agent_id, session_id, task_id, jsonrpc)


def build_heartbeat(agent_id: str) -> str:
    """构造心跳消息"""
    return json.dumps({
        "msgType": "heartbeat",
        "agentId": agent_id,
        "timestamp": int(time.time() * 1000),
    })


# ── 消息解析 ──────────────────────────────────────────────

def parse_user_message(data: dict) -> tuple:
    """
    解析小艺发来的 A2A 消息，提取用户文本。

    Returns: (task_id, session_id, message_id, user_text, message_parts)
    """
    # JSON-RPC 格式
    params = data.get("params", {})
    task_id = data.get("id", params.get("id", ""))
    session_id = data.get("sessionId", params.get("sessionId", ""))
    message_id = data.get("id", "")  # JSON-RPC id 作为 message_id

    message = params.get("message", {})
    parts = message.get("parts", [])

    text_parts = []
    for part in parts:
        kind = part.get("kind", "")
        if kind == "text" and part.get("text"):
            text_parts.append(part["text"])
        elif kind == "file":
            f = part.get("file", {})
            text_parts.append("[文件: {} ({})]".format(
                f.get("name", "unknown"), f.get("mimeType", "")))
        elif kind == "data":
            d = part.get("data", {})
            event = d.get("event", {})
            if event:
                text_parts.append("[事件: {} → {}]".format(
                    event.get("intentName", ""), event.get("status", "")))

    user_text = " ".join(text_parts).strip()
    return task_id, session_id, message_id, user_text, parts


# ── Agent 调用 ──────────────────────────────────────────────

_assistant = None


def call_agent_claude_code(text: str) -> str:
    """通过 Claude Code CLI 调用 agent67"""
    global _assistant
    if _assistant is None:
        project_root = Path(__file__).resolve().parent.parent
        agent_dir = project_root / "agentexample"
        if str(agent_dir) not in sys.path:
            sys.path.insert(0, str(agent_dir))
        if str(project_root) not in sys.path:
            sys.path.insert(0, str(project_root))
        from agent67.core.assistant import PersonalAssistant
        _assistant = PersonalAssistant(model="sonnet", use_api=False)
    return _assistant.chat(text)


def call_agent_paas(agent_id: str, text: str) -> str:
    """通过 AgentPaaS HTTP API 调用 agent"""
    from agentpaas.__main__ import api
    try:
        data = api("POST", "/api/v1/agents/{}/run".format(agent_id),
                    {"input": text})
        return data.get("output", "(无回复)")
    except Exception:
        return "(Agent 调用失败)"


async def call_agent_async(text: str, paas_agent_id: str = "",
                           claude_code: bool = True) -> str:
    """异步包装 agent 调用"""
    loop = asyncio.get_event_loop()
    if claude_code:
        return await loop.run_in_executor(None, call_agent_claude_code, text)
    return await loop.run_in_executor(None, call_agent_paas, paas_agent_id, text)


# ── WebSocket 桥接 ────────────────────────────────────────

async def run_bridge(creds: dict, paas_agent_id: str = "",
                     claude_code: bool = True):
    """
    启动小艺 ↔ Agent67 WebSocket 桥接。

    使用双 WebSocket 连接到华为云，接收小艺消息，
    调用 agent67 处理，然后通过 WebSocket 回复。
    """
    try:
        import websockets
    except ImportError:
        print("websockets 未安装。请运行: pip install websockets")
        return

    ak = creds["ak"]
    sk = creds["sk"]
    xiaoyi_agent_id = creds["xiaoyi_agent_id"]
    backend = "Claude Code CLI" if claude_code else "PaaS API"

    print("\n🟢 小艺 ↔ Agent67 桥接已启动 (双服务器)")
    print("  小艺 Agent:  {}".format(xiaoyi_agent_id))
    print("  后端:        {}".format(backend))
    print("  Server 1:    {}".format(DEFAULT_WS_URL_1))
    print("  Server 2:    {}".format(DEFAULT_WS_URL_2))
    print("  按 Ctrl+C 停止\n")

    # 双服务器并行监听
    await asyncio.gather(
        _listen_server("S1", DEFAULT_WS_URL_1, ak, sk, xiaoyi_agent_id,
                       paas_agent_id, claude_code, skip_cert=False),
        _listen_server("S2", DEFAULT_WS_URL_2, ak, sk, xiaoyi_agent_id,
                       paas_agent_id, claude_code, skip_cert=True),
    )


async def _listen_server(name, ws_url, ak, sk, xiaoyi_agent_id,
                         paas_agent_id, claude_code, skip_cert=False):
    """连接单个 WebSocket 服务器并监听消息。自动重连。"""
    try:
        import websockets
    except ImportError:
        print("websockets 未安装")
        return

    reconnect_attempts = 0
    while reconnect_attempts < MAX_RECONNECT_ATTEMPTS:
        try:
            headers = generate_auth_headers(ak, sk, xiaoyi_agent_id)

            ssl_ctx = None
            if skip_cert:
                ssl_ctx = ssl.create_default_context()
                ssl_ctx.check_hostname = False
                ssl_ctx.verify_mode = ssl.CERT_NONE

            connect_kwargs = {
                "extra_headers": headers,
                # 小艺不响应 WebSocket 层 ping frame，必须禁用
                # 改用应用层 heartbeat (_heartbeat_loop) 保活
                "ping_interval": None,
                "ping_timeout": None,
                "close_timeout": 5,
            }
            if ssl_ctx:
                connect_kwargs["ssl"] = ssl_ctx

            async with websockets.connect(ws_url, **connect_kwargs) as ws:
                print("  ✓ [{}] WebSocket 已连接".format(name))
                reconnect_attempts = 0

                # 关键: 发送 clawd_bot_init 注册消息，小艺才会路由消息过来
                init_msg = json.dumps({
                    "msgType": "clawd_bot_init",
                    "agentId": xiaoyi_agent_id,
                })
                await ws.send(init_msg)
                print("  ✓ [{}] clawd_bot_init 已发送".format(name))

                heartbeat_task = asyncio.create_task(
                    _heartbeat_loop(ws, xiaoyi_agent_id))

                try:
                    async for raw_msg in ws:
                        try:
                            data = json.loads(raw_msg)
                        except json.JSONDecodeError:
                            continue

                        # 忽略心跳响应
                        msg_type = data.get("msgType", data.get("type", ""))
                        if msg_type == "heartbeat":
                            continue

                        # 处理 clear context
                        action = data.get("action", "")
                        method = data.get("method", "")
                        if action == "clear" or method == "tasks/cancel":
                            print("  🔄 收到清除/取消请求")
                            continue

                        # 解析用户消息
                        if "params" not in data and "message" not in data.get("params", {}):
                            continue

                        task_id, session_id, message_id, user_text, _ = parse_user_message(data)
                        if not user_text:
                            continue

                        print("  📩 [{}] 小艺 → \"{}\" (task={}, session={})".format(
                            name, user_text[:80], task_id[:12], session_id[:12]))

                        # 起后台任务: 不阻塞主循环收新消息
                        asyncio.create_task(_handle_message(
                            ws, name, xiaoyi_agent_id, session_id, task_id,
                            message_id, user_text, paas_agent_id, claude_code))

                except websockets.exceptions.ConnectionClosed as e:
                    print("  ⚠️  WebSocket 断开: {}".format(e))
                finally:
                    heartbeat_task.cancel()

        except Exception as e:
            reconnect_attempts += 1
            print("  ✗ 连接失败 ({}/{}): {}".format(
                reconnect_attempts, MAX_RECONNECT_ATTEMPTS, e))
            await asyncio.sleep(RECONNECT_DELAY * min(reconnect_attempts, 6))

    print("🔴 [{}] 达到最大重连次数".format(name))


async def _handle_message(ws, name, xiaoyi_agent_id, session_id, task_id,
                          message_id, user_text, paas_agent_id, claude_code):
    """处理单条用户消息: keepalive + 调用 agent + 三段式回复"""
    reason_msgs = ["正在思考...", "分析问题中...", "调用工具中...",
                   "处理数据...", "即将完成..."]
    done_event = asyncio.Event()

    async def keepalive():
        tick = 0
        while not done_event.is_set():
            try:
                await asyncio.wait_for(done_event.wait(), timeout=8)
                return
            except asyncio.TimeoutError:
                pass
            tick += 1
            try:
                await ws.send(build_reasoning_chunk(
                    xiaoyi_agent_id, session_id, task_id, message_id,
                    reason_msgs[tick % len(reason_msgs)]))
                print("  ♻ [{}] reasoning #{}".format(name, tick))
            except Exception:
                return

    ka_task = asyncio.create_task(keepalive())
    t0 = time.time()
    try:
        reply = await call_agent_async(user_text, paas_agent_id, claude_code)
    except Exception as e:
        reply = "抱歉，处理出错: {}".format(e)
    finally:
        done_event.set()
        try:
            await ka_task
        except Exception:
            pass

    elapsed = time.time() - t0
    print("  💬 [{}] ({:.1f}s) → \"{}\"".format(name, elapsed, reply[:80]))

    artifact_id = "artifact_{}".format(int(time.time() * 1000))
    try:
        await ws.send(build_text_chunk(
            xiaoyi_agent_id, session_id, task_id, message_id,
            artifact_id, reply))
        await ws.send(build_final_marker(
            xiaoyi_agent_id, session_id, task_id, message_id, artifact_id))
        await ws.send(build_status_response(
            xiaoyi_agent_id, session_id, task_id, message_id, "completed"))
        print("  ✓ [{}] 已回复 (task={})\n".format(name, task_id[:12]))
    except Exception as e:
        print("  ✗ [{}] 回复失败: {}".format(name, e))


async def _heartbeat_loop(ws, agent_id: str):
    """定时发送心跳"""
    while True:
        try:
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            await ws.send(build_heartbeat(agent_id))
        except asyncio.CancelledError:
            break
        except Exception:
            break


# ── 公开入口 ────────────────────────────────────────────────

def cmd_xiaoyi_start(paas_agent_id: str, ak: str, sk: str,
                     xiaoyi_agent_id: str, claude_code: bool = True,
                     daemon: bool = False):
    """启动小艺桥接。"""
    creds = {
        "ak": ak,
        "sk": sk,
        "xiaoyi_agent_id": xiaoyi_agent_id,
        "paas_agent_id": paas_agent_id,
        "ws_url": DEFAULT_WS_URL_1,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
    }
    _save_credentials(creds)

    if daemon:
        pid = os.fork()
        if pid > 0:
            print("后台运行中，PID: {}".format(pid))
            PID_FILE.parent.mkdir(parents=True, exist_ok=True)
            PID_FILE.write_text(str(pid))
            return
        os.setsid()

    try:
        asyncio.run(run_bridge(creds, paas_agent_id, claude_code))
    except KeyboardInterrupt:
        print("\n🔴 桥接已停止")
    finally:
        if daemon and PID_FILE.exists():
            PID_FILE.unlink()


def cmd_xiaoyi_status():
    """查看小艺桥接状态。"""
    creds = _load_credentials()
    if not creds:
        print("未配置小艺 OpenClaw，请先运行: agentpaas xiaoyi start ...")
        return

    print("=== 小艺 OpenClaw 状态 ===")
    print("  Agent ID:    {}".format(creds.get("xiaoyi_agent_id", "未知")))
    print("  PaaS Agent:  {}".format(creds.get("paas_agent_id", "未绑定")))
    print("  配置时间:    {}".format(creds.get("created_at", "未知")))

    pid = None
    if PID_FILE.exists():
        try:
            pid = int(PID_FILE.read_text().strip())
        except ValueError:
            pass
    if pid:
        try:
            os.kill(pid, 0)
            print("  守护进程:    运行中 (PID {})".format(pid))
        except OSError:
            print("  守护进程:    已停止")
    else:
        print("  守护进程:    未运行")


def cmd_xiaoyi_stop():
    """停止后台桥接进程。"""
    if not PID_FILE.exists():
        print("没有运行中的守护进程")
        return
    try:
        pid = int(PID_FILE.read_text().strip())
        os.kill(pid, signal.SIGTERM)
        print("已发送停止信号到 PID {}".format(pid))
        for _ in range(10):
            try:
                os.kill(pid, 0)
                time.sleep(0.5)
            except OSError:
                break
        PID_FILE.unlink()
        print("守护进程已停止")
    except (ValueError, OSError) as e:
        print("停止失败: {}".format(e))
        if PID_FILE.exists():
            PID_FILE.unlink()
