from __future__ import annotations

import asyncio
import json
import os
import shutil
import uuid
from pathlib import Path


class WorkerError(Exception):
    def __init__(self, message: str, code: str = "worker_error"):
        super().__init__(message)
        self.code = code


class BrowserWorker:
    """Serialized RPC to the single browser owner; never imports the old application."""

    def __init__(self, root: Path, home: Path):
        self.root, self.home = root, home
        self.process: asyncio.subprocess.Process | None = None
        self.lock = asyncio.Lock()
        self.stderr_task: asyncio.Task | None = None
        self.closing = False
        self.close_task: asyncio.Task | None = None

    async def _start(self) -> None:
        if self.process and self.process.returncode is None:
            return
        node = os.environ.get("ASSISTANT_NODE") or shutil.which("node")
        if not node:
            raise WorkerError("找不到 Node.js，请通过项目启动脚本运行后台", "node_not_found")
        script = self.root / "tools" / "learning" / "browser-worker.mjs"
        if not script.is_file():
            raise WorkerError("本机学习执行器尚未就绪", "worker_unavailable")
        environment = {**os.environ, "ASSISTANT_LEARNING_HOME": str(self.home)}
        self.process = await asyncio.create_subprocess_exec(node, str(script), cwd=self.root, env=environment, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=8 * 1024 * 1024)
        if self.closing:
            await self._stop()
            raise WorkerError("学习执行器正在正常退出", "worker_closed")
        process = self.process

        async def consume_stderr() -> None:
            assert process.stderr
            while await process.stderr.readline():
                pass  # Do not forward browser diagnostic text or URLs into business logs.

        self.stderr_task = asyncio.create_task(consume_stderr())

    async def call(self, method: str, params: dict) -> dict:
        async with self.lock:
            if self.closing:
                raise WorkerError("学习执行器正在正常退出", "worker_closed")
            await self._start()
            process = self.process
            assert process and process.stdin and process.stdout
            request_id = uuid.uuid4().hex
            try:
                process.stdin.write((json.dumps({"id": request_id, "method": method, "params": params}, ensure_ascii=False) + "\n").encode())
                await process.stdin.drain()
                async with asyncio.timeout(120):
                    while True:
                        line = await process.stdout.readline()
                        if not line:
                            raise WorkerError("本机学习浏览器已退出，请重新打开会话", "worker_closed")
                        try:
                            response = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        if response.get("id") != request_id:
                            continue
                        if response.get("error"):
                            error = response["error"]
                            raise WorkerError(str(error.get("message", "学习执行失败")), str(error.get("code", "worker_error")))
                        return response.get("result") or {}
            except (TimeoutError, BrokenPipeError, ConnectionResetError):
                await self._stop()
                raise WorkerError("学习执行器响应中断；请核对页面与步骤后重新打开，不会自动重放动作", "worker_timeout") from None

    async def _stop(self) -> None:
        process, self.process = self.process, None
        if not process:
            return
        if process.stdin:
            process.stdin.close()
        try:
            await asyncio.wait_for(process.wait(), 15)
        except TimeoutError:
            process.terminate()
            try:
                await asyncio.wait_for(process.wait(), 3)
            except TimeoutError:
                process.kill()
                await process.wait()
        if self.stderr_task:
            self.stderr_task.cancel()
            self.stderr_task = None

    async def close(self) -> None:
        # EOF must reach the child even while an RPC holds the request lock.
        # Node closes its contexts on EOF, unblocking that exact active action.
        self.closing = True
        if self.close_task is None:
            self.close_task = asyncio.create_task(self._stop())
        await asyncio.shield(self.close_task)