"""Independent loopback broker. No simulation, training or LeRobot imports."""

import asyncio
import contextlib
import hmac
import re
import secrets
import time
from collections import deque

from aiohttp import WSMsgType, web

from . import protocol as p
from .protocol import RobotError

MAX_BYTES = 65536
PREFIX = "/api/control/v1"
ORIGIN = re.compile(r"http://(?:127\.0\.0\.1|localhost)(?::[0-9]{1,5})?\Z")
HOST = re.compile(r"(?:127\.0\.0\.1|localhost)(?::[0-9]{1,5})?\Z")
BROKER = web.AppKey("broker", object)


class Broker:
    def __init__(self, token):
        if (
            not isinstance(token, str)
            or not 16 <= len(token) <= 4096
            or not all(33 <= ord(c) <= 126 for c in token)
        ):
            raise ValueError("控制 token 必须为16–4096字符的可打印 ASCII，且不含空白")
        self.token = token
        self.ws = None
        self.descriptor = None
        self.observation = None
        self.observed_at = 0.0
        self.authorization_generation = 0
        self.blocked_generation = -1
        self.authorized = False
        self.lease = None
        self.last_action_at = 0.0
        self.last_action_seq = 0
        self.rate = deque()
        self.pending = {}

    def robot(self):
        if self.ws is None or self.ws.closed or self.descriptor is None:
            raise RobotError("DISCONNECTED", "浏览器机器人未连接")
        return self.descriptor

    def fresh_observation(self):
        self.robot()
        # Paused samples stop advancing by design; do not hide the required play
        # step behind a freshness error merely because the user waited to connect.
        if self.observation is not None and self.observation["paused"]:
            raise RobotError("PAUSED", "仿真已暂停；请先在浏览器点击“播放”，再点击“允许外部控制”")
        if self.observation is None or time.monotonic() - self.observed_at > 0.5:
            raise RobotError(
                "STALE",
                "机器人观测超过 500ms 未更新；请保持仿真页面可见，"
                "检查仿真是否卡顿，再播放并重新允许外部控制",
            )
        return self.observation

    async def state(self, packet):
        p.record(packet, {"type", "observation", "enabled", "authorizationGeneration"})
        if type(packet["enabled"]) is not bool:
            p.invalid("授权标记必须为 bool")
        generation = p.integer(packet["authorizationGeneration"])
        if generation < self.authorization_generation:
            raise RobotError("STALE", "授权代次倒退")
        observed = p.observation(packet["observation"], self.robot())
        old = self.observation
        if old:
            p.same_identity(observed, old, ("sessionId",))
            if observed["modelEpoch"] < old["modelEpoch"] or (
                observed["modelEpoch"] == old["modelEpoch"]
                and observed["sequence"] < old["sequence"]
            ):
                raise RobotError("STALE", "观测代次/序号倒退")
        # Sequence may restart only in a new epoch. Same-sample heartbeats aren't fresh.
        if (
            old is None
            or observed["modelEpoch"] > old["modelEpoch"]
            or observed["sequence"] > old["sequence"]
        ):
            self.observation = observed
            self.observed_at = time.monotonic()
        if self.lease and generation != self.authorization_generation:
            # Cancel the previous owner without consuming the new explicit grant.
            await self.revoke("浏览器已重新授权", notify=False)
        self.authorization_generation = generation
        self.authorized = packet["enabled"] and generation > self.blocked_generation
        if self.lease and (
            not self.authorized
            or observed["paused"]
            or any(observed[k] != self.lease[k] for k in ("sessionId", "modelEpoch"))
        ):
            await self.revoke("浏览器已撤销控制/重置模型", notify=False)

    async def rpc(self, operation, payload):
        self.robot()
        if operation == "action":
            for key, (future, op) in list(self.pending.items()):
                if op == "action":
                    if not future.done():
                        future.set_exception(RobotError("SUPERSEDED", "已由更新的目标替代"))
                    self.pending.pop(key, None)
        if len(self.pending) >= 8:
            raise RobotError("CONFLICT", "待确认请求已满")
        request_id = secrets.token_hex(16)
        future = asyncio.get_running_loop().create_future()
        ws, lease, generation = self.ws, self.lease, self.authorization_generation
        self.pending[request_id] = (future, operation)
        try:
            await asyncio.wait_for(
                self.ws.send_json(
                    {"type": "request", "id": request_id, "op": operation, "payload": payload}
                ),
                1,
            )
            return await asyncio.wait_for(future, 1)
        except TimeoutError as exc:
            if (
                self.ws is ws
                and self.lease == lease
                and self.authorization_generation == generation
            ):
                await self.revoke("动作应用确认超时")
            raise RobotError("TIMEOUT", "浏览器未在一秒内确认应用") from exc
        finally:
            self.pending.pop(request_id, None)
            if not future.done():
                future.cancel()
            elif not future.cancelled():
                future.exception()  # Also consume failures delivered during a blocked send.

    async def revoke(self, reason, notify=True):
        self.lease = None
        self.authorized = False
        self.blocked_generation = self.authorization_generation
        for future, _ in self.pending.values():
            if not future.done():
                future.set_exception(RobotError("DISCONNECTED", reason))
        self.pending.clear()
        if notify and self.ws is not None and not self.ws.closed and self.observation:
            with contextlib.suppress(ConnectionError, TimeoutError):
                await asyncio.wait_for(
                    self.ws.send_json(
                        {
                            "type": "stop",
                            "reason": reason,
                            "sessionId": self.observation["sessionId"],
                            "authorizationGeneration": self.authorization_generation,
                        }
                    ),
                    0.1,
                )

    def require_lease(self, request):
        provided = request.headers.get("X-Control-Lease", "")
        if (
            not self.lease
            or not provided.isascii()
            or not hmac.compare_digest(provided.encode(), self.lease["leaseId"].encode())
        ):
            raise RobotError("UNAUTHORIZED", "控制租约无效")
        return dict(self.lease)

    async def monitor(self):
        while True:
            await asyncio.sleep(0.05)
            if self.lease and (
                time.monotonic() - self.last_action_at > 0.5
                or time.monotonic() - self.observed_at > 0.5
            ):
                await self.revoke("动作或观测看门狗超时 (500ms)")


@web.middleware
async def security(request, handler):
    broker = request.app[BROKER]
    try:
        if not HOST.fullmatch(request.headers.get("Host", "")):
            raise RobotError("UNAUTHORIZED", "Host 不被允许")
        origin = request.headers.get("Origin")
        if origin is not None and not ORIGIN.fullmatch(origin):
            raise RobotError("UNAUTHORIZED", "Origin 不被允许")
        if request.query:
            raise RobotError("INVALID_MESSAGE", "禁止 URL 查询参数及 URL 中的 token")
        if request.path != "/ws/control/v1":
            expected = f"Bearer {broker.token}".encode()
            provided = request.headers.get("Authorization", "")
            if not provided.isascii() or not hmac.compare_digest(provided.encode(), expected):
                raise RobotError("UNAUTHORIZED", "需要本机控制 Bearer token")
        elif origin is None:
            raise RobotError("UNAUTHORIZED", "浏览器 WebSocket 必须提供 Origin")
        return await handler(request)
    except RobotError as exc:
        status = {
            "UNAUTHORIZED": 401,
            "DISCONNECTED": 503,
            "TIMEOUT": 504,
            "INVALID_MESSAGE": 400,
            "UNSUPPORTED": 400,
        }.get(exc.code, 409)
        return web.json_response({"error": exc.as_dict()}, status=status)
    except web.HTTPRequestEntityTooLarge:
        return web.json_response(
            {"error": {"code": "INVALID_MESSAGE", "message": "消息不能超过64KiB"}}, status=413
        )


async def body(request):
    if request.content_type != "application/json":
        p.invalid("需要 application/json")
    return p.loads(await request.read())


async def health(request):
    b = request.app[BROKER]
    return web.json_response(
        {
            "protocolVersion": 1,
            "backendConnected": b.descriptor is not None,
            "authorized": b.authorized,
            "hasLease": b.lease is not None,
        }
    )


async def robot(request):
    return web.json_response(request.app[BROKER].robot())


async def observation(request):
    return web.json_response(request.app[BROKER].fresh_observation())


async def claim(request):
    b = request.app[BROKER]
    data = p.record(await body(request), {"sessionId", "modelEpoch", "modelFingerprint"})
    p.text(data["sessionId"])
    p.integer(data["modelEpoch"])
    b.robot()
    observed = b.fresh_observation()
    p.same_identity(data, observed, ("sessionId", "modelEpoch"))
    if data["modelFingerprint"] != b.descriptor["modelFingerprint"]:
        raise RobotError("INCOMPATIBLE_MODEL", "模型指纹已变化")
    if b.lease:
        raise RobotError("CONFLICT", "已有一个 Python 控制者")
    if not b.authorized:
        raise RobotError("UNAUTHORIZED", "请先在浏览器显式允许外部控制")
    lease = {
        "sessionId": data["sessionId"],
        "modelEpoch": data["modelEpoch"],
        "leaseId": secrets.token_hex(24),
    }
    b.lease = lease
    b.last_action_at = time.monotonic()
    b.last_action_seq = 0
    b.rate.clear()
    try:
        result = p.record(
            await b.rpc(
                "claim",
                {
                    **lease,
                    "modelFingerprint": data["modelFingerprint"],
                    "authorizationGeneration": b.authorization_generation,
                },
            ),
            p.IDENTITY,
        )
        p.same_identity(result, lease)
        if b.lease != lease:
            raise RobotError("STALE", "租约在申请期间失效")
        return web.json_response(result)
    except BaseException:
        if b.lease == lease:
            await b.revoke("租约申请失败")
        raise


async def action(request):
    b = request.app[BROKER]
    lease = b.require_lease(request)
    b.fresh_observation()
    data = p.action(await body(request), b.robot())
    p.same_identity(data, lease)
    if data["actionSeq"] <= b.last_action_seq:
        raise RobotError("STALE", "拒绝重复/乱序动作")
    now = time.monotonic()
    while b.rate and now - b.rate[0] >= 1:
        b.rate.popleft()
    if len(b.rate) >= 100:
        raise RobotError("CONFLICT", "动作频率不能超过100Hz")
    b.rate.append(now)
    b.last_action_seq = data["actionSeq"]
    b.last_action_at = now
    try:
        result = p.action_result(await b.rpc("action", data), b.robot())
        p.same_identity(result, data, (*p.IDENTITY, "actionSeq"))
    except RobotError as exc:
        if exc.code != "SUPERSEDED" and b.lease == lease:
            await b.revoke("动作确认失败")
        raise
    if b.lease != lease:
        raise RobotError("STALE", "动作确认来自失效租约")
    return web.json_response(result)


async def release(request):
    b = request.app[BROKER]
    lease = b.require_lease(request)
    try:
        await b.rpc("release", lease)
    finally:
        if b.lease == lease:
            await b.revoke("控制者断开连接")
    return web.json_response({"released": True})


async def reset(request):
    b = request.app[BROKER]
    lease = b.require_lease(request)
    p.record(await body(request), set())
    if not b.robot()["capabilities"]["reset"]:
        raise RobotError("UNSUPPORTED", "机器人不支持 reset")
    try:
        result = p.observation(await b.rpc("reset", lease), b.robot())
        if (
            result["sessionId"] != lease["sessionId"]
            or result["modelEpoch"] <= lease["modelEpoch"]
            or not result["paused"]
        ):
            p.invalid("reset 未返回新代次的暂停观测")
    finally:
        if b.lease == lease:
            await b.revoke("reset 后需要重新授权")
    return web.json_response(result)


async def websocket(request):
    b = request.app[BROKER]
    ws = web.WebSocketResponse(
        max_msg_size=MAX_BYTES, heartbeat=2, receive_timeout=5, compress=False
    )
    await ws.prepare(request)
    registered = False
    times = deque()
    try:
        # Outer deadline: WS ping/pong must NOT restart the authentication clock.
        first = await asyncio.wait_for(ws.receive(), 5)
        if first.type != WSMsgType.TEXT:
            raise RobotError("UNAUTHORIZED", "必须在5秒内通过认证")
        auth = p.record(p.loads(first.data), {"type", "token", "protocolVersion"})
        p.version(auth["protocolVersion"])
        if (
            auth["type"] != "auth"
            or not isinstance(auth["token"], str)
            or not auth["token"].isascii()
            or not hmac.compare_digest(auth["token"].encode(), b.token.encode())
        ):
            raise RobotError("UNAUTHORIZED", "WebSocket token 无效")
        if b.ws is not None:
            raise RobotError("CONFLICT", "已有浏览器后端连接")
        b.ws = ws  # Reserve the single backend before any further await.
        await ws.send_json({"type": "authenticated"})
        while True:
            message = (
                await asyncio.wait_for(ws.receive(), 5) if not registered else await ws.receive()
            )
            if message.type != WSMsgType.TEXT:
                break
            now = time.monotonic()
            while times and now - times[0] > 1:
                times.popleft()
            if len(times) >= 256:
                raise RobotError("CONFLICT", "浏览器消息过于频繁")
            times.append(now)
            packet = p.loads(message.data)
            if not isinstance(packet, dict):
                p.invalid("消息必须是对象")
            kind = packet.get("type")
            if kind == "register" and not registered:
                p.record(
                    packet,
                    {"type", "descriptor", "observation", "enabled", "authorizationGeneration"},
                )
                b.descriptor = p.descriptor(packet["descriptor"])
                await b.state({k: v for k, v in packet.items() if k != "descriptor"})
                registered = True
                await ws.send_json({"type": "ready"})
            elif kind == "state" and registered:
                await b.state(packet)
            elif kind == "result" and registered:
                p.record(packet, {"type", "id", "ok", "value"})
                p.text(packet["id"])
                if type(packet["ok"]) is not bool:
                    p.invalid("ok 必须为 bool")
                pending = b.pending.get(packet["id"])
                if pending and not pending[0].done():
                    if packet["ok"]:
                        pending[0].set_result(packet["value"])
                    else:
                        error = p.record(packet["value"], {"code", "message"})
                        if (
                            not isinstance(error["code"], str)
                            or not isinstance(error["message"], str)
                            or len(error["message"]) > 1024
                        ):
                            p.invalid("错误格式无效")
                        pending[0].set_exception(RobotError(error["code"], error["message"]))
            else:
                p.invalid("未注册后端或未知消息类型")
    except (RobotError, TimeoutError, ConnectionError) as exc:
        error = (
            exc
            if isinstance(exc, RobotError)
            else RobotError("DISCONNECTED", "浏览器连接中断/超时")
        )
        if not ws.closed:
            with contextlib.suppress(ConnectionError):
                await ws.send_json({"type": "error", "error": error.as_dict()})
    finally:
        if b.ws is ws:
            await b.revoke("浏览器已断开", notify=False)
            b.ws = b.descriptor = b.observation = None
            b.authorization_generation = 0
            b.blocked_generation = -1
        await ws.close()
    return ws


def create_app(token):
    app = web.Application(middlewares=[security], client_max_size=MAX_BYTES)
    app[BROKER] = Broker(token)
    app.add_routes(
        [
            web.get(f"{PREFIX}/health", health),
            web.get(f"{PREFIX}/robot", robot),
            web.get(f"{PREFIX}/observation", observation),
            web.post(f"{PREFIX}/lease", claim),
            web.delete(f"{PREFIX}/lease", release),
            web.post(f"{PREFIX}/action", action),
            web.post(f"{PREFIX}/reset", reset),
            web.get("/ws/control/v1", websocket),
        ]
    )

    async def lifetime(application):
        b = application[BROKER]
        task = asyncio.create_task(b.monitor())
        yield
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task
        await b.revoke("桥接服务关闭")
        if b.ws is not None:
            await b.ws.close()

    app.cleanup_ctx.append(lifetime)
    return app