# utils/reporting_jinja.py
from __future__ import annotations
from datetime import datetime
from html import escape
from pathlib import Path
from typing import Any, Dict, List, Set, Union

from jinja2 import Environment, FileSystemLoader, select_autoescape

from utils.models import HostResult, HostReport, GroupedReports  # <-- add HostReport, GroupedReports
from ipaddress import ip_address

def fmt_ports(ports: List[int]) -> str:
    if not ports:
        return "none"
    return ", ".join(str(p) for p in sorted(set(int(x) for x in ports)))

def badge(text: str, bg: str, fg: str = "#ffffff") -> str:
    return (
        f'<span style="display:inline-block;padding:2px 6px;border-radius:12px;'
        f'font-size:12px;line-height:16px;background:{bg};color:{fg};'
        f'font-family:Segoe UI,Arial,sans-serif">{escape(text)}</span>'
    )

def pill_state(state: str) -> str:
    s = (state or "").lower()
    if s == "open":
        return badge("open", "#16a34a")
    if s in ("open|filtered",):
        return badge("open|filtered", "#0ea5e9")
    if s == "filtered":
        return badge("filtered", "#f59e0b", "#111111")
    if s == "closed":
        return badge("closed", "#ef4444")
    return badge(s, "#6b7280")

def pill_proto(proto: str) -> str:
    return badge((proto or "").lower(), "#334155")

def _env(templates_dir: Path) -> Environment:
    env = Environment(
        loader=FileSystemLoader(str(templates_dir)),
        autoescape=select_autoescape(["html", "xml"]),
        trim_blocks=True,
        lstrip_blocks=True,
    )
    env.globals.update(badge=badge, pill_state=pill_state, pill_proto=pill_proto, fmt_ports=fmt_ports)
    return env


# --- NEW: compatibility shim so renderer can accept either shape ---

ReportsArg = Union[GroupedReports, Dict[str, Dict[str, List[int]]]]

def _coerce_to_grouped(reports: ReportsArg) -> GroupedReports:
    """
    Accept either:
      - GroupedReports (new, preferred), or
      - legacy dict: { ip: {"unexpected_tcp":[], "missing_tcp":[], "unexpected_udp":[], "missing_udp":[] } }
    and return a GroupedReports instance.
    """
    if isinstance(reports, GroupedReports):
        return reports

    # Legacy dict -> build HostReport objects
    by_ip: Dict[str, HostReport] = {}
    for ip, d in reports.items():
        hr = HostReport(
            ip=str(ip),
            unexpected_tcp=list(d.get("unexpected_tcp", []) or []),
            missing_tcp=list(d.get("missing_tcp", []) or []),
            unexpected_udp=list(d.get("unexpected_udp", []) or []),
            missing_udp=list(d.get("missing_udp", []) or []),
        )
        by_ip[ip] = hr

    # Split & sort by IP
    issues: List[HostReport] = []
    expected: List[HostReport] = []
    for hr in by_ip.values():
        (issues if hr.has_issues() else expected).append(hr)

    def _ip_key(hr: HostReport):
        try:
            return ip_address(hr.ip)
        except ValueError:
            return hr.ip

    issues.sort(key=_ip_key)
    expected.sort(key=_ip_key)

    return GroupedReports(issues=issues, expected=expected, by_ip=by_ip)


def render_html_report_jinja(
    reports: ReportsArg,
    host_results: List[HostResult],
    templates_dir: Path,
    template_name: str = "report.html.j2",
    title: str = "Port Compliance Report",
    only_issues: bool = False,
) -> str:
    env = _env(templates_dir)
    template = env.get_template(template_name)

    grouped = _coerce_to_grouped(reports)

    total_hosts = len(grouped.by_ip)
    ok_hosts = len(grouped.expected)
    hosts_with_issues = [hr.ip for hr in grouped.issues]

    # Build a mapping of IP -> HostResult and sort port rows for stable output
    by_ip_results = {hr.address: hr for hr in host_results}
    for hr in by_ip_results.values():
        if hr and hr.ports:
            hr.ports.sort(key=lambda p: (p.protocol, p.port))

    html = template.render(
        title=title,
        generated=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),

        # Summary bar
        total_hosts=total_hosts,
        ok_hosts=ok_hosts,
        hosts_with_issues=hosts_with_issues,

        # New grouped context for simple “issues first, expected later” loops
        issues=grouped.issues,        # list[HostReport]
        expected=grouped.expected,    # list[HostReport]
        by_ip=grouped.by_ip,          # dict[str, HostReport]

        # Legacy context (kept for compatibility if your template still uses it)
        reports=reports,

        # Host scan details
        host_results_by_ip=by_ip_results,

        # Existing behavior switch
        only_issues=only_issues,
    )
    return html


def write_html_report_jinja(
    reports: ReportsArg,
    host_results: List[HostResult],
    out_path: Path,
    templates_dir: Path = Path("templates"),
    template_name: str = "report.html.j2",
    title: str = "Port Compliance Report",
    only_issues: bool = False,
) -> Path:
    html = render_html_report_jinja(
        reports, host_results, templates_dir, template_name, title, only_issues
    )
    out_path.write_text(html, encoding="utf-8")
    return out_path