from random import choice, randint
from database import init_database, insert_equipment, fetch_all_equipments
from classes import Equipment, Interface

NUMBER_OF_EQUIPMENTS_TO_POPULATE = 15

COMMON_SERVICES: dict[str, int] = {
    "FTP": 21,
    "SSH": 22,
    "Telnet": 23,
    "SMTP": 25,
    "DNS": 53,
    "HTTP": 80,
    "POP3": 110,
    "NetBIOS": 139,
    "IMAP": 143,
    "HTTPS": 443,
    "LDAPS": 636,
    "RDP": 3389,
}

TYPES = [
    "Server",
    "Instance",
    "Switch",
    "Gateway",
    "Router",
    "WAF",
    "Database",
]

def random_ipv4() -> str:
    return ".".join(str(randint(0, 255)) for _ in range(4))

def load_words(path: str = "wordlist.txt") -> list[str]:
    with open(path, "r", encoding="utf-8") as f:
        return [line.strip() for line in f if line.strip()]

def random_hostname(words: list[str]) -> str:
    triplet = "-".join(choice(words) for _ in range(3))
    return f"{triplet}-{randint(1, 999):03d}"

def random_interfaces() -> list[Interface]:
    count = randint(2, 5)
    services = list(COMMON_SERVICES.items())
    return [Interface(port=port, service=svc) for svc, port in
            (choice(services) for _ in range(count))]

def main(n: int = NUMBER_OF_EQUIPMENTS_TO_POPULATE) -> None:
    init_database()
    if fetch_all_equipments():
        print("Base déjà peuplée, peuplement ignoré.")
        return
    words = load_words()
    for _ in range(n):
        insert_equipment(Equipment(
            hostname=random_hostname(words),
            type=choice(TYPES),
            ip_address=random_ipv4(),
            is_active=True,
            latency=round(randint(1, 500) / 10, 1),
            interfaces=random_interfaces(),
        ))
    print(f"{n} équipements insérés.")

if __name__ == "__main__":
    main()