#fs imports
from fastapi import FastAPI, HTTPException, WebSocket
from fastapi.responses import HTMLResponse
from uvicorn import run as UR
from contextlib import asynccontextmanager

#sys
import asyncio
from typing import List, Dict
import os, sys, traceback

#internal
from api_models import API_CORE

class API(API_CORE):
    exts = {}

    def __init__(self):
        self.WORKDIR = os.path.dirname(os.path.abspath(__file__))
        print(f"Change workdir to {self.WORKDIR}")
        os.chdir(self.WORKDIR)

        self.app = FastAPI(lifespan=self.lifespan)
        self.extensionLoader(["test", "system"])

    def run(self):#todo filter
        UR(self.app, host="0.0.0.0", port=14088)

    def ws(self):
        async def wshello(websocket: WebSocket):
            for ext_path, ext in self.exts.items():
                await ext.wshello(websocket)

        @self.app.webhooks("/api/ws")
        async def ws(websocket: WebSocket):
            await websocket.accept()
            self.ws_clients.append(websocket)
            await wshello(websocket)
            try:
                while True:
                    await websocket.receive_json()
            except:
                pass
            finally:
                self.ws_clients.remove(websocket)
                try:
                    await websocket.close()
                except:
                    pass

    def extensionLoader(self, search_paths = []):
        if type(search_paths) == str:
            search_paths = [search_paths]

        for path in search_paths:
            print(f"Try found extensions in {path}")
            if not os.path.exists(path) or not os.path.isdir(path):
                print(f"Directory is not exists or not directory, skip")
                continue

            sys.path.insert(0, path)
            for extension in os.listdir(path):
                extension, ext = os.path.splitext(extension)
                if ext != ".py":
                    continue
                print(f"Found ext: {extension}")
                self.exts[f"{path}/{extension}"] = __import__(extension).API_EXTENSION(self)
            sys.path.pop(0)

        print(f"Found {self.exts.keys().__len__()} extension")
        for ext_path, ext in self.exts.items():
            print(f"Register routes")
            ext.routes()

    @asynccontextmanager
    async def lifespan(self, app: FastAPI):
        print("server is started")
        for ext_path, ext in self.exts.items():
            print(f"start backgrounds tasks in {ext_path}")
            for task in ext.tasks:
                asyncio.create_task(task())
        yield
        print("server is shutdown now")

if __name__ == "__main__":
    API().run()