import os
import sys
import time
import importlib.util

from pathlib import Path
from threading import Thread
from multiprocessing import JoinableQueue

from milkman.db import save_flags
from milkman.logger import logger
from milkman.config import Config
from milkman.singleton import SingletonMeta

# This path is relative to CWD, not to the file. I expect this to only be called from app.py
p = Path("../exploits")
conf = Config()


class Exploits(list, metaclass=SingletonMeta):
    def __init__(self, *args, **kwargs):
        self.load_exploits()

    def load_exploits(self, *args, **kwargs):
        log = logger.bind(file="exploits.log")
        exploitsfiles = list(p.glob("**/[!_]*.py"))
        exploits = []
        for exp in exploitsfiles:
            name = exp.stem

            try:
                spec = importlib.util.spec_from_file_location(name, exp.resolve())
                module = importlib.util.module_from_spec(spec)  # type: ignore
                sys.modules[name] = module
                spec.loader.exec_module(module)  # type: ignore
                exploitfun = getattr(sys.modules[name], "main")
                exploits.append(exploitfun)
            except Exception as e:
                log.error(f"Error while importing {name}: {e}")

        log.debug(f"{exploits = }")

        # This may rarely cause bugs if some other process uses the list before we extend it with the new ones.
        # log.critical(f"exp = {[x.__module__ for x in self]}")
        gametickQueue = JoinableQueue()
        for exploit in exploits:
            if exploit.__module__ not in [x.__module__ for x in self]:
                log.info(f"Found new exploit: {exploit.__module__}")
                addlogger(name)
                threads = runExploit(exploit, gametickQueue)
                Thread(
                    target=timeCounterThread,
                    args=(gametickQueue, threads),
                    name=f"{exploit.__module__}",
                ).start()
        self.clear()
        self.extend(exploits)


def timeCounterThread(gametickQueue, threads):
    log = logger.bind(file="exploits.log")
    duration = conf["max_thread_time"]
    for t in threads:
        gametickQueue.put("")
    start = time.time()
    while gametickQueue._notempty:
        pass
    gametickQueue.join()
    for t in threads:
        t.join()
    gametickQueue.close()
    elapsed = time.time() - start
    remaining = duration - elapsed
    if remaining < 0:
        remaining = 0

    log.info(f"Took {elapsed:.2f} seconds. Waiting for {remaining:.2f} seconds")

    return


def addlogger(name: str):
    logger.add(
        sink=f"../logs/exploits/{name}.log",
        format="<d>{time:HH:mm:ss.SS}</d> | <level>{level:^8}</level> | {message}",
        enqueue=True,
        filter=lambda record: record["extra"]["file"] == f"{name}.log",
        level="DEBUG",
        colorize=True,
    )


def runExploit(exploit, gametickQueue):
    log = logger.bind(file="exploits.log")

    threads = []
    for id in range(1, conf["highest_id"] + 1):
        if id == conf["team_id"]:
            continue
        target_ip = conf["baseip"].format(id=id)

        t = Thread(
            target=launchAttack,
            args=(
                exploit,
                target_ip,
                gametickQueue,
            ),
            name=f"{exploit.__module__}_{target_ip}",
        )
        t.start()
        threads.append(t)
    return threads


def launchAttack(exploit, target_ip: str, gametickQueue):
    log = logger.bind(file=f"exploits.log")
    gametickQueue.get()
    if (
        conf["ping_before_exploit"]
        and os.system("ping -c 1 " + target_ip + " > /dev/null") != 0
    ):
        log.info(f"{target_ip} is unreachable")
        gametickQueue.task_done()
        return

    level = "INFO"
    start = time.time()
    flags_iterable = exploit(target_ip)
    if isinstance(flags_iterable, str):
        flags_iterable: list[str] = [flags_iterable]
    for flag in flags_iterable:
        save_flags(flag, exploit.__module__)  # type: ignore

    elapsed = time.time() - start
    gametickQueue.task_done()

    if elapsed > conf["max_thread_time"]:
        level = "WARNING"
    log.log(level, f"{exploit.__module__} took {elapsed:.2f} seconds")
    return


if __name__ == "__main__":
    # print(f"{exploitsfiles = }")
    exploits = Exploits()
    for ex in exploits:
        logger.info(f"{ex.__module__}")
    logger.info(f"{exploits = }")
