from __future__ import annotations

import os
import json
import struct
from datetime import datetime, timezone
from pathlib import Path
import threading
import time
from typing import Any


PROJECT_ROOT = Path(__file__).resolve().parents[2]
WORKSPACE_ROOT = PROJECT_ROOT.parent
JPEG_FRAME_DIR = WORKSPACE_ROOT / "RobotDataShow" / "jpeg-frames"
OMNISOCKET_CONFIG_PATH = PROJECT_ROOT / "config" / "omnisocket_demo.yaml"

VIDEO_SOURCE_MODE = os.getenv("VIDEO_SOURCE_MODE", "auto").strip().lower()
OMNISOCKET_FRAME_FRESH_SECONDS = 2.0
VIDEO_TIMESTAMP_SAMPLE_SIZE = 10
VIDEO_TRAILER_ENDIANNESS = "little"
VIDEO_TRAILER_TIMESTAMP_UNIT = "ms"
VIDEO_TRAILER_TIMESTAMP_MULTIPLIER_NS = 1_000_000
VIDEO_TRAILER_TIMESTAMP_MAX_SKEW_NS = 7 * 24 * 60 * 60 * 1_000_000_000
VIDEO_TRAILER_COORDINATE_FORMAT = (
    "uint64 timestamp_ms + float64 latitude + float64 longitude + uint32 capture_to_send_ms (little-endian)"
)
VIDEO_TRAILER_STRUCT = struct.Struct("<QddI")
VIDEO_TRAILER_BYTES = VIDEO_TRAILER_STRUCT.size

CONTROL_PACKET = struct.Struct("<6f")
CONTROL_PACKET_SIZE = CONTROL_PACKET.size
CONTROL_SOURCE_NATIVE_UDP = "native_udp"
CONTROL_SOURCE_WEB = "web"
CONTROL_SOURCE_PRIORITY = (CONTROL_SOURCE_NATIVE_UDP, CONTROL_SOURCE_WEB)
ZERO_CONTROL_PAYLOAD = CONTROL_PACKET.pack(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
BLITZ_RUN_DIR_RAW = os.getenv("BLITZ_RUN_DIR", "").strip()
BLITZ_RUN_DIR = Path(BLITZ_RUN_DIR_RAW).expanduser() if BLITZ_RUN_DIR_RAW else None
BLITZ_INSTANCE_ID = os.getenv("BLITZ_INSTANCE_ID", "").strip() or f"backend-{os.getpid()}"


def utc_iso_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")


def parse_simple_yaml_scalar(value: str) -> Any:
    if value in {'""', "''"}:
        return ""
    if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
        return value[1:-1]
    if value.lower() == "true":
        return True
    if value.lower() == "false":
        return False
    if value and value.lstrip("-").isdigit():
        return int(value)
    return value


def load_simple_yaml_config(path: Path) -> dict[str, Any]:
    parsed: dict[str, Any] = {}
    current_section: str | None = None

    with path.open("r", encoding="utf-8") as file:
        for raw_line in file:
            line = raw_line.split("#", 1)[0].rstrip()
            if not line.strip():
                continue

            if not line.startswith(" "):
                if not line.endswith(":"):
                    raise ValueError(f"invalid top-level yaml line: {raw_line.strip()}")
                current_section = line[:-1].strip()
                parsed[current_section] = {}
                continue

            if current_section is None:
                raise ValueError(f"yaml key outside section: {raw_line.strip()}")

            stripped = line.strip()
            if ":" not in stripped:
                raise ValueError(f"invalid yaml key line: {raw_line.strip()}")

            key, value = stripped.split(":", 1)
            parsed[current_section][key.strip()] = parse_simple_yaml_scalar(value.strip())

    return parsed


def load_omnisocket_config() -> dict[str, Any]:
    config: dict[str, Any] = {}
    if OMNISOCKET_CONFIG_PATH.exists():
        try:
            try:
                import yaml  # type: ignore

                with OMNISOCKET_CONFIG_PATH.open("r", encoding="utf-8") as file:
                    config = yaml.safe_load(file) or {}
            except ImportError:
                config = load_simple_yaml_config(OMNISOCKET_CONFIG_PATH)
        except Exception:
            config = {}

    transport_cfg = dict(config.get("transport", {}))
    video_receiver_cfg = dict(config.get("video_receiver", {}))
    control_sender_cfg = dict(config.get("control_sender", {}))
    control_ack_receiver_cfg = dict(config.get("control_ack_receiver", {}))
    control_ingress_cfg = dict(config.get("control_ingress", {}))
    video_sender_cfg = dict(config.get("video_sender", {}))
    telemetry_receiver_cfg = dict(config.get("telemetry_receiver", {}))

    transport_cfg["server_addr"] = os.getenv(
        "OMNISOCKET_SERVER_ADDR",
        str(transport_cfg.get("server_addr", "127.0.0.1:10909")),
    )
    transport_cfg["relay_via"] = os.getenv(
        "OMNISOCKET_RELAY_VIA",
        str(transport_cfg.get("relay_via", "")),
    )
    transport_cfg["bind_ip"] = os.getenv(
        "OMNISOCKET_BIND_IP",
        str(transport_cfg.get("bind_ip", "")),
    )
    transport_cfg["bind_device"] = os.getenv(
        "OMNISOCKET_BIND_DEVICE",
        str(transport_cfg.get("bind_device", "")),
    )

    video_receiver_cfg["peer_id"] = os.getenv(
        "OMNISOCKET_VIDEO_PEER_ID",
        str(video_receiver_cfg.get("peer_id", "peer-a-video")),
    )
    video_receiver_cfg["buffer_bytes"] = int(
        os.getenv(
            "OMNISOCKET_BUFFER_BYTES",
            str(video_receiver_cfg.get("buffer_bytes", 1024 * 1024)),
        )
    )

    control_sender_cfg["peer_id"] = os.getenv(
        "OMNISOCKET_CONTROL_PEER_ID",
        str(control_sender_cfg.get("peer_id", "peer-a-ctrl")),
    )
    control_sender_cfg["target_peer"] = os.getenv(
        "OMNISOCKET_CONTROL_TARGET_PEER",
        str(control_sender_cfg.get("target_peer", "peer-b-ctrl")),
    )

    control_ack_receiver_cfg["peer_id"] = os.getenv(
        "OMNISOCKET_CONTROL_ACK_RECEIVER_PEER_ID",
        str(control_ack_receiver_cfg.get("peer_id", "peer-a-ctrl-ack")),
    )
    control_ack_receiver_cfg["expected_sender"] = os.getenv(
        "OMNISOCKET_CONTROL_ACK_EXPECTED_SENDER",
        str(control_ack_receiver_cfg.get("expected_sender", "peer-b-ctrl-ack")),
    )

    video_sender_cfg["peer_id"] = os.getenv(
        "OMNISOCKET_VIDEO_SENDER_PEER_ID",
        str(video_sender_cfg.get("peer_id", "peer-b-video")),
    )
    video_sender_cfg["target_peer"] = os.getenv(
        "OMNISOCKET_VIDEO_TARGET_PEER_ID",
        str(video_sender_cfg.get("target_peer", "peer-a-video")),
    )

    control_ingress_cfg["native_udp_bind"] = os.getenv(
        "OMNISOCKET_CONTROL_NATIVE_UDP_BIND",
        str(control_ingress_cfg.get("native_udp_bind", "127.0.0.1:10921")),
    )
    control_ingress_cfg["source_lease_ms"] = int(
        os.getenv(
            "OMNISOCKET_CONTROL_SOURCE_LEASE_MS",
            str(control_ingress_cfg.get("source_lease_ms", 300)),
        )
    )
    control_ingress_cfg["send_rate_hz"] = float(
        os.getenv(
            "OMNISOCKET_CONTROL_SEND_RATE_HZ",
            str(control_ingress_cfg.get("send_rate_hz", 20.0)),
        )
    )
    control_ingress_cfg["zero_burst_packets"] = int(
        os.getenv(
            "OMNISOCKET_CONTROL_ZERO_BURST_PACKETS",
            str(control_ingress_cfg.get("zero_burst_packets", 3)),
        )
    )

    telemetry_receiver_cfg["peer_id"] = os.getenv(
        "OMNISOCKET_TELEMETRY_PEER_ID",
        str(telemetry_receiver_cfg.get("peer_id", "peer-a-telemetry")),
    )
    telemetry_receiver_cfg["interval_ms"] = int(
        os.getenv(
            "OMNISOCKET_TELEMETRY_INTERVAL_MS",
            str(telemetry_receiver_cfg.get("interval_ms", 500)),
        )
    )
    telemetry_receiver_cfg["stale_after_ms"] = int(
        os.getenv(
            "OMNISOCKET_TELEMETRY_STALE_AFTER_MS",
            str(telemetry_receiver_cfg.get("stale_after_ms", telemetry_receiver_cfg["interval_ms"] * 3)),
        )
    )

    return {
        "transport": transport_cfg,
        "video_receiver": video_receiver_cfg,
        "control_sender": control_sender_cfg,
        "control_ack_receiver": control_ack_receiver_cfg,
        "control_ingress": control_ingress_cfg,
        "video_sender": video_sender_cfg,
        "telemetry_receiver": telemetry_receiver_cfg,
    }


class JsonlRunLogger:
    def __init__(self, stem_env: str, default_stem: str) -> None:
        explicit_path = os.getenv(stem_env, "").strip()
        self._path = Path(explicit_path) if explicit_path else (
            BLITZ_RUN_DIR / f"{default_stem}.{BLITZ_INSTANCE_ID}.jsonl" if BLITZ_RUN_DIR is not None else None
        )
        self._lock = threading.Lock()
        self._file = None
        self._buffered_bytes = 0
        self._current_bytes = 0
        self._flush_bytes = self._positive_int_env("BLITZ_JSONL_FLUSH_BYTES", 262144)
        self._flush_interval_ms = self._positive_int_env("BLITZ_JSONL_FLUSH_INTERVAL_MS", 1000)
        self._max_bytes = self._positive_int_env("BLITZ_JSONL_ROTATE_BYTES", 134217728)
        self._max_files = self._positive_int_env("BLITZ_JSONL_ROTATE_FILES", 8)
        self._last_flush_monotonic_ms = self._now_ms()

        if self._path is not None:
            try:
                self._path.parent.mkdir(parents=True, exist_ok=True)
                self._file = self._path.open("a", encoding="utf-8")
                self._current_bytes = self._path.stat().st_size if self._path.exists() else 0
            except OSError:
                self._file = None

    @property
    def path(self) -> str | None:
        return str(self._path) if self._path is not None else None

    @property
    def enabled(self) -> bool:
        return self._file is not None

    def write(self, payload: dict[str, Any]) -> None:
        if self._file is None:
            return
        line = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
        line_bytes = len(line.encode("utf-8")) + 1
        with self._lock:
            if self._file is None:
                return
            try:
                self._file.write(line)
                self._file.write("\n")
                self._buffered_bytes += line_bytes
                self._current_bytes += line_bytes
                now_ms = self._now_ms()
                if (
                    self._buffered_bytes >= self._flush_bytes
                    or (self._flush_interval_ms > 0 and now_ms - self._last_flush_monotonic_ms >= self._flush_interval_ms)
                ):
                    self._flush_locked(now_ms)
                if self._max_bytes > 0 and self._max_files > 0 and self._current_bytes >= self._max_bytes:
                    self._rotate_locked()
            except OSError:
                if self._file is not None:
                    try:
                        self._file.close()
                    except OSError:
                        pass
                    self._file = None

    def close(self) -> None:
        with self._lock:
            if self._file is not None:
                try:
                    self._flush_locked(self._now_ms())
                except OSError:
                    pass
                self._file.close()
                self._file = None

    def _flush_locked(self, now_ms: int) -> None:
        if self._file is None:
            return
        self._file.flush()
        self._buffered_bytes = 0
        self._last_flush_monotonic_ms = now_ms

    def _rotate_locked(self) -> None:
        if self._path is None or self._file is None or self._max_files <= 0:
            return
        self._flush_locked(self._now_ms())
        self._file.close()
        self._file = None

        oldest = self._path.with_name(f"{self._path.name}.{self._max_files}")
        if oldest.exists():
            oldest.unlink()

        for index in range(self._max_files - 1, 0, -1):
            src = self._path.with_name(f"{self._path.name}.{index}")
            if src.exists():
                dst = self._path.with_name(f"{self._path.name}.{index + 1}")
                src.replace(dst)

        if self._path.exists():
            rotated = self._path.with_name(f"{self._path.name}.1")
            self._path.replace(rotated)

        self._file = self._path.open("a", encoding="utf-8")
        self._buffered_bytes = 0
        self._current_bytes = self._path.stat().st_size if self._path.exists() else 0
        self._last_flush_monotonic_ms = self._now_ms()

    @staticmethod
    def _now_ms() -> int:
        return int(time.monotonic() * 1000)

    @staticmethod
    def _positive_int_env(name: str, default: int) -> int:
        raw = os.getenv(name, "").strip()
        try:
            value = int(raw)
        except ValueError:
            return default
        return value if value > 0 else default


def parse_host_port(bind_addr: str) -> tuple[str, int]:
    host, port_text = bind_addr.rsplit(":", 1)
    host = host.strip() or "127.0.0.1"
    port = int(port_text)
    if port <= 0 or port > 65535:
        raise ValueError(f"invalid port in bind address: {bind_addr}")
    return host, port