from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import Any

import websockets

from .sync import StorageAgent


class StorageAgentTunnel:
    def __init__(self, agent: StorageAgent):
        self.agent = agent
        self.config = agent.config
        self.logger = agent.logger

    async def run_forever(self) -> None:
        self.logger.info("Starting tunnel mode for node=%s", self.config.node_code)
        while True:
            try:
                await self._connect_once()
            except KeyboardInterrupt:
                raise
            except Exception as exc:
                self.logger.exception("Tunnel disconnected/error: %s", exc)
            await asyncio.sleep(max(5, self.config.tunnel_reconnect_seconds))

    async def _connect_once(self) -> None:
        async with websockets.connect(
            self.config.tunnel_url,
            ping_interval=30,
            ping_timeout=30,
            close_timeout=10,
            max_size=1024 * 1024,
        ) as websocket:
            await websocket.send(self._json_text({"type": "ready", "agent_time_utc": self._now()}))
            self.agent._maybe_heartbeat(force=True)
            async for raw in websocket:
                message = self._parse_json(raw)
                if not isinstance(message, dict):
                    continue
                msg_type = message.get("type")
                if msg_type == "connected":
                    self.logger.info("Tunnel connected: %s", message.get("node_code"))
                    continue
                if msg_type == "sync":
                    jobs = message.get("jobs") or []
                    requests_ = message.get("download_requests") or message.get("requests") or []
                    await asyncio.to_thread(self.agent.process_storage_jobs_from_payload, jobs)
                    await asyncio.to_thread(self.agent.process_download_requests_from_payload, requests_)
                    await websocket.send(self._json_text({
                        "type": "agent_status",
                        "jobs_seen": len(jobs),
                        "requests_seen": len(requests_),
                        "agent_time_utc": self._now(),
                    }))
                    continue
                if msg_type == "error":
                    self.logger.error("Tunnel server error: %s", message.get("error"))

    def _now(self) -> str:
        return datetime.now(timezone.utc).isoformat()

    def _json_text(self, payload: dict[str, Any]) -> str:
        import json
        return json.dumps(payload, separators=(",", ":"))

    def _parse_json(self, raw: str) -> Any:
        import json
        try:
            return json.loads(raw)
        except Exception:
            return None
