#!/usr/bin/env python3

import asyncio
import contextlib
import contextvars
import datetime
import functools
import ipaddress
import itertools
import logging
import os
import pathlib
import random
import string
from dataclasses import dataclass, field

import asgineer

import not_my_board._auth as auth
import not_my_board._http as http
import not_my_board._jsonrpc as jsonrpc
import not_my_board._metrics as metrics
import not_my_board._models as models
import not_my_board._util as util

logger = logging.getLogger(__name__)
client_ip_var = contextvars.ContextVar("client_ip")
connection_id_var = contextvars.ContextVar("connection_id")
authenticator_var = contextvars.ContextVar("authenticator")


def run_hub():
    import socket  # noqa: PLC0415

    import uvicorn  # noqa: PLC0415

    with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
        s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)

        host = "::"
        port = 2092
        s.bind((host, port))

        s.listen()
        print("ready", flush=True)  # noqa: T201

        fd = s.fileno()
        uvicorn.main([f"--fd={fd}", __name__ + ":asgi_app"])


async def asgi_app(scope, receive, send):
    if scope["type"] == "lifespan":
        # asgineer doesn't expose the lifespan hooks. Handle them here
        # before handing over to asgineer
        await _handle_lifespan(scope, receive, send)
    elif scope["type"] == "http" and scope["path"] == "/metrics":
        await metrics.asgi_app(scope, receive, send)
    else:
        await _handle_request(scope, receive, send)


async def _handle_lifespan(scope, receive, send):
    while True:
        message = await receive()
        if message["type"] == "lifespan.startup":
            try:
                hub = await _on_startup()
                scope["state"]["hub"] = hub
            except Exception as err:
                await send({"type": "lifespan.startup.failed", "message": str(err)})
            else:
                await send({"type": "lifespan.startup.complete"})
        elif message["type"] == "lifespan.shutdown":
            try:
                await hub.shutdown()
            except Exception as err:
                await send({"type": "lifespan.shutdown.failed", "message": str(err)})
            else:
                await send({"type": "lifespan.shutdown.complete"})
            return
        else:
            logger.warning("Unknown lifespan message %s", message["type"])


async def _on_startup():
    config_file = os.environ.get("NOT_MY_BOARD_HUB_CONFIG")
    if not config_file:
        config_file = "/etc/not-my-board/hub.toml"
    config_file = pathlib.Path(config_file)

    if config_file.exists():
        config = util.toml_loads(config_file.read_text())
    else:
        logger.warning('Config file "%s" not found', config_file)
        config = {}

    hub = Hub(config, http.Client())
    await hub.startup()

    return hub


@asgineer.to_asgi
async def _handle_request(request):
    util.generate_log_request_id()
    hub = request.scope["state"]["hub"]
    response = (404, {}, "Page not found")

    if isinstance(request, asgineer.WebsocketRequest):
        if request.path == "/ws":
            await _handle_websocket(hub, request)
        else:
            await request.close()
        response = None
    elif isinstance(request, asgineer.HttpRequest):
        if request.path == "/api/v1/places":
            response = await hub.get_places()
        elif request.path == "/api/v1/auth-info":
            response = hub.auth_info()
        elif request.path == "/oidc-callback":
            response = await hub.oidc_callback(request.querydict)
    return response


async def _handle_websocket(hub, ws):
    await ws.accept()
    client_ip = ws.scope["client"][0]
    channel = jsonrpc.Channel(ws.send, ws.receive_iter())
    await hub.communicate(client_ip, channel)


def require_role(role):
    def decorator(func):
        @functools.wraps(func)
        async def require_role_wrapper(*args, **kwargs):
            authenticator = authenticator_var.get()
            await authenticator.require_role(role)
            return await func(*args, **kwargs)

        return require_role_wrapper

    return decorator


class Hub:
    def __init__(self, config=None, http_client=None):
        self._places = {}
        self._place_names = set()
        self._place_stats = {}
        self._exporters = {}
        self._available = set()
        self._wait_queue = []
        self._reservations = {}
        self._authenticators = {}
        self._pending_callbacks = {}

        if config is None:
            config = {}

        if "log_level" in config:
            log_level_str = config["log_level"]
            log_level_map = {
                "debug": logging.DEBUG,
                "info": logging.INFO,
                "warning": logging.WARNING,
                "error": logging.ERROR,
            }
            log_level = log_level_map[log_level_str]
            util.configure_logging(log_level)

        self._uniqueness_tolerance = config.get("uniqueness_tolerance", 0.01)

        auth_config = config.get("auth")
        if auth_config:
            trusted_issuers = {auth_config["issuer"]}
            for permission in auth_config["permissions"]:
                issuer = permission["claims"].get("iss")
                if issuer:
                    trusted_issuers.add(issuer)

            def make_issuer_config(issuer):
                issuer_config = auth_config.get("issuers", {}).get(issuer, {})
                return IssuerConfig.from_dict(issuer_config)

            self._issuer_configs = {
                issuer: make_issuer_config(issuer) for issuer in trusted_issuers
            }

            keys = {"issuer", "client_id"}
            self._auth_info = {k: auth_config[k] for k in keys}
            issuer_config = self._issuer_configs[auth_config["issuer"]]
            if issuer_config.show_claims is not None:
                self._auth_info["show_claims"] = issuer_config.show_claims

            def make_permission(d):
                d["claims"].setdefault("iss", auth_config["issuer"])
                return Permission.from_dict(d)

            self._permissions = [make_permission(d) for d in auth_config["permissions"]]
            self._validator = auth.Validator(
                auth_config["client_id"], http_client, trusted_issuers
            )
        else:
            logger.warning("Authentication is disabled")
            self._auth_info = {}
            self._permissions = []
            self._validator = None
            self._issuer_configs = {}

        self._id_generator = itertools.count(start=1)

        metrics.PLACES_REGISTERED.set_function(lambda: len(self._places))
        metrics.PLACES_RESERVED.set_function(
            lambda: len(self._places) - len(self._available)
        )
        metrics.RESERVATION_QUEUE_LENGTH.set_function(lambda: len(self._wait_queue))

    @jsonrpc.hidden
    async def startup(self):
        pass

    @jsonrpc.hidden
    async def shutdown(self):
        pass

    @jsonrpc.hidden
    async def get_places(self):
        return {"places": [p.dict() for p in self._places.values()]}

    @jsonrpc.hidden
    async def communicate(self, client_ip, channel):
        client_ip_var.set(client_ip)
        async with self._connection_context(channel):
            channel.set_api_object(self)
            await channel.communicate_forever()

    @contextlib.asynccontextmanager
    async def _connection_context(self, channel):
        id_ = next(self._id_generator)
        connection_id_var.set(id_)
        self._reservations[id_] = set()
        authenticator = Authenticator(
            self._permissions, self._validator, channel, self._issuer_configs
        )

        authenticator_var.set(authenticator)
        self._authenticators[id_] = authenticator

        try:
            async with authenticator:
                yield
        finally:
            if id_ in self._places:
                name = self._places[id_].name
                logger.info("Place @%s disappeared", name)
                self._place_names.discard(name)
                del self._place_stats[id_]
                del self._places[id_]
                del self._exporters[id_]
                self._available.discard(id_)
                for candidates, _, future in self._wait_queue:
                    candidates.discard(id_)
                    if not candidates and not future.done():
                        future.set_exception(
                            RuntimeError("All candidate places are gone")
                        )

            coros = [self.return_reservation(id_) for id_ in self._reservations[id_]]
            results = await asyncio.gather(*coros, return_exceptions=True)
            del self._reservations[id_]
            del self._authenticators[id_]
            for result in results:
                if isinstance(result, Exception):
                    logger.warning("Error while deregistering agent: %s", result)

    @require_role("exporter")
    async def register_place(self, export_desc):
        id_ = connection_id_var.get()
        client_ip = client_ip_var.get()
        place = models.Place.init_recursive(
            id=id_, host=_unmap_ip(client_ip), **export_desc
        )

        if id_ in self._places:
            raise RuntimeError("Place already registered")

        if place.name in self._place_names:
            raise RuntimeError(f'Place with name "{place.name}" already registered')

        self._places[id_] = place
        self._place_stats[id_] = PlaceStats()
        self._place_names.add(place.name)
        self._exporters[id_] = jsonrpc.get_current_channel()
        self._available.add(id_)
        logger.info("New place registered: @%s", place.name)
        return id_

    @require_role("importer")
    async def reserve(self, candidate_ids):
        id_ = connection_id_var.get()
        existing_candidates = {c_id for c_id in candidate_ids if c_id in self._places}
        if not existing_candidates:
            raise RuntimeError("None of the candidates exist anymore")

        # update uniqueness: If a place is requested with many alternative
        # candidates we assume it is a common place. If only a few or no
        # alternative candidates are requested, then we assume it is a unique
        # place. This uniqueness score is only used if more than one candidate
        # is still available. Then we reserve the most common available place
        # to keep the unique places free as long as possible.
        uniqueness = 1 - (len(existing_candidates) / len(self._places))
        for c_id in existing_candidates:
            self._place_stats[c_id].uniqueness.update(uniqueness)

        available_candidates = existing_candidates & self._available
        if available_candidates:
            # Select the most common places, not just the single most common
            # one, to avoid that small uniqueness differences lead to always
            # reserving the same place. Configured by uniqueness_tolerance.
            min_uniqueness = min(
                float(self._place_stats[c_id].uniqueness)
                for c_id in available_candidates
            )
            common_candidates = [
                c_id
                for c_id in available_candidates
                if float(self._place_stats[c_id].uniqueness) - min_uniqueness
                <= self._uniqueness_tolerance
            ]

            reserved_id = random.choice(common_candidates)  # noqa: S311

            self._available.remove(reserved_id)
            self._reservations[id_].add(reserved_id)

            user_name = self._format_user_name()
            logger.info(
                "Place @%s reserved by %s", self._places[reserved_id].name, user_name
            )
        else:
            logger.debug(
                "No places available, adding request to queue: %s",
                str(existing_candidates),
            )
            future = asyncio.get_running_loop().create_future()
            entry = (existing_candidates, id_, future)
            self._wait_queue.append(entry)
            try:
                reserved_id = await future
            finally:
                self._wait_queue.remove(entry)

        client_ip = client_ip_var.get()
        async with util.on_error(self.return_reservation, reserved_id):
            exporter = self._exporters[reserved_id]
            await exporter.set_allowed_ips([_unmap_ip(client_ip)])

        return reserved_id

    @require_role("importer")
    async def return_reservation(self, place_id):
        id_ = connection_id_var.get()
        self._reservations[id_].remove(place_id)
        if place_id in self._places:
            name = self._places[place_id].name
            returned_by = self._format_user_name()

            for candidates, agent_id, future in self._wait_queue:
                if place_id in candidates and not future.done():
                    self._reservations[agent_id].add(place_id)

                    reserved_by = self._format_user_name(agent_id)
                    logger.info(
                        "Place @%s returned by %s. Now reserved by %s",
                        name,
                        returned_by,
                        reserved_by,
                    )
                    future.set_result(place_id)
                    break
            else:
                logger.info("Place @%s returned by %s", name, returned_by)
                self._available.add(place_id)
                rpc = self._exporters[place_id]
                await rpc.set_allowed_ips([])
        else:
            logger.debug("Place %d returned, but it doesn't exist", place_id)

    @require_role("importer")
    async def get_reservations(self):
        def reservations():
            for agent_id, place_ids in self._reservations.items():
                for place_id in place_ids:
                    yield {
                        "place_id": place_id,
                        "user": self._format_user_name(agent_id),
                    }

        return {"reservations": list(reservations())}

    @require_role("importer")
    async def get_place_stats(self, place_id):
        return {
            "uniqueness": float(self._place_stats[place_id].uniqueness),
        }

    def _format_user_name(self, agent_id=None):
        if agent_id is None:
            agent_id = connection_id_var.get()

        return self._authenticators[agent_id].get_user_name() or f"<{agent_id}>"

    async def get_authentication_response(self, state):
        future = asyncio.get_running_loop().create_future()
        self._pending_callbacks[state] = future
        try:
            channel = jsonrpc.get_current_channel()
            await channel.oidc_callback_registered(_notification=True)
            return await future
        finally:
            del self._pending_callbacks[state]

    @jsonrpc.hidden
    async def oidc_callback(self, query):
        future = self._pending_callbacks[query["state"]]
        if not future.done():
            future.set_result(query)

        return "Continue in not-my-board CLI"

    @jsonrpc.hidden
    def auth_info(self):
        return self._auth_info


class Authenticator(util.ContextStack):
    _leeway = datetime.timedelta(seconds=30)
    _timeout = datetime.timedelta(seconds=30)

    def __init__(self, permissions, validator, channel, issuer_configs):
        self._permissions = permissions
        self._validator = validator
        self._channel = channel
        self._issuer_configs = issuer_configs
        self._required_roles = set()
        self._refresh_start_event = asyncio.Event()
        self._token = None
        self._token_lock = asyncio.Lock()
        self._previous_claims = None

    async def _context_stack(self, stack):
        if self._validator:
            coro = self._refresh_token()
            await stack.enter_async_context(util.background_task(coro))

    async def require_role(self, role):
        if not self._validator:
            # authentication is disabled
            return

        def check(token):
            if role not in token.roles:
                raise RuntimeError(f'Permission denied: requires role "{role}"')

        async with self._token_lock:
            if self._token is None:
                token = await self._request_token()
                check(token)
                self._token = token
                self._refresh_start_event.set()
            else:
                check(self._token)

            self._required_roles.add(role)

    async def _request_token(self):
        id_token = await self._channel.get_id_token()
        token_claims = await self._validator.extract_claims(
            id_token, leeway=self._leeway.total_seconds()
        )
        self._log_claims_once(token_claims)

        expires = datetime.datetime.fromtimestamp(
            int(token_claims["exp"]), tz=datetime.timezone.utc
        )

        return TokenInfo(
            claims=token_claims,
            roles=self._derive_roles(token_claims),
            expires=expires,
        )

    def _log_claims_once(self, token_claims):
        if logger.isEnabledFor(logging.INFO):
            show_claims = self._issuer_configs[token_claims["iss"]].show_claims
            if show_claims is not None:
                filtered_claims = [
                    (c, token_claims[c]) for c in show_claims if c in token_claims
                ]
            else:
                filtered_claims = list(token_claims.items())

            if filtered_claims and filtered_claims != self._previous_claims:
                claims_str = ", ".join([f"{k!r}: {v!r}" for k, v in filtered_claims])
                logger.info("Token claims: %s", claims_str)
                self._previous_claims = filtered_claims

    def _derive_roles(self, token_claims):
        roles = set()
        for permission in self._permissions:
            if permission.roles <= roles:
                # permission rule has no new roles, skip
                continue

            for key, required_claim in permission.claims.items():
                token_claim = token_claims.get(key)
                if isinstance(required_claim, set):
                    if (
                        not isinstance(token_claim, list)
                        or set(token_claim) < required_claim
                    ):
                        break
                elif token_claim != required_claim:
                    break
            else:
                # Token has all needed claims. Add roles.
                roles.update(permission.roles)

        return roles

    async def _refresh_token(self):
        await self._refresh_start_event.wait()

        # This is run as a background task. Every exception closes the
        # connection.
        while True:
            now = datetime.datetime.now(tz=datetime.timezone.utc)
            delay = (self._token.expires - now) + self._leeway
            logger.debug(
                "Token expires at %s, refreshing in %d seconds",
                str(self._token.expires),
                delay.total_seconds(),
            )
            await asyncio.sleep(delay.total_seconds())

            async with util.timeout(self._timeout.total_seconds()):
                self._token = await self._request_token()

            if self._required_roles > self._token.roles:
                raise RuntimeError("Permission lost")

    def get_user_name(self):
        if self._token is None:
            return None

        issuer = self._token.claims["iss"]
        templates = self._issuer_configs[issuer].user_name_formats

        for template in templates:
            with contextlib.suppress(Exception):
                return template.substitute(self._token.claims)

        # fall back to formatting the name from claims, that are guaranteed to
        # be present
        template = string.Template("${sub}@${iss}")
        return template.substitute(self._token.claims)


@dataclass
class Permission:
    claims: dict[str, dict | str | int | float | set | bool]
    roles: set

    @classmethod
    def from_dict(cls, d):
        claims = {}
        for key, value in d["claims"].items():
            if isinstance(value, list):
                claims[key] = set(value)
            else:
                claims[key] = value

        return cls(claims, set(d["roles"]))


@dataclass
class IssuerConfig:
    show_claims: list[str] | None = None
    user_name_formats: list[string.Template] = field(default_factory=list)

    @classmethod
    def from_dict(cls, d):
        show_claims = d.get("show_claims")

        user_name_formats = d.get("user_name_formats", [])
        user_name_formats = [string.Template(s) for s in user_name_formats]

        return cls(show_claims, user_name_formats)


@dataclass
class TokenInfo:
    claims: dict[str, dict | str | int | float | list | bool]
    roles: set[str]
    expires: datetime.datetime


class EMA:
    """Exponential Moving Average for float values"""

    def __init__(self, alpha=0.25):
        self._alpha = alpha
        self._value = None

    def __float__(self):
        return self._value if self._value is not None else 0.0

    def update(self, new_value):
        if self._value is None:
            self._value = new_value
        else:
            self._value = (1 - self._alpha) * self._value + self._alpha * new_value


@dataclass
class PlaceStats:
    uniqueness: EMA = field(default_factory=EMA)


def _unmap_ip(ip_str):
    """Resolve IPv4-mapped-on-IPv6 to an IPv4 address"""
    ip = ipaddress.ip_address(ip_str)
    unmapped = getattr(ip, "ipv4_mapped", None) or ip
    return str(unmapped)


class ProtocolError(Exception):
    pass
