"""
PushStreamManager - RTSP/RTMP push stream management via ffmpeg.

Python 3.8 compatible, zero-transcode (-c copy) by default.
"""

import asyncio
import dataclasses
import json
import os
import re
import time
import uuid
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any, Dict, List, Optional


def _set_low_priority():
    """Run after fork to lower child process priority."""
    try:
        os.nice(19)
    except Exception:
        pass


@dataclass
class PushConfig:
    """Configuration and state for a single push stream."""

    # Identity / behaviour
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    enabled: bool = True
    auto_start: bool = False
    stop_on_flight: bool = True

    # Stream endpoint
    source_url: str = ""
    server_host: str = ""
    server_port: int = 554
    stream_key: str = ""
    protocol: str = "rtsp"
    codec: str = "copy"
    extra_args: List[str] = field(default_factory=list)

    # Persistent stats (survive restarts)
    total_push_count: int = 0
    total_push_seconds: float = 0.0
    total_bytes_sent: int = 0

    # Runtime fields (reset on load)
    start_timestamp: Optional[float] = None
    last_bytes_sent: Optional[int] = None
    current_bandwidth_kbps: Optional[float] = None
    pid: Optional[int] = None
    status: str = "idle"          # idle | pushing | error | blocked
    error: Optional[str] = None
    retry_count: int = 0


class PushStreamManager:
    """Manage multiple ffmpeg push streams with flight-safety awareness."""

    CONFIG_FILE = Path("push_streams.json")

    def __init__(self, config_file: Optional[Path] = None):
        self.config_file = (
            config_file or Path(__file__).parent / "push_streams.json"
        )
        self.configs: Dict[str, PushConfig] = {}
        self._processes: Dict[str, asyncio.subprocess.Process] = {}
        self._tasks: Dict[str, asyncio.Task] = {}

        self._flight_monitor_task: Optional[asyncio.Task] = None
        self._stop_flight_monitor: asyncio.Event = asyncio.Event()

        self.global_auto_push: bool = True
        self._lock = asyncio.Lock()

        self._load()

    # ------------------------------------------------------------------
    # Persistence helpers
    # ------------------------------------------------------------------
    def _load(self) -> None:
        if not self.config_file.exists():
            return
        try:
            with open(self.config_file, "r", encoding="utf-8") as fh:
                raw = json.load(fh)
        except Exception:
            return

        known = {f.name for f in dataclasses.fields(PushConfig)}
        for item in raw:
            if not isinstance(item, dict):
                continue
            filtered = {k: v for k, v in item.items() if k in known}
            if "extra_args" in filtered and filtered["extra_args"] is None:
                filtered["extra_args"] = []
            try:
                cfg = PushConfig(**filtered)
            except Exception:
                continue
            # Reset runtime fields
            cfg.start_timestamp = None
            cfg.last_bytes_sent = None
            cfg.current_bandwidth_kbps = None
            cfg.pid = None
            cfg.status = "idle"
            cfg.error = None
            cfg.retry_count = 0
            self.configs[cfg.id] = cfg

    def _save(self) -> None:
        try:
            self.config_file.parent.mkdir(parents=True, exist_ok=True)
            payload = [asdict(c) for c in self.configs.values()]
            with open(self.config_file, "w", encoding="utf-8") as fh:
                json.dump(payload, fh, indent=2, default=str)
                os.fsync(fh.fileno())
            os.sync()
        except Exception:
            pass

    # ------------------------------------------------------------------
    # CRUD
    # ------------------------------------------------------------------
    async def list_configs(self) -> List[PushConfig]:
        async with self._lock:
            return list(self.configs.values())

    async def get_config(self, config_id: str) -> Optional[PushConfig]:
        async with self._lock:
            return self.configs.get(config_id)

    async def create_config(self, **kwargs: Any) -> PushConfig:
        async with self._lock:
            cfg = PushConfig(**kwargs)
            self.configs[cfg.id] = cfg
            self._save()
            return cfg

    async def update_config(self, config_id: str, **kwargs: Any) -> PushConfig:
        async with self._lock:
            cfg = self.configs.get(config_id)
            if cfg is None:
                raise ValueError(f"Config {config_id} not found")
            for key, value in kwargs.items():
                if hasattr(cfg, key):
                    setattr(cfg, key, value)
            self._save()
            return cfg

    async def delete_config(self, config_id: str) -> None:
        async with self._lock:
            cfg = self.configs.get(config_id)
            if cfg is None:
                return
            if config_id in self._processes:
                await self._do_stop(config_id)
            del self.configs[config_id]
            self._save()

    # ------------------------------------------------------------------
    # Global settings
    # ------------------------------------------------------------------
    def get_global_auto_push(self) -> bool:
        return self.global_auto_push

    async def set_global_auto_push(self, enabled: bool) -> None:
        async with self._lock:
            self.global_auto_push = enabled

    # ------------------------------------------------------------------
    # Start / stop
    # ------------------------------------------------------------------
    async def start(self, config_id: str) -> PushConfig:
        async with self._lock:
            cfg = self.configs.get(config_id)
            if cfg is None:
                raise ValueError(f"Config {config_id} not found")
            if not cfg.enabled:
                raise ValueError(f"Config {config_id} is disabled")
            if cfg.status == "pushing":
                return cfg

            # Flight safety
            if cfg.stop_on_flight:
                if os.path.exists("/tmp/arm.lock") or os.path.exists("/tmp/flight.lock"):
                    cfg.status = "blocked"
                    cfg.error = "Flight lock active"
                    return cfg

            output_url = self._build_output_url(cfg)
            cmd = [
                "ffmpeg",
                "-re",
                "-i", cfg.source_url,
                "-c", cfg.codec,
            ] + cfg.extra_args + [output_url]

            try:
                proc = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.DEVNULL,
                    stderr=asyncio.subprocess.PIPE,
                    preexec_fn=_set_low_priority,
                )
            except Exception as exc:
                cfg.status = "error"
                cfg.error = str(exc)
                self._save()
                raise

            cfg.pid = proc.pid
            cfg.status = "pushing"
            cfg.error = None
            cfg.start_timestamp = time.time()
            cfg.total_push_count += 1
            cfg.retry_count = 0

            self._processes[cfg.id] = proc
            task = asyncio.create_task(self._stats_reader(cfg.id, proc))
            self._tasks[cfg.id] = task
            self._save()
            return cfg

    async def stop(self, config_id: str) -> PushConfig:
        async with self._lock:
            return await self._do_stop(config_id)

    async def _do_stop(self, config_id: str) -> PushConfig:
        cfg = self.configs.get(config_id)
        if cfg is None:
            raise ValueError(f"Config {config_id} not found")

        proc = self._processes.pop(config_id, None)
        task = self._tasks.pop(config_id, None)

        if proc is not None and proc.returncode is None:
            try:
                proc.terminate()
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
            except Exception:
                pass

        if task is not None and not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
            except Exception:
                pass

        # Accumulate stats
        if cfg.start_timestamp is not None:
            cfg.total_push_seconds += time.time() - cfg.start_timestamp
        if cfg.last_bytes_sent is not None:
            cfg.total_bytes_sent += cfg.last_bytes_sent

        # Determine final status
        if proc is not None and proc.returncode not in (0, None, -15, -9, -2):
            cfg.status = "error"
            cfg.error = f"ffmpeg exited with code {proc.returncode}"
        elif cfg.status not in ("error",):
            cfg.status = "idle"

        cfg.start_timestamp = None
        cfg.last_bytes_sent = None
        cfg.current_bandwidth_kbps = None
        cfg.pid = None
        self._save()
        return cfg

    async def _cleanup_after_exit(self, config_id: str) -> None:
        async with self._lock:
            if config_id not in self._processes:
                return
            try:
                await self._do_stop(config_id)
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Stats reader
    # ------------------------------------------------------------------
    async def _stats_reader(
        self, config_id: str, proc: asyncio.subprocess.Process
    ) -> None:
        cfg = self.configs.get(config_id)
        if cfg is None or proc.stderr is None:
            return

        try:
            while True:
                line = await proc.stderr.readline()
                if not line:
                    break
                line_str = line.decode("utf-8", errors="ignore")

                size_match = re.search(
                    r"size=\s*([\d.]+)\s*([kMGT]?)B", line_str
                )
                if size_match:
                    val = float(size_match.group(1))
                    unit = size_match.group(2)
                    size_mult = {"": 1, "k": 1000, "M": 1000 ** 2, "G": 1000 ** 3, "T": 1000 ** 4}
                    cfg.last_bytes_sent = int(val * size_mult.get(unit, 1))

                bitrate_match = re.search(
                    r"bitrate=\s*([\d.]+)\s*([kMGT]?)bits/s", line_str
                )
                if bitrate_match:
                    val = float(bitrate_match.group(1))
                    unit = bitrate_match.group(2)
                    bw_mult = {"": 0.001, "k": 1.0, "M": 1000.0, "G": 1000000.0}
                    cfg.current_bandwidth_kbps = val * bw_mult.get(unit, 1.0)
        except asyncio.CancelledError:
            raise
        except Exception:
            pass
        finally:
            if config_id in self._processes:
                asyncio.create_task(self._cleanup_after_exit(config_id))

    # ------------------------------------------------------------------
    # Lifespan helpers
    # ------------------------------------------------------------------
    async def start_all_auto(self) -> None:
        async with self._lock:
            ids_to_start = [
                c.id
                for c in self.configs.values()
                if c.enabled
                and c.auto_start
                and self.global_auto_push
                and c.status == "idle"
            ]
        for cid in ids_to_start:
            try:
                await self.start(cid)
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Flight monitor
    # ------------------------------------------------------------------
    async def start_flight_monitor(self) -> None:
        async with self._lock:
            if (
                self._flight_monitor_task is not None
                and not self._flight_monitor_task.done()
            ):
                return
            self._stop_flight_monitor.clear()
            self._flight_monitor_task = asyncio.create_task(
                self._flight_monitor()
            )

    async def stop_flight_monitor(self) -> None:
        async with self._lock:
            if (
                self._flight_monitor_task is None
                or self._flight_monitor_task.done()
            ):
                return
            self._stop_flight_monitor.set()
            self._flight_monitor_task.cancel()
            try:
                await self._flight_monitor_task
            except asyncio.CancelledError:
                pass
            self._flight_monitor_task = None

    async def _flight_monitor(self) -> None:
        while not self._stop_flight_monitor.is_set():
            try:
                is_flying = os.path.exists("/tmp/arm.lock") or os.path.exists(
                    "/tmp/flight.lock"
                )

                actions: List[tuple] = []
                async with self._lock:
                    for cfg in list(self.configs.values()):
                        if is_flying and cfg.stop_on_flight and cfg.status == "pushing":
                            actions.append(("stop", cfg.id))
                        elif (
                            not is_flying
                            and cfg.auto_start
                            and cfg.enabled
                            and self.global_auto_push
                            and cfg.status in ("idle", "blocked", "error")
                            and cfg.retry_count < 5
                        ):
                            actions.append(("start", cfg.id))

                for action, cid in actions:
                    if action == "stop":
                        try:
                            await self.stop(cid)
                        except Exception:
                            pass
                    else:
                        try:
                            await self.start(cid)
                        except Exception:
                            async with self._lock:
                                c = self.configs.get(cid)
                                if c is not None:
                                    c.retry_count += 1

                await asyncio.wait_for(
                    self._stop_flight_monitor.wait(), timeout=5.0
                )
            except asyncio.TimeoutError:
                pass
            except asyncio.CancelledError:
                raise
            except Exception:
                await asyncio.sleep(5.0)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    def _build_output_url(self, cfg: PushConfig) -> str:
        return (
            f"{cfg.protocol}://{cfg.server_host}:{cfg.server_port}/{cfg.stream_key}"
        )
