"""scrcpy server adapter — manages scrcpy H.264 stream for WebRTC.

Launches scrcpy-server on the Android device via ADB, reads the H.264 NAL
unit stream from an ADB-forwarded socket, and pushes NAL units into a
``ScrcpyVideoTrack`` for WebRTC streaming.
"""

from __future__ import annotations

import asyncio
import logging
import socket
import struct
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .video_source import ScrcpyVideoTrack

logger = logging.getLogger(__name__)

_DEFAULT_SERVER_JAR = "/data/local/tmp/scrcpy-server"
_DEFAULT_VERSION = "2.7"
_DEFAULT_LOCAL_PORT = 27183
_ADB_MAX_RETRIES = 3


class ScrcpyAdapter:
    """Manages scrcpy server lifecycle and H.264 frame extraction.

    This adapter:
    1. Pushes scrcpy-server.jar to the device via ADB
    2. Starts scrcpy-server in raw-stream mode with tunnel_forward
    3. Connects to the ADB-forwarded local socket
    4. Reads H.264 NAL units with frame-meta header (PTS + size)
    5. Pushes NAL units into a ``ScrcpyVideoTrack`` for WebRTC
    """

    def __init__(
        self,
        adb_serial: str,
        max_fps: int = 30,
        max_width: int = 1080,
        bitrate: int = 2_000_000,
        adb_host: str = "127.0.0.1",
        adb_port: int = 5037,
        server_jar: str = _DEFAULT_SERVER_JAR,
        scrcpy_version: str = _DEFAULT_VERSION,
        local_port: int = _DEFAULT_LOCAL_PORT,
    ) -> None:
        self._serial = adb_serial
        self._max_fps = max_fps
        self._max_width = max_width
        self._bitrate = bitrate
        self._adb_host = adb_host
        self._adb_port = adb_port
        self._server_jar = server_jar
        self._scrcpy_version = scrcpy_version
        self._local_port = local_port
        self._process: asyncio.subprocess.Process | None = None
        self._reader_task: asyncio.Task | None = None
        self._video_track: ScrcpyVideoTrack | None = None
        self._running = False

    @property
    def running(self) -> bool:
        return self._running

    @property
    def video_track(self) -> ScrcpyVideoTrack | None:
        return self._video_track

    async def start(self) -> ScrcpyVideoTrack:
        """Push scrcpy-server, start it, and return a video track."""
        from .video_source import ScrcpyVideoTrack as _Track

        self._video_track = _Track(target_fps=self._max_fps)

        await self._push_server()
        await self._setup_forward()
        await self._start_server()
        await asyncio.sleep(0.5)

        self._reader_task = asyncio.create_task(self._read_loop())
        self._running = True
        logger.info("scrcpy adapter started (serial=%s)", self._serial)
        return self._video_track

    async def stop(self) -> None:
        """Stop the scrcpy server process and reader."""
        self._running = False
        if self._reader_task and not self._reader_task.done():
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass
            self._reader_task = None

        if self._video_track:
            self._video_track.stop()
            self._video_track = None

        if self._process:
            try:
                self._process.terminate()
                await asyncio.wait_for(self._process.wait(), timeout=5.0)
            except (asyncio.TimeoutError, ProcessLookupError):
                if self._process:
                    self._process.kill()
            self._process = None

        await self._remove_forward()
        logger.info("scrcpy adapter stopped")

    async def _adb_exec(self, *args: str) -> tuple[bytes, bytes, int]:
        """Run an ADB command with retry logic, returning (stdout, stderr, returncode)."""
        base = ["adb", "-H", self._adb_host, "-P", str(self._adb_port), "-s", self._serial]
        cmd = base + list(args)
        last_err: Exception | None = None
        for attempt in range(_ADB_MAX_RETRIES):
            try:
                proc = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                stdout, stderr = await proc.communicate()
                if proc.returncode == 0:
                    return stdout, stderr, 0
                last_err = RuntimeError(
                    f"ADB command {args} failed (exit {proc.returncode}): {stderr.decode(errors='replace')}"
                )
            except Exception as e:
                last_err = e
            if attempt < _ADB_MAX_RETRIES - 1:
                await asyncio.sleep(0.5 * (attempt + 1))
        return b"", b"", -1

    async def _push_server(self) -> None:
        """Push scrcpy-server.jar to the device if not already present."""
        stdout, _, rc = await self._adb_exec(
            "shell", f"test -f {self._server_jar} && echo exists"
        )
        if b"exists" in stdout:
            return

        import shutil
        local_jar = shutil.which("scrcpy-server")
        if local_jar is None:
            import importlib.resources
            try:
                ref = importlib.resources.files("device_agent") / "data" / "scrcpy-server"
                local_jar = str(ref)
            except Exception:
                local_jar = None

        if local_jar and __import__("os").path.isfile(local_jar):
            logger.info("Pushing scrcpy-server.jar to device %s ...", self._serial)
            _, _, rc = await self._adb_exec("push", local_jar, self._server_jar)
            if rc == 0:
                return

        raise RuntimeError(
            f"scrcpy-server.jar not found on device {self._serial} and auto-push failed. "
            f"Push it manually: adb -s {self._serial} push scrcpy-server {self._server_jar}"
        )

    async def _setup_forward(self) -> None:
        """Set up ADB port-forward for scrcpy tunnel."""
        _, stderr, rc = await self._adb_exec(
            "forward", f"tcp:{self._local_port}", "localabstract:scrcpy"
        )
        if rc != 0:
            raise RuntimeError(f"ADB forward failed: {stderr.decode(errors='replace')}")

    async def _remove_forward(self) -> None:
        await self._adb_exec("forward", "--remove", f"tcp:{self._local_port}")

    async def _start_server(self) -> None:
        """Start scrcpy-server via ADB shell."""
        cmd = [
            "adb", "-H", self._adb_host, "-P", str(self._adb_port),
            "-s", self._serial, "shell",
            f"CLASSPATH={self._server_jar}",
            "app_process", "/",
            "com.genymobile.scrcpy.Server",
            self._scrcpy_version,
            f"max_fps={self._max_fps}",
            f"max_size={self._max_width}",
            f"video_bit_rate={self._bitrate}",
            "tunnel_forward=true",
            "audio=false",
            "control=false",
            "send_frame_meta=true",
        ]
        self._process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

    async def _read_loop(self) -> None:
        """Connect to scrcpy forwarded socket and read H.264 NAL units.

        scrcpy with ``send_frame_meta=true`` sends each frame as:
          - 8 bytes: PTS (int64 big-endian, microseconds)
          - 4 bytes: packet size (uint32 big-endian)
          - N bytes: H.264 NAL unit data
        """
        loop = asyncio.get_event_loop()
        sock: socket.socket | None = None
        try:
            for attempt in range(10):
                try:
                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    sock.setblocking(False)
                    await loop.sock_connect(sock, ("127.0.0.1", self._local_port))
                    break
                except (ConnectionRefusedError, OSError):
                    if sock:
                        sock.close()
                        sock = None
                    await asyncio.sleep(0.3 * (attempt + 1))
            else:
                logger.error("Failed to connect to scrcpy socket after 10 attempts")
                return

            device_name = await self._read_exactly(loop, sock, 64)
            logger.info("scrcpy device: %s", device_name.rstrip(b"\x00").decode(errors="replace"))

            while self._running:
                header = await self._read_exactly(loop, sock, 12)
                if len(header) < 12:
                    break
                _pts = struct.unpack(">q", header[:8])[0]
                pkt_size = struct.unpack(">I", header[8:12])[0]
                if pkt_size == 0 or pkt_size > 10 * 1024 * 1024:
                    continue
                nal_data = await self._read_exactly(loop, sock, pkt_size)
                if len(nal_data) < pkt_size:
                    break
                if self._video_track:
                    self._video_track.push_nal(nal_data)
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("scrcpy read loop error")
        finally:
            if sock:
                sock.close()

    @staticmethod
    async def _read_exactly(loop: asyncio.AbstractEventLoop, sock: socket.socket, n: int) -> bytes:
        buf = bytearray()
        while len(buf) < n:
            chunk = await loop.sock_recv(sock, n - len(buf))
            if not chunk:
                break
            buf.extend(chunk)
        return bytes(buf)

    def configure(self, max_fps: int | None = None, max_width: int | None = None,
                  bitrate: int | None = None) -> None:
        """Update stream parameters (takes effect on next restart)."""
        if max_fps is not None:
            self._max_fps = max_fps
        if max_width is not None:
            self._max_width = max_width
        if bitrate is not None:
            self._bitrate = bitrate
