from fastapi import FastAPI
from fastapi import Request
from fastapi.responses import Response
from fastapi import BackgroundTasks
import uvicorn
import os
from time import time
import asyncio

class Collector:
    app = FastAPI()
    secrets = []
    hosts = []
    store = {}
    def __init__(self) -> None:
        self.secrets = os.getenv("KEYS", "secret").split(",")
        self.hosts = os.getenv("HOSTS", "srv").split(",")
        self.dead_after = int(os.getenv("DEAD_AFTER", "60"))
        self.build_routes()

    def check_headers(self, request: Request):
        if not request.headers.get("Secret", ""):
            print("secret empty")
            return Response(status_code=400)
        if not request.headers.get("Host", ""):
            print("host empty")
            return Response(status_code=400)

        if not request.headers.get("Secret", "") in self.secrets:
            return Response(status_code=400)

        if not request.headers.get("Host", "").split(":")[0] in self.hosts:
            print(request.headers.get("Host", ""))
            print(request.headers)
            return Response(status_code=400)

        return None

    async def deadline(self, host):
        await asyncio.sleep(self.dead_after)
        if host in self.store:
            if time()+1 - self.store[host]['timestamp'] > self.dead_after:
                print(f"{host} is dead")
                del self.store[host]


    def build_routes(self):
        @self.app.get("/metrics")
        async def get(request: Request):
            response: Response = self.check_headers(request)
            if response:
                return response

            if request.headers.get("Host", "").split(":")[0] in self.store and self.store[request.headers.get("Host", "").split(":")[0]]['data']:
                return Response(content=self.store[request.headers.get("Host", "").split(":")[0]]['data'], status_code=200)
            else:
                return Response(status_code=404)

        @self.app.post("/metrics")
        async def post(request: Request, background_tasks: BackgroundTasks):
            response: Response = self.check_headers(request)
            if response:
                return response

            host = request.headers.get("Host", "").split(":")[0]

            self.store[host] = {
                "data":await request.body(), 
                "timestamp": time()
                }

            background_tasks.add_task(self.deadline, host)
            return Response(status_code=200)

if __name__ == "__main__":
    collector = Collector()
    uvicorn.run(collector.app, host = os.getenv("NEC_HOST", "0.0.0.0"), port = int(os.getenv("NEC_PORT", "9100")))