#!/usr/bin/env python3
"""
port_checker.py
- expects `expected.json` in same dir (see format below)
- writes nmap XML to a temp file, parses, compares, prints a report
"""
import os
import json
import subprocess
import tempfile
from datetime import datetime
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, List, Set

from utils.scanner import nmap_scanner
from utils.models import HostResult
from reporting_jinja import write_html_report_jinja

EXPECTED_FILE = Path() / "data" / "expected.json"
HTML_REPORT_FILE = Path() / "data" / "report.html"

def load_expected(path: Path) -> Dict[str, Dict[str, Set[int]]]:
    with path.open() as fh:
        arr = json.load(fh)
    out = {}
    for entry in arr:
        ip = entry["ip"]
        out[ip] = {
            "expected_tcp": set(entry.get("expected_tcp", [])),
            "expected_udp": set(entry.get("expected_udp", [])),
        }
    return out

# def write_targets(expected: Dict[str, Dict[str, Set[int]]], path: Path) -> None:
    path.write_text("\n".join(sorted(expected.keys())) + "\n")

# 
def results_to_open_sets(
        results: List[HostResult],
        count_as_open: Set[str] = frozenset({"open", "open|filtered"})) -> Dict[str, Dict[str, Set[int]]]:
    """
    Convert HostResult list to:
    { ip: {"tcp": {open ports}, "udp": {open ports}} }
    Only include ports whose state is in `count_as_open`.
    """
    out: Dict[str, Dict[str, Set[int]]] = {}
    for hr in results:
        tcp = set()
        udp = set()
        for p in hr.ports:
            if p.state.lower() in count_as_open:
                (tcp if p.protocol == "tcp" else udp).add(p.port)
        out[hr.address] = {"tcp": tcp, "udp": udp}
    return out

# Build the "reports" dict (what the HTML renderer expects)
def build_reports(
    expected: Dict[str, Dict[str, Set[int]]],
    discovered: Dict[str, Dict[str, Set[int]]],
) -> Dict[str, Dict[str, List[int]]]:
    """
    Create the per-IP delta structure:
    {
      ip: {
        "unexpected_tcp": [...],
        "missing_tcp": [...],
        "unexpected_udp": [...],
        "missing_udp": [...]
      }
    }
    """
    reports: Dict[str, Dict[str, List[int]]] = {}
    all_ips = set(expected.keys()) | set(discovered.keys())

    for ip in sorted(all_ips):
        exp_tcp = expected.get(ip, {}).get("expected_tcp", set())
        exp_udp = expected.get(ip, {}).get("expected_udp", set())
        disc_tcp = discovered.get(ip, {}).get("tcp", set())
        disc_udp = discovered.get(ip, {}).get("udp", set())

        reports[ip] = {
            "unexpected_tcp": sorted(disc_tcp - exp_tcp),
            "missing_tcp": sorted(exp_tcp - disc_tcp),
            "unexpected_udp": sorted(disc_udp - exp_udp),
            "missing_udp": sorted(exp_udp - disc_udp),
        }
    return reports

def main():

    # repo = ScanConfigRepository()

    if not EXPECTED_FILE.exists():
        print("Expected File not found")
        return

    expected = load_expected(EXPECTED_FILE)
    targets = sorted(expected.keys())
    scanner = nmap_scanner(targets)
    scan_results = scanner.scan_targets()
    discovered_sets = results_to_open_sets(scan_results, count_as_open={"open", "open|filtered"})
    reports = build_reports(expected, discovered_sets)
    write_html_report_jinja(reports=reports,host_results=scan_results,out_path=HTML_REPORT_FILE,title="Compliance Report",only_issues=True)
    scanner.cleanup()

if __name__ == "__main__":
    main()