# utils/reporting_jinja.py
from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from html import escape
from ipaddress import ip_address
from pathlib import Path
from typing import Dict, List

from jinja2 import Environment, FileSystemLoader, select_autoescape

from utils.models import HostResult, HostReport, GroupedReports


# ---------------------------
# Helper render-time functions
# ---------------------------

def fmt_ports(ports: List[int]) -> str:
    """
    Render a simple, comma-separated list of port numbers, deduped and sorted.
    """
    if not ports:
        return "none"
    unique_ports: List[int] = []
    seen: set[int] = set()
    for x in ports:
        p = int(x)
        if p not in seen:
            seen.add(p)
            unique_ports.append(p)

    unique_ports.sort()
    return ", ".join(str(p) for p in unique_ports)


def badge(text: str, kind: str = "") -> str:
    """
    Return a class-based badge element (no inline color).
    'kind' should map to CSS variants defined in your stylesheet (e.g., 'ok', 'issue').
    Unknown kinds fall back to the base 'badge' styling.

    Examples:
      badge("OK", "ok")        -> <span class="badge ok">OK</span>
      badge("ISSUE", "issue")  -> <span class="badge issue">ISSUE</span>
      badge("info")            -> <span class="badge">info</span>
    """
    text_s = escape(str(text))
    kind_s = str(kind).strip()
    if kind_s:
        return f'<span class="badge {kind_s}">{text_s}</span>'
    return f'<span class="badge">{text_s}</span>'


def pill_state(state: str) -> str:
    """
    Map port states to neutral/semantic classes (no inline color).
    Keeps semantics simple so CSS controls appearance.
    """
    s = (state or "").strip().lower()

    if s == "open":
        return badge("open", "ok")  # green-ish in CSS
    if s == "closed":
        return badge("closed", "issue")  # red-ish in CSS

    # For less common states, keep a neutral chip/badge so style stays consistent
    if s in ("open|filtered", "filtered"):
        return f'<span class="chip">{escape(s)}</span>'

    # Fallback
    return f'<span class="chip">{escape(s or "-")}</span>'


def pill_proto(proto: str) -> str:
    """
    Render protocol as a neutral 'chip'.
    """
    return f'<span class="chip">{escape((proto or "").lower())}</span>'


def _env(templates_dir: Path) -> Environment:
    """
    Build a Jinja2 environment with our helpers registered as globals.
    """
    env = Environment(
        loader=FileSystemLoader(str(templates_dir)),
        autoescape=select_autoescape(["html", "xml", "j2"]),
        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


# ---------------------------
# Rendering entry points
# ---------------------------

def render_html_report_jinja(
    grouped: GroupedReports,
    host_results: List[HostResult],
    templates_dir: Path,
    *,
    template_name: str = "report_body.html.j2",
    title: str = "Port Compliance Report",
    only_issues: bool = False,
    dark_mode: bool = True,
) -> str:
    """
    Render the HTML report using the new structure/palette templates.

    Args:
        grouped: GroupedReports with .issues, .expected, .by_ip already populated.
        host_results: Raw scan results (per-host) used for the detailed port tables.
        templates_dir: Path to the Jinja templates root (contains base_report.html.j2, palettes/, etc.).
        template_name: Structure template to render (defaults to 'report_body.html.j2').
        title: Report title.
        only_issues: If True, suppresses the 'expected/OK' section.
        dark_mode: Controls which palette is included by base_report.html.j2.

    Returns:
        Rendered HTML as a string.
    """
    env = _env(templates_dir)
    template = env.get_template(template_name)

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

    # Build mapping of IP -> HostResult and sort port rows for stable output
    by_ip_results: Dict[str, HostResult] = {}
    for hr in host_results:
        by_ip_results[hr.address] = hr

    for hr in by_ip_results.values():
        if hr and hr.ports:
            # Sort by (protocol, then port, then service)
            hr.ports.sort(key=lambda p: (str(p.protocol).lower(), int(p.port), str(p.service or "")))

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

        # Palette switch (consumed by base_report.html.j2)
        dark_mode=dark_mode,

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

        # Sections
        issues=grouped.issues,       # list[HostReport]
        expected=grouped.expected,   # list[HostReport]
        by_ip=grouped.by_ip,         # dict[str, HostReport]

        # Host scan details (used for port tables)
        host_results_by_ip=by_ip_results,

        # Behavior flag
        only_issues=only_issues,
    )
    return html


def write_html_report_jinja(
    grouped: GroupedReports,
    host_results: List[HostResult],
    out_path: Path,
    *,
    templates_dir: Path = Path("templates"),
    template_name: str = "report_body.html.j2",
    title: str = "Port Compliance Report",
    only_issues: bool = False,
    dark_mode: bool = True,
) -> Path:
    """
    Render and write the HTML report to disk.

    Notes:
      - Assumes your templates directory contains:
          base_report.html.j2
          report_body.html.j2
          palettes/light.css.j2
          palettes/dark.css.j2
      - 'dark_mode' toggles which palette is inlined by the base template.
    """
    html = render_html_report_jinja(
        grouped=grouped,
        host_results=host_results,
        templates_dir=templates_dir,
        template_name=template_name,
        title=title,
        only_issues=only_issues,
        dark_mode=dark_mode,
    )
    out_path.write_text(html, encoding="utf-8")
    return out_path