"""Asynchronous session manager backed by Docker containers."""

from __future__ import annotations

import asyncio
import logging
import socket
from contextlib import suppress
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Dict, Iterable, Optional

import docker
import httpx
from docker.errors import APIError, NotFound
from docker.models.containers import Container

from .config import Settings, settings
from .models import (
    DATASET_LABEL,
    EXPIRY_LABEL,
    MANAGED_LABEL,
    SESSION_ID_LABEL,
    SessionCreateRequest,
    SessionRecord,
    default_labels,
)

logger = logging.getLogger(__name__)


class SessionError(RuntimeError):
    """Base class for orchestration failures."""


class SessionExistsError(SessionError):
    """Raised when a duplicate session identifier is requested."""


class SessionNotFoundError(SessionError):
    """Raised when attempting to operate on a missing session."""


class DatasetDownloadError(SessionError):
    """Raised when the dataset cannot be fetched."""


class PortUnavailableError(SessionError):
    """Raised when the requested port is already in use."""


class SessionManager:
    """Coordinate embedding-atlas containers per incoming request."""

    def __init__(self, app_settings: Settings | None = None) -> None:
        self.settings = app_settings or settings
        self.client = docker.DockerClient(base_url=self.settings.docker_url)
        self._sessions: Dict[str, SessionRecord] = {}
        self._lock = asyncio.Lock()
        self._cleanup_task: Optional[asyncio.Task[None]] = None

    async def start(self) -> None:
        """Initialise caches and background cleanup."""

        await self._load_existing_sessions()
        if self._cleanup_task is None:
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def stop(self) -> None:
        """Stop background tasks and close Docker handles."""

        if self._cleanup_task is not None:
            self._cleanup_task.cancel()
            with suppress(asyncio.CancelledError):
                await self._cleanup_task
            self._cleanup_task = None
        await asyncio.to_thread(self.client.close)

    async def create_session(self, payload: SessionCreateRequest) -> SessionRecord:
        """Download the dataset and spin up an embedding-atlas container."""

        session_id = payload.normalised_session_id()
        async with self._lock:
            if session_id in self._sessions:
                raise SessionExistsError(f"Session {session_id} already active.")

        host = payload.host or self.settings.container_host
        if payload.port is not None:
            is_free = await asyncio.to_thread(self._port_available, payload.port)
            if not is_free:
                raise PortUnavailableError(f"Requested port {payload.port} is unavailable")
            port = payload.port
        else:
            port = await asyncio.to_thread(self._allocate_port)
        expiry_seconds = payload.auto_remove_seconds or self.settings.auto_remove_seconds
        expires_at = datetime.now(tz=timezone.utc) + timedelta(seconds=expiry_seconds)
        dataset_path = await self._download_dataset(session_id, payload)

        try:
            record = await asyncio.to_thread(
                self._run_container,
                session_id,
                dataset_path,
                payload,
                host,
                port,
                expires_at,
            )
        except Exception:
            # Drop downloaded artefacts if startup failed
            with suppress(Exception):
                self._purge_dataset(dataset_path)
            raise

        async with self._lock:
            self._sessions[session_id] = record
        return record

    async def close_session(self, session_id: str) -> bool:
        """Terminate the associated container and remove cached data."""

        record = await self._get_session(session_id)
        await asyncio.to_thread(self._stop_container, record.container_name)
        await asyncio.to_thread(self._purge_dataset, Path(record.dataset_path))
        async with self._lock:
            self._sessions.pop(session_id, None)
        return True

    async def list_sessions(self) -> Iterable[SessionRecord]:
        """Return the internal session registry."""

        async with self._lock:
            return list(self._sessions.values())

    async def cleanup_expired(self) -> None:
        """Stop containers that exceeded their TTL."""

        now = datetime.now(tz=timezone.utc)
        async with self._lock:
            sessions = list(self._sessions.values())
        for record in sessions:
            if record.expires_at <= now:
                logger.info("Session %s expired; stopping container %s", record.session_id, record.container_name)
                with suppress(SessionError, APIError, NotFound):
                    await self.close_session(record.session_id)

    async def _cleanup_loop(self) -> None:
        """Periodic cleanup coroutine."""

        try:
            while True:
                await asyncio.sleep(self.settings.cleanup_interval_seconds)
                await self.cleanup_expired()
        except asyncio.CancelledError:
            logger.debug("Cleanup loop cancelled")
            raise

    async def _get_session(self, session_id: str) -> SessionRecord:
        async with self._lock:
            record = self._sessions.get(session_id)
        if record is None:
            raise SessionNotFoundError(f"Session {session_id} not found")
        return record

    async def _load_existing_sessions(self) -> None:
        """Recover managed containers into the local cache on startup."""

        def _list() -> list[Container]:
            return self.client.containers.list(all=True, filters={"label": MANAGED_LABEL})

        containers = await asyncio.to_thread(_list)
        restored: Dict[str, SessionRecord] = {}
        for container in containers:
            labels = container.labels or {}
            session_id = labels.get(SESSION_ID_LABEL)
            if not session_id:
                continue
            try:
                port = self._extract_port(container)
            except ValueError:
                logger.warning("Unable to restore port mapping for container %s", container.name)
                continue
            expires_at = self._parse_expiry(labels.get(EXPIRY_LABEL))
            dataset_path = labels.get(DATASET_LABEL, "")
            restored[session_id] = SessionRecord(
                session_id=session_id,
                container_id=container.id,
                container_name=container.name,
                port=port,
                host=self.settings.container_host,
                started_at=self._parse_created(container.attrs.get("Created")),
                expires_at=expires_at,
                dataset_path=dataset_path,
            )
        async with self._lock:
            self._sessions.update(restored)
        if restored:
            logger.info("Restored %d managed sessions from existing containers", len(restored))

    def _run_container(
        self,
        session_id: str,
        dataset_path: Path,
        payload: SessionCreateRequest,
        host: str,
        port: int,
        expires_at: datetime,
    ) -> SessionRecord:
        container_name = f"{self.settings.container_name_prefix}_{session_id}"
        self._guard_duplicate_container(container_name)

        bind_root = dataset_path.parent
        container_dataset_path = Path("/session") / dataset_path.name

        command = [
            "embedding-atlas",
            str(container_dataset_path),
            "--host",
            host,
            "--port",
            str(port),
            "--no-auto-port",
        ]
        if payload.extra_args:
            command.extend(payload.extra_args)

        env = payload.environment.copy()
        env.setdefault("SESSION_ID", session_id)

        labels = default_labels(session_id, expires_at, str(dataset_path))
        labels.update(payload.labels)

        logger.info(
            "Starting embedding-atlas container %s on port %s (session %s)",
            container_name,
            port,
            session_id,
        )

        port_bindings = {f"{port}/tcp": port}
        container = self.client.containers.run(
            payload.image or self.settings.container_image,
            command,
            name=container_name,
            detach=True,
            environment=env,
            labels=labels,
            volumes={
                str(bind_root): {"bind": "/session", "mode": "ro"}
            },
            ports=port_bindings,
            network=self.settings.container_network,
        )

        return SessionRecord(
            session_id=session_id,
            container_id=container.id,
            container_name=container.name,
            port=port,
            host=host,
            started_at=datetime.now(tz=timezone.utc),
            expires_at=expires_at,
            dataset_path=str(dataset_path),
        )

    def _stop_container(self, container_name: str) -> None:
        try:
            container = self.client.containers.get(container_name)
        except NotFound as exc:
            raise SessionNotFoundError(f"Container {container_name} not found") from exc
        container.stop(timeout=30)
        container.remove(v=True)
        logger.info("Removed container %s", container_name)

    def _purge_dataset(self, dataset_path: Path) -> None:
        try:
            dataset_path.unlink(missing_ok=True)
            parent = dataset_path.parent
            if parent.exists() and not any(parent.iterdir()):
                parent.rmdir()
        except Exception as exc:
            logger.warning("Failed to clean dataset path %s: %s", dataset_path, exc)

    async def _download_dataset(self, session_id: str, payload: SessionCreateRequest) -> Path:
        filename = payload.input_filename or Path(payload.data_url.path).name
        if not filename:
            filename = f"dataset-{session_id}"
        session_dir = self.settings.session_root / session_id
        session_dir.mkdir(parents=True, exist_ok=True)
        target = session_dir / filename

        logger.info("Downloading dataset for session %s from %s", session_id, payload.data_url)
        timeout = httpx.Timeout(self.settings.download_timeout_seconds)
        async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
            try:
                async with client.stream("GET", str(payload.data_url)) as response:
                    response.raise_for_status()
                    with target.open("wb") as file_handle:
                        async for chunk in response.aiter_bytes(chunk_size=self.settings.download_chunk_size):
                            file_handle.write(chunk)
            except Exception as exc:
                raise DatasetDownloadError(f"Failed to fetch dataset: {exc}") from exc
        return target

    def _allocate_port(self) -> int:
        if self.settings.port_range_start >= self.settings.port_range_end:
            return self._random_free_port()
        candidates = range(self.settings.port_range_start, self.settings.port_range_end + 1)
        for candidate in candidates:
            if self._port_available(candidate):
                return candidate
        return self._random_free_port()

    @staticmethod
    def _random_free_port() -> int:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.bind(("", 0))
            return sock.getsockname()[1]

    @staticmethod
    def _port_available(port: int) -> bool:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            result = sock.connect_ex(("127.0.0.1", port))
            return result != 0

    def _guard_duplicate_container(self, container_name: str) -> None:
        try:
            container = self.client.containers.get(container_name)
        except NotFound:
            return
        raise SessionExistsError(
            f"Container name {container_name} already exists (id={container.id})."
        )

    @staticmethod
    def _extract_port(container: Container) -> int:
        ports = container.attrs.get("NetworkSettings", {}).get("Ports", {})
        if not ports:
            raise ValueError("No exposed ports found")
        # Expect single port binding
        (container_port, host_bindings), = ports.items()
        if not host_bindings:
            raise ValueError("Container port not published")
        host_port = host_bindings[0].get("HostPort")
        if host_port is None:
            raise ValueError("Missing host port")
        return int(host_port)

    @staticmethod
    def _parse_expiry(value: Optional[str]) -> datetime:
        if not value:
            return datetime.now(tz=timezone.utc)
        try:
            return datetime.fromtimestamp(int(value), tz=timezone.utc)
        except ValueError:
            return datetime.now(tz=timezone.utc)

    @staticmethod
    def _parse_created(value: Optional[str]) -> datetime:
        if not value:
            return datetime.now(tz=timezone.utc)
        try:
            return datetime.fromisoformat(value.replace('Z', '+00:00'))
        except ValueError:
            return datetime.now(tz=timezone.utc)


__all__ = [
    "SessionManager",
    "SessionError",
    "SessionExistsError",
    "SessionNotFoundError",
    "DatasetDownloadError",
    "PortUnavailableError",
]
