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
import datetime

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", "90"))
        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", "") and not request.headers.get("Host-replace", ""):
            print("host empty")
            return Response(status_code=400)

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

        host = request.headers.get("Host-replace", "") or request.headers.get("Host", "")
        host = host.split(":")[0]
        if not host in self.hosts:
            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

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

            if host in self.store and self.store[host]['data']:
                return Response(
                    content=self.store[host]['data'], 
                    status_code=200,
                    headers=self.store[host]['headers'])
            else:
                print(f"cannot find {host}")
                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-replace", "") or request.headers.get("Host", "")
            host = host.split(":")[0]

            self.store[host] = {
                "data":await request.body(), 
                "timestamp": time(),
                "headers": {
                    "Content-Type": "text/plain; version=0.0.4; charset=utf-8; escaping=underscores",
                    "Date": datetime.datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S GMT")
                }
            }

            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")))