import asyncio
import logging
from abc import ABC, abstractmethod
from functools import partial
from typing import Awaitable, Callable, Dict, List, Optional, Type, Union

from aiohttp import web
from aiohttp.web_log import AccessLogger
from smb.common.http_client import HttpClient
from smb.common.pgswim import Migrator, SwimEngine
from smb.common.sensors import BaseBuilder, HistRateBuilder, MetricsHub, RateBuilder

from smb.common.aiotvm import TvmClient
from maps_adv.warden.client.lib import PeriodicalTask, TaskContext, TaskMaster

from . import sensors, tvm
from .exceptions import DbAutomigrationError

__all__ = ["Lasagna", "LasagnaAccessLogger"]


class LasagnaAccessLogger(AccessLogger):
    ignore_handles = ("/ping", "/sensors/")

    def log(self, request: web.Request, response: web.Response, time: float):
        if request.path not in self.ignore_handles:
            super().log(request, response, time)

        self.sensor_log(request, response, time)

    def sensor_log(self, request: web.Request, response: web.Response, time: float):
        lasagna: Lasagna = request.app["lasagna"]

        if lasagna.SENSORS is None:
            return

        info = request.match_info.get_info()
        path = info.get("formatter", info.get("path"))
        if not path:
            return

        if "rps" in lasagna.SENSORS:
            sensors.log_request("rps", request, response.status, 1)

        if "response_time" in lasagna.SENSORS:
            sensors.log_request("response_time", request, response.status, time * 1000)


class Lasagna(ABC):
    __slots__ = "config", "db", "tvm", "sensors", "_clients"

    SWIM_ENGINE_CLS: Type[SwimEngine] = SwimEngine
    MIGRATIONS_PATH: Optional[str] = "lib/db/migrations"

    TASKS: Optional[Dict[str, Callable[[TaskContext, dict], Awaitable[None]]]] = None
    TASKS_KWARGS_KEYS: List[str] = []

    SENSORS: Optional[Dict[str, BaseBuilder]] = {
        "rps": RateBuilder(),
        "response_time": HistRateBuilder().linear(size=200),
        "warden_tasks_status": RateBuilder(),
        "logging": RateBuilder(),
        "exceptions": RateBuilder(),
    }

    config: dict
    db: Optional[SwimEngine]
    tvm: Optional[TvmClient]
    sensors: Optional[MetricsHub]

    _clients: List[HttpClient]

    def __init__(self, config: dict):
        self.config = config

        self._clients = []

    async def setup(self, db: Optional[SwimEngine]):
        self.db = db

        self.tvm = await self._setup_tvm()
        self.sensors = self._setup_sensors()

        api = await self._setup_layers(self.db)

        if self.tvm:
            api.on_shutdown.append(lambda *a: self.tvm.close())
            self._setup_tvm_middleware(api)

        self._setup_sensors_integration(api)
        self._setup_tasks(api)

        if self._clients:
            api.on_shutdown.append(
                lambda _: asyncio.gather(*[c.close() for c in self._clients])
            )

        api["lasagna"] = self
        return api

    def register_client(self, client: HttpClient) -> HttpClient:
        self.sensors.merge_with(client.sensors, prefix=client.__class__.__name__)
        self._clients.append(client)
        return client

    async def _setup_tvm(self) -> Optional[TvmClient]:
        if not all([self.config.get("TVM_DAEMON_URL"), self.config.get("TVM_TOKEN")]):
            return

        return await TvmClient(self.config["TVM_DAEMON_URL"], self.config["TVM_TOKEN"])

    @abstractmethod
    async def _setup_layers(self, db: Union[SwimEngine, None]) -> web.Application:
        raise NotImplementedError()

    def _setup_sensors(self) -> Optional[MetricsHub]:
        if self.SENSORS is not None:
            return MetricsHub(**self.SENSORS)

    def _setup_sensors_integration(self, api: web.Application):
        if self.SENSORS is None:
            return

        if "rps" in self.SENSORS:
            api.middlewares.insert(0, sensors.rps_middleware)
        if "response_time" in self.SENSORS:
            api.middlewares.insert(0, sensors.response_time_middleware)

        api.add_routes([web.get("/sensors/", sensors.handler)])

        if "logging" in self.SENSORS:
            _group = self.sensors.take_group("logging")
            logging_handler = sensors.SolomonLoggingHandler(metric_group=_group)

            root_logger = logging.getLogger()
            root_logger.addHandler(logging_handler)

    def _setup_tvm_middleware(self, api: web.Application) -> web.Application:
        _middleware = tvm.Middleware(self.tvm, self.config)
        api.middlewares.insert(0, _middleware)

        return api

    def _setup_tasks(self, api):
        if (
            not self.config.get("WARDEN_URL")
            or not self.TASKS
            or not self.config.get("WARDEN_TASKS")
        ):
            return

        kwargs = {k: getattr(self, k) for k in self.TASKS_KWARGS_KEYS}
        tasks = {
            name: func
            for name, func in self.TASKS.items()
            if name in self.config["WARDEN_TASKS"]
        }

        task_master = TaskMaster(
            server_url=self.config["WARDEN_URL"],
            tasks=[
                PeriodicalTask(
                    name,
                    partial(task, **kwargs),
                    sensors=self.sensors.take_group("warden_tasks_status"),
                )
                for name, task in tasks.items()
            ],
        )

        api.on_startup.append(lambda *a: task_master.run())
        api.on_shutdown.append(lambda *a: task_master.stop())

    async def _setup_db(self) -> Union[SwimEngine, None]:
        if not self.SWIM_ENGINE_CLS:
            return

        db = await self.SWIM_ENGINE_CLS.create(self.config["DATABASE_URL"])

        if self.config.get("DB_AUTOMIGRATE") is not False:
            if not self.MIGRATIONS_PATH:
                raise DbAutomigrationError("Migrations path is not configured.")

            migrator = Migrator(db=db, migrations_path=self.MIGRATIONS_PATH)
            await migrator.upgrade()

        return db

    def run(
        self, *, host: Optional[str] = None, port: Union[int, str, None] = None
    ) -> None:
        async def _make():
            db = await self._setup_db()
            api = await self.setup(db)
            if db is not None:
                api.on_shutdown.append(lambda *a: db.close())
            return api

        return web.run_app(
            _make(),
            access_log_class=LasagnaAccessLogger,
            host=host,
            port=port,
            loop=asyncio.get_event_loop(),
        )
