#!/usr/bin/env python3
from os import makedirs
from time import sleep
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Any, Dict, List

import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse

from . import LOGGER
from .cache import OktaCache
from .models import GraphQuery, QueryResult
from .plugins import NodePlugin, NullNodePlugin


class OktaGraphServer:
    def __init__(
        self,
        cache: OktaCache,
        static_content_dir: str = "./static",
        temp_dir: str = "/temp",
        node_plugin: NodePlugin = NullNodePlugin(),
    ):
        #: A callable that can modify nodes before they are inserted into the graph
        self.node_plugin = node_plugin

        #: Root of the static content directories
        self.static_content_dir: str = static_content_dir
        #: The location for temporary files generated by Pydot
        self.temp_dir: str = temp_dir
        #: The cache object to use
        self.cache: OktaCache = cache

        #: The FastAPI app that our server will run
        self.app: FastAPI = FastAPI(lifespan=self.__lifespan)

        self.app.mount(
            "/static", StaticFiles(directory=self.static_content_dir), name="static"
        )

        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
        )

        @self.app.options("/graph")
        async def _() -> None:
            """Returns an empty response to allow for CORS requests"""
            return None

        @self.app.get("/")
        async def _() -> FileResponse:
            """Returns the index page."""
            return FileResponse(f"{self.static_content_dir}/index.html")

        @self.app.get("/group_names")
        async def _() -> List[str]:
            """Returns a list containing all group names from Okta"""
            return [x["profile"]["name"] for x in self.cache.groups.values()]

        @self.app.post("/graph")
        async def _(query: GraphQuery) -> QueryResult:
            res = await query.make_graph(
                temp_dir=self.temp_dir,
                icon_dir=f"{self.static_content_dir}/icons",
                groups=self.groups,
                rules=self.rules,
                node_plugin=self.node_plugin,
            )
            return res.model_dump()

    @asynccontextmanager
    async def __lifespan(self, _: FastAPI) -> AsyncGenerator[None, None]:
        """
        Manages the lifespan for the FastAPI app
        """
        makedirs(f"{self.temp_dir}/maps", exist_ok=True)
        makedirs(f"{self.temp_dir}/dots", exist_ok=True)
        makedirs(f"{self.temp_dir}/images", exist_ok=True)

        await self.cache.start()
        await self.node_plugin.start()

        yield

        await self.cache.terminate()
        await self.node_plugin.terminate()

        while self.cache.status not in (
            "Terminated",
            "Complete",
        ) or self.node_plugin.status not in ("Terminated", "Complete"):
            LOGGER.info(f"Node Plugin Status: {self.node_plugin.status}")
            LOGGER.info(f"Cache Status: {self.cache.status}")

            sleep(1)

    def run(self, **kwargs: Dict[str, Any]) -> None:
        """
        Starts the server. Accepts all of the arguments that `uvicorn.run` does and passes them directly
        into the call to start the server
        """
        if kwargs.get("reload"):
            app = "main:self.app"
            kwargs["reload_delay"] = kwargs.get("reload_delay", 1)
        else:
            app = self.app

        uvicorn.run(app, lifespan="on", **kwargs)

    @property
    def groups(self):
        return self.cache.groups

    @property
    def rules(self):
        return self.cache.rules
