#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import asyncio
import json
import os
import ssl
import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import aiohttp
from aiohttp import web
from zeroconf import ServiceBrowser, Zeroconf

MDNS_SERVICE = "_yandexio._tcp.local."
QUASAR_TOKEN_URL = "https://quasar.yandex.net/glagol/token"
YANDEX_INFO_URL = "https://api.iot.yandex.net/v1.0/user/info"


@dataclass
class Station:
    device_id: str
    platform: str
    host: str
    port: int
    name: str = ""


class _Listener:
    def __init__(self):
        self.stations: Dict[str, Station] = {}

    def remove_service(self, zeroconf: Zeroconf, service_type: str, name: str) -> None:
        pass

    def add_service(self, zeroconf: Zeroconf, service_type: str, name: str) -> None:
        self._handle(zeroconf, service_type, name)

    def update_service(self, zeroconf: Zeroconf, service_type: str, name: str) -> None:
        self._handle(zeroconf, service_type, name)

    def _handle(self, zeroconf: Zeroconf, service_type: str, name: str) -> None:
        info = zeroconf.get_service_info(service_type, name)
        if not info:
            return

        props: Dict[str, Any] = {}
        for k, v in (info.properties or {}).items():
            try:
                kk = k.decode("utf-8") if isinstance(k, (bytes, bytearray)) else str(k)
                vv = v.decode("utf-8", errors="ignore") if isinstance(v, (bytes, bytearray)) else v
                props[kk] = vv
            except Exception:
                continue

        device_id = str(props.get("deviceId") or props.get("device_id") or "").strip()
        platform = str(props.get("platform") or "").strip()
        if not device_id or not platform:
            return

        try:
            addr = info.addresses[0]
            host = ".".join(str(b) for b in addr)
        except Exception:
            return

        port = int(info.port or 0)
        if not port:
            return

        name_friendly = str(props.get("name") or props.get("deviceName") or name).strip()

        self.stations[device_id] = Station(
            device_id=device_id,
            platform=platform,
            host=host,
            port=port,
            name=name_friendly,
        )


async def discover(timeout_sec: float = 2.5) -> List[Station]:
    zc = Zeroconf()
    listener = _Listener()
    browser = ServiceBrowser(zc, MDNS_SERVICE, listener)
    try:
        await asyncio.sleep(timeout_sec)
    finally:
        try:
            browser.cancel()
        except Exception:
            pass
        try:
            zc.close()
        except Exception:
            pass
    return list(listener.stations.values())


def _ssl_ctx() -> ssl.SSLContext:
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx


def _volume01_from_0_100(v: float) -> float:
    if v < 0 or v > 100:
        raise RuntimeError("volume must be 0..100")
    return round(v / 100.0, 3)


def pick_station(stations: List[Station], selector: str) -> Station:
    selector = (selector or "").strip()
    if not selector:
        raise RuntimeError("empty selector")

    for s in stations:
        if s.device_id == selector:
            return s
    for s in stations:
        if s.host == selector or f"{s.host}:{s.port}" == selector:
            return s
    low = selector.lower()
    for s in stations:
        if low in (s.name or "").lower():
            return s
    raise RuntimeError(f"station not found by selector: {selector}")


async def get_conversation_token(session: aiohttp.ClientSession, oauth_token: str, station: Station) -> str:
    params = {"device_id": station.device_id, "platform": station.platform}
    headers = {"Authorization": f"oauth {oauth_token}"}
    async with session.get(QUASAR_TOKEN_URL, params=params, headers=headers, timeout=aiohttp.ClientTimeout(total=10)) as r:
        text = await r.text()
        try:
            data = json.loads(text)
        except Exception:
            raise RuntimeError(f"Quasar token: bad json: {text[:200]}")
        if data.get("status") != "ok" or not data.get("token"):
            raise RuntimeError(f"Quasar token: unexpected response: {data}")
        return str(data["token"])


async def ws_send(session: aiohttp.ClientSession, station: Station, conv_token: str, payload: dict, wait_sec: float = 5.0) -> dict:
    url = f"wss://{station.host}:{station.port}"
    req_id = str(uuid.uuid4())
    msg = {
        "conversationToken": conv_token,
        "id": req_id,
        "payload": payload,
        "sentTime": int(round(time.time() * 1000)),
    }

    async with session.ws_connect(url, ssl=_ssl_ctx(), heartbeat=55, timeout=aiohttp.ClientTimeout(total=10)) as ws:
        await ws.send_json(msg)

        deadline = time.time() + wait_sec
        while time.time() < deadline:
            in_msg = await ws.receive(timeout=wait_sec)
            if in_msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(in_msg.data)
                except Exception:
                    continue
                if data.get("requestId") == req_id or data.get("id") == req_id:
                    return data
            elif in_msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                break

    return {"status": "timeout"}


async def do_list(args) -> int:
    stations = await discover(args.timeout)
    for s in stations:
        print(f"{s.name}\tdevice_id={s.device_id}\tplatform={s.platform}\t{s.host}:{s.port}")
    return 0


async def do_stations_cloud(args) -> int:
    oauth = args.token or os.environ.get("YANDEX_TOKEN", "").strip()
    if not oauth:
        raise RuntimeError("No token. Use --token or env YANDEX_TOKEN")

    async with aiohttp.ClientSession() as session:
        headers = {"Authorization": f"OAuth {oauth}"}
        async with session.get(YANDEX_INFO_URL, headers=headers, timeout=aiohttp.ClientTimeout(total=15)) as r:
            data = await r.json(content_type=None)

    rooms = {}
    for rm in (data.get("rooms") or []):
        rid = str(rm.get("id") or "")
        if rid:
            rooms[rid] = str(rm.get("name") or "")

    out = []
    for d in (data.get("devices") or []):
        tp = str(d.get("type") or "").lower()
        if "speaker" not in tp and "smart_speaker" not in tp:
            continue
        rid = str(d.get("room") or "")
        out.append({
            "name": str(d.get("name") or ""),
            "id": str(d.get("id") or ""),
            "room": rooms.get(rid, ""),
            "model": str(d.get("model") or d.get("type") or ""),
        })

    print(json.dumps(out, ensure_ascii=False, indent=2))
    return 0


class StationConn:
    def __init__(self, station: Station, oauth_token: str, wait_sec: float):
        self.station = station
        self.oauth_token = oauth_token
        self.wait_sec = wait_sec
        self._session: Optional[aiohttp.ClientSession] = None
        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self._conv_token: str = ""
        self._conv_exp_ts: float = 0.0
        self._lock = asyncio.Lock()
        self._playing: Optional[int] = None
        self._volume: Optional[int] = None
        self._last_probe_ts: float = 0.0

    @property
    def playing(self) -> Optional[int]:
        if self._playing is None:
            return None
        return 1 if self._playing else 0

    @property
    def volume(self) -> Optional[int]:
        if self._volume is None:
            return None
        v = int(self._volume)
        if v < 0:
            return 0
        if v > 100:
            return 100
        return v

    @staticmethod
    def _to_playing_flag(value: Any) -> Optional[int]:
        if isinstance(value, bool):
            return 1 if value else 0
        if isinstance(value, (int, float)):
            v = int(round(float(value)))
            if v == 2:
                return 1
            if v in (0, 1):
                return v
        if isinstance(value, str):
            low = value.strip().lower()
            if low in {"playing", "play", "started", "start", "on", "true", "yes"}:
                return 1
            if low in {"paused", "pause", "stopped", "stop", "idle", "off", "false", "no"}:
                return 0
            try:
                vv = int(float(low))
                if vv == 2:
                    return 1
                if vv in (0, 1):
                    return vv
            except Exception:
                return None
        return None

    @classmethod
    def _infer_playing(cls, obj: Any) -> Optional[int]:
        if isinstance(obj, dict):
            for k in ("playing", "isPlaying", "playState", "playerState", "state"):
                if k in obj:
                    out = cls._to_playing_flag(obj.get(k))
                    if out is not None:
                        return out
            for v in obj.values():
                out = cls._infer_playing(v)
                if out is not None:
                    return out
        elif isinstance(obj, list):
            for v in obj:
                out = cls._infer_playing(v)
                if out is not None:
                    return out
        return None

    @staticmethod
    def _to_volume_percent(value: Any) -> Optional[int]:
        if isinstance(value, bool):
            return 100 if value else 0
        if isinstance(value, (int, float)):
            v = float(value)
            if 0.0 <= v <= 1.0:
                v = v * 100.0
            iv = int(round(v))
            if iv < 0:
                iv = 0
            if iv > 100:
                iv = 100
            return iv
        if isinstance(value, str):
            s = value.strip().replace(",", ".")
            if not s:
                return None
            try:
                return StationConn._to_volume_percent(float(s))
            except Exception:
                return None
        return None

    @classmethod
    def _infer_volume(cls, obj: Any) -> Optional[int]:
        if isinstance(obj, dict):
            for k in ("volume", "currentVolume", "volumeLevel", "level"):
                if k in obj:
                    out = cls._to_volume_percent(obj.get(k))
                    if out is not None:
                        return out
            for v in obj.values():
                out = cls._infer_volume(v)
                if out is not None:
                    return out
        elif isinstance(obj, list):
            for v in obj:
                out = cls._infer_volume(v)
                if out is not None:
                    return out
        return None

    async def start(self):
        if self._session is None:
            self._session = aiohttp.ClientSession()
        await self._ensure_ws()

    async def close(self):
        try:
            if self._ws is not None:
                await self._ws.close()
        except Exception:
            pass
        self._ws = None
        try:
            if self._session is not None:
                await self._session.close()
        except Exception:
            pass
        self._session = None

    async def _ensure_conv_token(self):
        now = time.time()
        if self._conv_token and self._conv_exp_ts > now:
            return
        assert self._session is not None
        self._conv_token = await get_conversation_token(self._session, self.oauth_token, self.station)
        self._conv_exp_ts = now + 60.0

    async def _ensure_ws(self):
        assert self._session is not None
        await self._ensure_conv_token()

        if self._ws is not None and not self._ws.closed:
            return

        url = f"wss://{self.station.host}:{self.station.port}"
        self._ws = await self._session.ws_connect(
            url,
            ssl=_ssl_ctx(),
            heartbeat=55,
            timeout=aiohttp.ClientTimeout(total=10),
        )

    async def send(self, payload: dict) -> dict:
        async with self._lock:
            await self._ensure_ws()

            req_id = str(uuid.uuid4())
            cmd = str(payload.get("command") or "").strip().lower()
            sent_volume = None
            if cmd == "setvolume":
                sent_volume = self._to_volume_percent(payload.get("volume"))
                if sent_volume is not None:
                    # Optimistic update: status endpoint can immediately expose requested level.
                    self._volume = sent_volume
            msg = {
                "conversationToken": self._conv_token,
                "id": req_id,
                "payload": payload,
                "sentTime": int(round(time.time() * 1000)),
            }

            try:
                assert self._ws is not None
                await self._ws.send_json(msg)
            except Exception:
                try:
                    if self._ws is not None:
                        await self._ws.close()
                except Exception:
                    pass
                self._ws = None
                await self._ensure_ws()
                assert self._ws is not None
                await self._ws.send_json(msg)

            deadline = time.time() + self.wait_sec
            while time.time() < deadline:
                try:
                    in_msg = await self._ws.receive(timeout=self.wait_sec)
                except Exception:
                    break
                if in_msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(in_msg.data)
                    except Exception:
                        continue
                    if data.get("requestId") == req_id or data.get("id") == req_id:
                        inferred = self._infer_playing(data)
                        if inferred is not None:
                            self._playing = inferred
                        elif cmd == "play":
                            self._playing = 1
                        elif cmd in {"pause", "stop"}:
                            self._playing = 0
                        inferred_volume = self._infer_volume(data)
                        if inferred_volume is not None:
                            self._volume = inferred_volume
                        elif sent_volume is not None:
                            self._volume = sent_volume
                        return data
                elif in_msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break

            if sent_volume is not None:
                self._volume = sent_volume
            return {"status": "timeout", "id": req_id}

    async def probe_state(self) -> None:
        now = time.time()
        if now - self._last_probe_ts < 5.0:
            return
        self._last_probe_ts = now

        candidates = [
            {"command": "getState"},
            {"command": "getStatus"},
            {"command": "status"},
            {"command": "playerState"},
        ]
        for payload in candidates:
            try:
                data = await self.send(payload)
            except Exception:
                continue
            if self._infer_playing(data) is not None or self._infer_volume(data) is not None:
                break


async def serve_daemon(args) -> int:
    oauth = args.token or os.environ.get("YANDEX_TOKEN", "").strip()
    if not oauth:
        raise RuntimeError("No token. Use --token or env YANDEX_TOKEN")

    stations = await discover(args.timeout)
    selected: List[Station] = []

    if args.stations:
        for sel in args.stations.split(","):
            selected.append(pick_station(stations, sel.strip()))
    else:
        selected = stations

    if not selected:
        raise RuntimeError("No stations found")

    conns: Dict[str, StationConn] = {}
    for st in selected:
        conns[st.device_id] = StationConn(st, oauth, args.wait)
        await conns[st.device_id].start()

    async def resolve_conn(station_sel: str) -> Optional[StationConn]:
        station_sel = (station_sel or "").strip()
        if not station_sel:
            return None

        conn = conns.get(station_sel)
        if conn is not None:
            return conn

        low = station_sel.lower()
        for c in conns.values():
            if c.station.device_id == station_sel or low in (c.station.name or "").lower():
                return c
            if c.station.host == station_sel or f"{c.station.host}:{c.station.port}" == station_sel:
                return c
        return None

    async def execute_action(req: Dict[str, Any]) -> Dict[str, Any]:
        station_sel = str(req.get("station") or "").strip()
        if not station_sel:
            return {"ok": False, "error": "station is required"}

        action = str(req.get("action") or "").strip().lower()
        conn = await resolve_conn(station_sel)
        if conn is None:
            return {"ok": False, "error": "station not found"}

        try:
            if action == "tts":
                text = str(req.get("text") or "").strip()
                if not text:
                    raise RuntimeError("text is required")
                volume = req.get("volume")
                if volume is not None and str(volume).strip() != "":
                    vol_payload = {"command": "setVolume", "volume": _volume01_from_0_100(float(volume))}
                    await conn.send(vol_payload)
                payload = {
                    "command": "serverAction",
                    "serverActionEventPayload": {
                        "type": "server_action",
                        "name": "update_form",
                        "payload": {
                            "form_update": {
                                "name": "personal_assistant.scenarios.repeat_after_me",
                                "slots": [{"type": "string", "name": "request", "value": text}],
                            },
                            "resubmit": True,
                        },
                    },
                }
            elif action == "command":
                text = str(req.get("text") or "").strip()
                if not text:
                    raise RuntimeError("text is required")
                payload = {"command": "sendText", "text": text}
            elif action == "player":
                cmd = str(req.get("cmd") or "").strip().lower()
                if cmd not in {"play", "pause", "stop", "next", "prev"}:
                    raise RuntimeError("cmd must be play|pause|stop|next|prev")
                payload = {"command": cmd}
            elif action == "volume":
                lvl = float(req.get("level"))
                payload = {"command": "setVolume", "volume": _volume01_from_0_100(lvl)}
            elif action == "raw":
                payload = req.get("payload")
                if isinstance(payload, list):
                    if not payload:
                        raise RuntimeError("payload list is empty")
                    results = []
                    for idx, p in enumerate(payload):
                        if not isinstance(p, dict):
                            raise RuntimeError(f"payload[{idx}] must be object")
                        results.append(await conn.send(p))
                    return {"ok": True, "result": results}
                if not isinstance(payload, dict):
                    raise RuntimeError("payload must be object or array of objects")
            else:
                raise RuntimeError("action must be tts|command|player|volume|raw")

            result = await conn.send(payload)
            return {"ok": True, "result": result}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        try:
            line = await reader.readline()
            if not line:
                writer.close()
                return
            req = json.loads(line.decode("utf-8", errors="ignore"))
        except Exception as e:
            writer.write((json.dumps({"ok": False, "error": f"bad request: {e}"}, ensure_ascii=False) + "\n").encode("utf-8"))
            await writer.drain()
            writer.close()
            return
        resp = await execute_action(req)

        writer.write((json.dumps(resp, ensure_ascii=False) + "\n").encode("utf-8"))
        await writer.drain()
        writer.close()

    server = await asyncio.start_server(handle_client, host=args.host, port=args.port)
    addrs = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
    print(f"tcp server listening on {addrs}")

    app = web.Application()

    async def http_exec(request: web.Request) -> web.Response:
        data: Dict[str, Any] = {}
        if request.method == "POST":
            if request.content_type and "application/json" in request.content_type:
                body = await request.json()
                if isinstance(body, dict):
                    data.update(body)
            else:
                body = await request.post()
                data.update({k: v for k, v in body.items()})
        data.update({k: v for k, v in request.query.items()})

        if "payload" in data and isinstance(data["payload"], str):
            payload_raw = str(data["payload"]).strip()
            if payload_raw:
                try:
                    data["payload"] = json.loads(payload_raw)
                except Exception:
                    return web.json_response({"ok": False, "error": "bad payload json"}, status=400)

        result = await execute_action(data)
        return web.json_response(result, status=200 if result.get("ok") else 422)

    async def http_status(request: web.Request) -> web.Response:
        station_sel = str(request.query.get("station") or "").strip()
        conn = await resolve_conn(station_sel)
        if conn is None:
            return web.json_response({"ok": False, "error": "station not found"}, status=404)
        # Keep persistent WS connection: try to restore it on each status poll.
        try:
            await conn.start()
        except Exception:
            pass
        try:
            await conn.probe_state()
        except Exception:
            pass
        running = 1 if (conn._ws is not None and not conn._ws.closed) else 0
        return web.json_response({
            "ok": 1,
            "running": running,
            "playing": conn.playing,
            "volume": conn.volume,
            "station": conn.station.device_id,
        })

    async def http_health(_request: web.Request) -> web.Response:
        return web.json_response({"ok": 1})

    app.router.add_get("/health", http_health)
    app.router.add_get("/api/status", http_status)
    app.router.add_get("/api/exec", http_exec)
    app.router.add_post("/api/exec", http_exec)

    http_runner = web.AppRunner(app)
    await http_runner.setup()
    http_site = web.TCPSite(http_runner, host=args.http_host, port=args.http_port)
    await http_site.start()
    print(f"http server listening on {(args.http_host, args.http_port)}")

    try:
        async with server:
            await server.serve_forever()
    finally:
        try:
            await http_runner.cleanup()
        except Exception:
            pass
        for c in conns.values():
            await c.close()


async def client_call(args) -> int:
    req = {"station": args.station, "action": args.action}
    if args.action in {"tts", "command"}:
        req["text"] = args.text
    elif args.action == "player":
        req["cmd"] = args.cmd
    elif args.action == "volume":
        req["level"] = float(args.level)
    elif args.action == "raw":
        req["payload"] = json.loads(args.payload)

    reader, writer = await asyncio.open_connection(args.host, args.port)
    writer.write((json.dumps(req, ensure_ascii=False) + "\n").encode("utf-8"))
    await writer.drain()
    line = await reader.readline()
    writer.close()

    if not line:
        print(json.dumps({"ok": False, "error": "no response"}, ensure_ascii=False))
        return 2

    print(line.decode("utf-8", errors="ignore").rstrip("\n"))
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="yastation.py", description="Yandex station helper")
    p.add_argument("--timeout", type=float, default=2.5)
    p.add_argument("--token", type=str, default="")
    p.add_argument("--wait", type=float, default=5.0)

    sub = p.add_subparsers(dest="action", required=True)

    sp = sub.add_parser("list")
    sp.set_defaults(fn=do_list)

    sp = sub.add_parser("stations-cloud")
    sp.set_defaults(fn=do_stations_cloud)

    sp = sub.add_parser("serve")
    sp.add_argument("--host", default="127.0.0.1")
    sp.add_argument("--port", type=int, default=9123)
    sp.add_argument("--http-host", default="0.0.0.0")
    sp.add_argument("--http-port", type=int, default=9124)
    sp.add_argument("--stations", default="")
    sp.set_defaults(fn=serve_daemon)

    sp = sub.add_parser("client")
    sp.add_argument("--host", default="127.0.0.1")
    sp.add_argument("--port", type=int, default=9123)
    sp.add_argument("--station", required=True)
    sp.add_argument("action", help="tts|command|player|volume|raw")
    sp.add_argument("--text", default="")
    sp.add_argument("--cmd", default="")
    sp.add_argument("--level", default="0")
    sp.add_argument("--payload", default="{}")
    sp.set_defaults(fn=client_call)

    return p


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    try:
        return asyncio.run(args.fn(args))
    except KeyboardInterrupt:
        return 130


if __name__ == "__main__":
    raise SystemExit(main())