from contextlib import contextmanager, nullcontext
from dataclasses import dataclass, field
import logging
from importlib import import_module
from io import UnsupportedOperation
from mock import patch
import os
from random import randrange
import requests
from shutil import rmtree
import signal
import sys
from threading import Thread
import time
import uuid

import fastapi as fa
from prometheus_client import multiprocess, CollectorRegistry
from prometheus_fastapi_instrumentator import Instrumentator
from pydantic import BaseModel, Extra, Field
import uvicorn

from chaiverse.logging import get_service_logger_config

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


class ChaiverseServerGracefulShutdown(Exception):
    pass


class UvicornServer(uvicorn.Server):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Python does not raise an exception during SIGTERM
        # so we must do it ourselves to gracefully shutdown
        # the uvicorn thread
        signal.signal(signal.SIGTERM, self.graceful_shutdown)

    def install_signal_handlers(self):
        pass

    @contextmanager
    def run_in_thread(self):
        thread = Thread(target=self.run)
        thread.start()
        try:
            yield
        finally:
            self.should_exit = True
            thread.join()

    def graceful_shutdown(self, signum, frame):
        raise ChaiverseServerGracefulShutdown


class ChaiverseServerConfig(BaseModel):
    name: str = "chaiverse_server"
    port: int = 8080
    host: str = "0.0.0.0"
    workers: int = 1
    environment_variables: dict = Field(default_factory=dict)
    # Set to a positive number to limit the number of requests the uvicorn
    # process can handle. Used to avoid memory leaks. Random jitter is added
    # to ensure replicas don't shutdown at same time.
    limit_max_requests: int = -1
    log_level: str = "warning"
    loki_base_url: str = ""

    class Config:
        extra = Extra.allow


class ChaiverseServer:
    def __init__(self, fastapi_app: fa.FastAPI, config):
        self.config = config
        self.fastapi_app = self._prepare_fastapi(fastapi_app)

    @contextmanager
    def serve(self):
        try:
            logger.info(f"Starting server {self.config.name}.")
            with self._start():
                logger.info(f"Waiting for server {self.config.name} to be ready.")
                start = time.time()
                while not self._server_ready:
                    pass
                duration = round(time.time() - start, 2)
                logger.info(f"Server {self.config.name} ready after {duration}s")
                yield
        finally:
            logger.info(f"Shutting down server {self.config.name}.")

    @contextmanager
    def _start(self, **kwargs):
        config = {
            "app": self.fastapi_app,
            "port": self.config.port,
            "host": self.config.host,
            "workers": self.config.workers,
            "log_config": self.logging_config,
            "log_level": self.config.log_level,
        }
        if self.config.limit_max_requests > 0:
            config.update({"limit_max_requests": self.limit_max_requests})
        config = uvicorn.Config(**config)
        server = UvicornServer(config=config)
        with self._environment_patch(), server.run_in_thread():
            yield

    @property
    def logging_config(self):
        return get_service_logger_config(self.config.name, self.config.loki_base_url)

    @property
    def _server_ready(self):
        try:
            requests.get(f"http://localhost:{self.config.port}/test_endpoint")
            is_ready = True
        except requests.exceptions.ConnectionError:
            logger.info(f"{self.config.name} not yet ready!")
            is_ready = False
        return is_ready

    def _prepare_fastapi(self, fastapi_app):
        # Allow fastapi app to be passed in as import string
        # to enable specifying app via environment variable
        # in production
        if isinstance(fastapi_app, str):
            fastapi_app = _resolve_fastapi_app(fastapi_app)
        # We copy the app as we are modifying it, and may want to launch
        # multiple instances of it
        _install_general_exception_handler(fastapi_app)
        _install_prometheus_instrumentator(fastapi_app, self.config.name)
        fastapi_app.state.chaiverse_server_config = self.config
        return fastapi_app

    def _environment_patch(self):
        environment_patch = (
            patch.dict(os.environ, self.config.environment_variables)
            if self.config.environment_variables
            else nullcontext()
        )
        return environment_patch

    @property
    def limit_max_requests(self):
        base_limit = self.config.limit_max_requests
        adjusted_limit = randrange(base_limit, int(base_limit + 0.25 * base_limit))
        return adjusted_limit


def _install_general_exception_handler(fastapi_app):
    @fastapi_app.exception_handler(Exception)
    def runtime_exception_handler(request: fa.Request, exc: Exception):
        return exception_response(500, exc)


def _install_prometheus_instrumentator(fastapi_app, name):
    if _should_install_prometheus_middleware(fastapi_app):
        prometheus_dir = f"/tmp/prometheus/{uuid.uuid4()}"
        os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_dir
        if not os.path.exists(prometheus_dir):
            os.mkdir(prometheus_dir)

        registry = CollectorRegistry()
        multiprocess.MultiProcessCollector(registry)

        instrumentator = Instrumentator()
        instrumentator.instrument(fastapi_app, metric_subsystem=name)
        instrumentator.expose(fastapi_app)


def _should_install_prometheus_middleware(fastapi_app):
    # This avoids a bug on later version of fastapi where it crashes
    # if you try to add a middleware more than once (useful in testing)
    should_add_middleware = True
    if fastapi_app.middleware_stack is not None:
        should_add_middleware = False
    return should_add_middleware


def exception_response(status_code, exc):
    if (
        getattr(exc, "message", None)
        and "unhandled errors in a TaskGroup" in exc.message
    ):
        exc = exc.exceptions
    return fa.responses.JSONResponse(
        status_code=status_code, content={"detail": str(exc)}
    )


def _resolve_fastapi_app(fastapi_app):
    import_path, app = fastapi_app.rsplit(".", 1)
    module = import_module(import_path)
    fastapi_app = getattr(module, app)
    return fastapi_app
