from __future__ import annotations

import asyncio
import json
import logging
import struct
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Optional

import numpy as np
from aiortc import (
    RTCConfiguration,
    RTCIceServer,
    RTCPeerConnection,
    RTCSessionDescription,
)
from starlette.websockets import WebSocket, WebSocketState

from app.config import Settings
from app.engine import (
    FRAME_DIMENSION,
    MOTION_LAYOUT,
    GenerationRequest,
    GenerationResult,
    SignMotionEngine,
    split_sentences,
)

logger = logging.getLogger("realtime")

PROTOCOL_VERSION = 1
MOTION_BLOCK_HEADER = struct.Struct("<III")


def pack_motion_block(request_index: int, frame_offset: int, frames: np.ndarray) -> bytes:
    contiguous_frames = np.ascontiguousarray(frames, dtype="<f4")
    header = MOTION_BLOCK_HEADER.pack(request_index, frame_offset, int(contiguous_frames.shape[0]))
    return header + contiguous_frames.tobytes()


class Transport:
    async def send_text(self, text: str) -> None:
        raise NotImplementedError

    async def send_bytes(self, data: bytes) -> None:
        raise NotImplementedError

    async def wait_until_writable(self) -> None:
        return None

    def is_open(self) -> bool:
        raise NotImplementedError


class WebSocketTransport(Transport):
    def __init__(self, websocket: WebSocket) -> None:
        self.websocket = websocket

    async def send_text(self, text: str) -> None:
        await self.websocket.send_text(text)

    async def send_bytes(self, data: bytes) -> None:
        await self.websocket.send_bytes(data)

    def is_open(self) -> bool:
        client_connected = self.websocket.client_state == WebSocketState.CONNECTED
        application_connected = self.websocket.application_state == WebSocketState.CONNECTED
        return client_connected and application_connected


class DataChannelTransport(Transport):
    def __init__(self, channel: Any, high_water_mark: int) -> None:
        self.channel = channel
        self.high_water_mark = high_water_mark

    async def send_text(self, text: str) -> None:
        self.channel.send(text)

    async def send_bytes(self, data: bytes) -> None:
        self.channel.send(data)

    async def wait_until_writable(self) -> None:
        while self.channel.readyState == "open" and self.channel.bufferedAmount > self.high_water_mark:
            await asyncio.sleep(0.005)

    def is_open(self) -> bool:
        return self.channel.readyState == "open"


class ConnectionLimiter:
    def __init__(self, maximum: int) -> None:
        self.maximum = maximum
        self.active = 0

    def acquire(self) -> bool:
        if self.active >= self.maximum:
            return False
        self.active += 1
        return True

    def release(self) -> None:
        if self.active > 0:
            self.active -= 1


@dataclass
class RequestState:
    request_id: str
    request_index: int
    sentence_count: int
    completed_sentences: int = 0
    frame_offset: int = 0
    started_at: float = field(default_factory=time.perf_counter)
    cancelled: bool = False


@dataclass
class Job:
    session: "RealtimeSession"
    request_index: int
    sentence_index: int
    generation_request: GenerationRequest


@dataclass
class OutgoingItem:
    kind: str
    payload: Optional[dict] = None
    job: Optional[Job] = None
    result: Optional[GenerationResult] = None
    message: str = ""


class Scheduler:
    def __init__(self, engine: SignMotionEngine, settings: Settings) -> None:
        self.engine = engine
        self.settings = settings
        self.queue: asyncio.Queue = asyncio.Queue(maxsize=settings.queue_size)
        self.worker_task: Optional[asyncio.Task] = None

    async def start(self) -> None:
        self.worker_task = asyncio.create_task(self.run_worker())

    async def stop(self) -> None:
        if self.worker_task is None:
            return
        self.worker_task.cancel()
        try:
            await self.worker_task
        except asyncio.CancelledError:
            pass
        self.worker_task = None

    def free_slots(self) -> int:
        return self.settings.queue_size - self.queue.qsize()

    def submit_all(self, jobs: list) -> bool:
        if len(jobs) > self.free_slots():
            return False
        for job in jobs:
            self.queue.put_nowait(job)
        return True

    async def run_worker(self) -> None:
        while True:
            job = await self.queue.get()
            try:
                if job.session.should_skip(job.request_index):
                    continue
                result = await asyncio.to_thread(self.engine.generate, job.generation_request)
                job.session.enqueue_result(job, result)
            except asyncio.CancelledError:
                raise
            except Exception as error:
                logger.exception("generation failed")
                job.session.enqueue_failure(job, str(error))
            finally:
                self.queue.task_done()


class RealtimeSession:
    def __init__(
        self,
        transport: Transport,
        engine: SignMotionEngine,
        scheduler: Scheduler,
        settings: Settings,
        label: str,
    ) -> None:
        self.transport = transport
        self.engine = engine
        self.scheduler = scheduler
        self.settings = settings
        self.label = label
        self.session_id = uuid.uuid4().hex[:12]
        self.outbox: asyncio.Queue = asyncio.Queue()
        self.requests: dict = {}
        self.next_request_index = 1
        self.closed = False
        self.sender_task: Optional[asyncio.Task] = None

    def start(self) -> None:
        self.sender_task = asyncio.create_task(self.run_sender())
        self.enqueue_json(self.build_hello())

    def build_hello(self) -> dict:
        return {
            "type": "hello",
            "protocol_version": PROTOCOL_VERSION,
            "session_id": self.session_id,
            "transport": self.label,
            "ready": self.engine.ready,
            "fps": self.settings.fps,
            "frame_dimension": FRAME_DIMENSION,
            "layout": MOTION_LAYOUT,
            "binary_block_header": "uint32 request_index, uint32 frame_offset, uint32 frame_count, then float32 frames, all little endian",
            "limits": {
                "maximum_text_characters": self.settings.maximum_text_characters,
                "maximum_sentences_per_request": self.settings.maximum_sentences_per_request,
                "maximum_pending_requests": self.settings.maximum_pending_requests_per_connection,
                "frames_per_message": self.settings.frames_per_message,
            },
        }

    def enqueue_json(self, payload: dict) -> None:
        self.outbox.put_nowait(OutgoingItem(kind="json", payload=payload))

    def enqueue_error(self, request_id: Optional[str], code: str, message: str) -> None:
        self.enqueue_json(
            {"type": "error", "id": request_id, "code": code, "message": message}
        )

    def enqueue_result(self, job: Job, result: GenerationResult) -> None:
        if self.closed:
            return
        self.outbox.put_nowait(OutgoingItem(kind="sentence", job=job, result=result))

    def enqueue_failure(self, job: Job, message: str) -> None:
        if self.closed:
            return
        self.outbox.put_nowait(OutgoingItem(kind="failure", job=job, message=message))

    def should_skip(self, request_index: int) -> bool:
        if self.closed:
            return True
        state = self.requests.get(request_index)
        if state is None:
            return True
        return state.cancelled

    def handle_message(self, message: Any) -> None:
        if self.closed:
            return

        if isinstance(message, (bytes, bytearray)):
            self.enqueue_error(None, "invalid_message", "Binary messages are not accepted")
            return

        try:
            data = json.loads(message)
        except (TypeError, json.JSONDecodeError):
            self.enqueue_error(None, "invalid_json", "Message is not valid JSON")
            return

        if not isinstance(data, dict):
            self.enqueue_error(None, "invalid_message", "Message must be a JSON object")
            return

        message_type = data.get("type")

        if message_type == "generate":
            self.handle_generate(data)
        elif message_type == "cancel":
            self.handle_cancel(data)
        elif message_type == "ping":
            self.enqueue_json({"type": "pong", "time": data.get("time")})
        else:
            self.enqueue_error(None, "unknown_type", "Unknown message type")

    def handle_cancel(self, data: dict) -> None:
        request_id = str(data.get("id"))
        cancelled_indexes = []

        for request_index, state in self.requests.items():
            if state.request_id == request_id:
                state.cancelled = True
                cancelled_indexes.append(request_index)

        for request_index in cancelled_indexes:
            del self.requests[request_index]

        self.enqueue_json({"type": "cancelled", "id": request_id, "count": len(cancelled_indexes)})

    def handle_generate(self, data: dict) -> None:
        raw_request_id = data.get("id")
        if raw_request_id is None:
            request_id = uuid.uuid4().hex[:8]
        else:
            request_id = str(raw_request_id)

        if not self.engine.ready:
            self.enqueue_error(request_id, "loading", "The model is still loading")
            return

        text = data.get("text")
        if not isinstance(text, str) or text.strip() == "":
            self.enqueue_error(request_id, "invalid_text", "The text field must be a non empty string")
            return

        if len(text) > self.settings.maximum_text_characters:
            self.enqueue_error(
                request_id,
                "text_too_long",
                "The text is longer than " + str(self.settings.maximum_text_characters) + " characters",
            )
            return

        pace = data.get("pace", 1.0)
        pace_is_number = isinstance(pace, (int, float)) and not isinstance(pace, bool)
        if not pace_is_number or pace < 0.5 or pace > 2.0:
            self.enqueue_error(request_id, "invalid_pace", "The pace must be a number from 0.5 to 2.0")
            return

        use_retrieval = data.get("retrieval", self.settings.use_retrieval)
        if not isinstance(use_retrieval, bool):
            self.enqueue_error(request_id, "invalid_retrieval", "The retrieval field must be true or false")
            return

        do_sample = data.get("sample", self.settings.do_sample)
        if not isinstance(do_sample, bool):
            self.enqueue_error(request_id, "invalid_sample", "The sample field must be true or false")
            return

        seed = data.get("seed", self.settings.seed)
        if not isinstance(seed, int) or isinstance(seed, bool):
            self.enqueue_error(request_id, "invalid_seed", "The seed field must be an integer")
            return

        if len(self.requests) >= self.settings.maximum_pending_requests_per_connection:
            self.enqueue_error(request_id, "too_many_requests", "Too many unfinished requests on this connection")
            return

        sentences = split_sentences(text, self.settings.maximum_words_per_sentence)
        if len(sentences) == 0:
            self.enqueue_error(request_id, "invalid_text", "No sentences were found in the text")
            return

        if len(sentences) > self.settings.maximum_sentences_per_request:
            self.enqueue_error(request_id, "too_many_sentences", "The text has too many sentences")
            return

        request_index = self.next_request_index
        self.next_request_index += 1

        state = RequestState(
            request_id=request_id,
            request_index=request_index,
            sentence_count=len(sentences),
        )
        self.requests[request_index] = state

        jobs = []
        for sentence_index, sentence in enumerate(sentences):
            token_count = self.engine.plan_token_count(sentence, float(pace))
            generation_request = GenerationRequest(
                text=sentence,
                token_count=token_count,
                use_retrieval=use_retrieval,
                do_sample=do_sample,
                temperature=self.settings.temperature,
                top_k=self.settings.top_k,
                seed=seed,
            )
            jobs.append(
                Job(
                    session=self,
                    request_index=request_index,
                    sentence_index=sentence_index,
                    generation_request=generation_request,
                )
            )

        accepted = self.scheduler.submit_all(jobs)
        if not accepted:
            del self.requests[request_index]
            self.enqueue_error(request_id, "busy", "The server queue is full, try again shortly")
            return

        self.enqueue_json(
            {
                "type": "accepted",
                "id": request_id,
                "request_index": request_index,
                "sentences": sentences,
                "fps": self.settings.fps,
            }
        )

    async def run_sender(self) -> None:
        try:
            while True:
                item = await self.outbox.get()
                await self.dispatch(item)
        except asyncio.CancelledError:
            raise
        except Exception as error:
            logger.info("sender stopped: %s", error)
            self.closed = True

    async def dispatch(self, item: OutgoingItem) -> None:
        if not self.transport.is_open():
            raise ConnectionError("The transport is closed")

        if item.kind == "json":
            await self.transport.send_text(json.dumps(item.payload))
            return

        if item.kind == "failure":
            failed_job = item.job
            failed_state = self.requests.get(failed_job.request_index)
            if failed_state is None:
                return
            failed_state.cancelled = True
            del self.requests[failed_job.request_index]
            await self.transport.send_text(
                json.dumps(
                    {
                        "type": "error",
                        "id": failed_state.request_id,
                        "code": "generation_failed",
                        "message": item.message,
                    }
                )
            )
            return

        if item.kind == "sentence":
            await self.dispatch_sentence(item)

    async def dispatch_sentence(self, item: OutgoingItem) -> None:
        job = item.job
        result = item.result
        state = self.requests.get(job.request_index)

        if state is None or state.cancelled:
            return

        frames = result.motion
        frame_count = int(frames.shape[0])

        await self.transport.send_text(
            json.dumps(
                {
                    "type": "sentence",
                    "id": state.request_id,
                    "request_index": job.request_index,
                    "sentence_index": job.sentence_index,
                    "sentence_count": state.sentence_count,
                    "text": job.generation_request.text,
                    "frame_offset": state.frame_offset,
                    "frame_count": frame_count,
                    "fps": result.fps,
                    "seconds": frame_count / result.fps,
                    "generation_ms": round(result.elapsed_seconds * 1000.0, 1),
                    "cache_hit": result.from_cache,
                    "keywords": result.keywords,
                }
            )
        )

        frames_per_message = self.settings.frames_per_message
        for start_index in range(0, frame_count, frames_per_message):
            if state.cancelled:
                return
            block = frames[start_index:start_index + frames_per_message]
            await self.transport.wait_until_writable()
            await self.transport.send_bytes(
                pack_motion_block(job.request_index, state.frame_offset + start_index, block)
            )

        state.frame_offset += frame_count
        state.completed_sentences += 1

        if state.completed_sentences >= state.sentence_count:
            elapsed_milliseconds = (time.perf_counter() - state.started_at) * 1000.0
            await self.transport.send_text(
                json.dumps(
                    {
                        "type": "done",
                        "id": state.request_id,
                        "request_index": state.request_index,
                        "total_frames": state.frame_offset,
                        "total_seconds": state.frame_offset / result.fps,
                        "elapsed_ms": round(elapsed_milliseconds, 1),
                    }
                )
            )
            self.requests.pop(state.request_index, None)

    async def close(self) -> None:
        self.closed = True

        for state in self.requests.values():
            state.cancelled = True
        self.requests.clear()

        if self.sender_task is not None:
            self.sender_task.cancel()
            try:
                await self.sender_task
            except (asyncio.CancelledError, Exception):
                pass
            self.sender_task = None


class WebRTCManager:
    def __init__(
        self,
        settings: Settings,
        engine: SignMotionEngine,
        scheduler: Scheduler,
        limiter: ConnectionLimiter,
    ) -> None:
        self.settings = settings
        self.engine = engine
        self.scheduler = scheduler
        self.limiter = limiter
        self.peers: dict = {}

    def build_configuration(self) -> RTCConfiguration:
        ice_servers = []
        for entry in self.settings.ice_servers:
            ice_servers.append(
                RTCIceServer(
                    urls=entry["urls"],
                    username=entry.get("username"),
                    credential=entry.get("credential"),
                )
            )
        return RTCConfiguration(iceServers=ice_servers)

    async def handle_offer(self, sdp: str, offer_type: str) -> dict:
        if not self.limiter.acquire():
            raise ConnectionRefusedError("Too many active connections")

        peer = RTCPeerConnection(configuration=self.build_configuration())
        holder: dict = {"session": None, "watchdog": None}
        self.peers[peer] = holder

        @peer.on("datachannel")
        def on_datachannel(channel: Any) -> None:
            if holder["session"] is not None:
                channel.close()
                return

            transport = DataChannelTransport(channel, self.settings.data_channel_high_water_mark)
            session = RealtimeSession(transport, self.engine, self.scheduler, self.settings, "webrtc")
            holder["session"] = session

            @channel.on("message")
            def on_message(message: Any) -> None:
                session.handle_message(message)

            @channel.on("close")
            def on_close() -> None:
                asyncio.ensure_future(session.close())

            if channel.readyState == "open":
                session.start()
            else:

                @channel.on("open")
                def on_open() -> None:
                    session.start()

        @peer.on("connectionstatechange")
        async def on_connectionstatechange() -> None:
            if peer.connectionState in ("failed", "closed"):
                await self.release_peer(peer)

        holder["watchdog"] = asyncio.create_task(self.expire_if_idle(peer, holder))

        try:
            await peer.setRemoteDescription(RTCSessionDescription(sdp=sdp, type=offer_type))
            answer = await peer.createAnswer()
            await peer.setLocalDescription(answer)
        except Exception:
            await self.release_peer(peer)
            raise

        return {
            "sdp": peer.localDescription.sdp,
            "type": peer.localDescription.type,
        }

    async def expire_if_idle(self, peer: RTCPeerConnection, holder: dict) -> None:
        await asyncio.sleep(self.settings.idle_peer_timeout_seconds)
        if peer in self.peers and holder["session"] is None:
            await self.release_peer(peer)

    async def release_peer(self, peer: RTCPeerConnection) -> None:
        holder = self.peers.pop(peer, None)
        if holder is None:
            return

        session = holder["session"]
        if session is not None:
            await session.close()

        try:
            await peer.close()
        except Exception:
            logger.exception("closing peer failed")

        self.limiter.release()

    async def shutdown(self) -> None:
        for peer in list(self.peers.keys()):
            await self.release_peer(peer)