#!/usr/bin/env python3
"""
souleyez.ui.live_formatters - Transform raw tool output into user-friendly live summaries

Instead of showing raw log output, these formatters parse tool output in real-time
and present it as clean, categorized summaries with icons and progress indicators.
"""

import re
from typing import Callable, Optional

import click


# Registry of tool formatters
FORMATTERS: dict[str, Callable] = {}

# Severity color mapping
SEVERITY_COLORS = {
    "critical": ("red", True),  # (color, bold)
    "high": ("red", False),
    "medium": ("yellow", False),
    "low": ("bright_black", False),
    "info": ("cyan", False),
}


def severity_style(text: str, severity: str) -> str:
    """Apply color styling based on severity level."""
    severity = severity.lower()
    if severity in SEVERITY_COLORS:
        color, bold = SEVERITY_COLORS[severity]
        return click.style(text, fg=color, bold=bold)
    return text


def render_progress_bar(current: int, total: int, width: int = 20) -> str:
    """
    Render a text-based progress bar.

    Example: [████████░░░░░░░░░░░░] 42%
    """
    if total <= 0:
        return ""

    percent = min(100, int((current / total) * 100))
    filled = int((current / total) * width)
    empty = width - filled

    bar = "█" * filled + "░" * empty
    return f"[{bar}] {percent}%"


def strip_ansi(text: str) -> str:
    """Remove ANSI escape codes from text."""
    ansi_pattern = re.compile(r"\x1b\[[0-9;]*m")
    return ansi_pattern.sub("", text)


def register_formatter(tool_name: str):
    """Decorator to register a formatter for a specific tool."""

    def decorator(func: Callable):
        FORMATTERS[tool_name] = func
        return func

    return decorator


def format_live_output(lines: list, tool: str, target: str = "") -> list:
    """
    Format raw log lines into user-friendly output.

    Args:
        lines: Raw log lines from the tool
        tool: Tool name (e.g., 'nikto', 'nmap', 'msf_auxiliary')
        target: Target being scanned (for context)

    Returns:
        List of formatted, user-friendly lines
    """
    # Get the appropriate formatter
    formatter = FORMATTERS.get(tool)

    if formatter:
        return formatter(lines, target)
    else:
        # Fallback: generic formatter that cleans up common patterns
        return format_generic(lines, target)


def format_generic(lines: list, target: str = "") -> list:
    """
    Generic formatter for tools without specific handlers.
    Filters metadata and applies basic cleanup.
    """
    result = []

    # Metadata prefixes to skip
    metadata_prefixes = (
        "=== Plugin:",
        "=== Command Execution",
        "=== End Command",
        "Target:",
        "Args:",
        "Label:",
        "Started:",
        "Command:",
        "Timeout:",
        "Exit Code:",
        "Duration:",
        "=== Output",
        "=== Wrapper",
    )

    for line in lines:
        stripped = line.strip()

        # Skip empty lines
        if not stripped:
            continue

        # Skip metadata
        if stripped.startswith(metadata_prefixes):
            continue

        # Pass through other lines
        result.append(line)

    return result


# =============================================================================
# NIKTO FORMATTER
# =============================================================================
@register_formatter("nikto")
def format_nikto(lines: list, target: str = "") -> list:
    """Format nikto output into friendly summary."""
    result = []
    # Categorize findings by severity
    findings_by_severity = {
        "critical": [],
        "high": [],
        "medium": [],
        "low": [],
        "info": [],
    }
    server_info = None
    scanning_started = False

    for line in lines:
        stripped = line.strip()

        # Skip metadata and empty lines
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue
        if stripped.startswith(("Started:", "Command:", "Timeout:")):
            continue

        # Skip nikto banner/headers
        if stripped.startswith("- Nikto v") or stripped.startswith("---"):
            scanning_started = True
            continue

        # Skip target info lines (we already know the target)
        if stripped.startswith("+ Target IP:") or stripped.startswith(
            "+ Target Hostname:"
        ):
            continue
        if stripped.startswith("+ Target Port:") or stripped.startswith(
            "+ Start Time:"
        ):
            continue

        # Parse nikto findings (lines starting with +)
        if stripped.startswith("+"):
            content = stripped[1:].strip()

            # Server identification
            if content.startswith("Server:"):
                server_info = content.replace("Server:", "").strip()
                continue

            # Categorize findings by severity
            if "CVE" in content:
                # CVEs are typically high/critical
                finding = content[:70] + "..." if len(content) > 70 else content
                findings_by_severity["high"].append(finding)
            elif "OSVDB" in content:
                finding = content[:70] + "..." if len(content) > 70 else content
                findings_by_severity["medium"].append(finding)
            elif any(
                h in content
                for h in [
                    "X-Frame-Options",
                    "X-Content-Type-Options",
                    "X-XSS-Protection",
                    "Content-Security-Policy",
                ]
            ):
                # Missing security headers = medium
                if "X-Frame-Options" in content:
                    findings_by_severity["medium"].append(
                        "Missing X-Frame-Options header"
                    )
                elif "X-Content-Type-Options" in content:
                    findings_by_severity["medium"].append(
                        "Missing X-Content-Type-Options header"
                    )
                elif "X-XSS-Protection" in content:
                    findings_by_severity["medium"].append(
                        "Missing X-XSS-Protection header"
                    )
                elif "Content-Security-Policy" in content:
                    findings_by_severity["medium"].append(
                        "Missing Content-Security-Policy header"
                    )
            elif "Strict-Transport-Security" in content:
                findings_by_severity["medium"].append("Missing HSTS header")
            elif "directory" in content.lower() or "folder" in content.lower():
                match = re.search(r"(/[^\s]+)", content)
                if match:
                    findings_by_severity["info"].append(f"Directory: {match.group(1)}")
            elif "No CGI" in content:
                findings_by_severity["info"].append("No CGI directories found")
            elif "allowed HTTP" in content.lower() or "methods" in content.lower():
                findings_by_severity["low"].append(content[:60])
            else:
                # Other findings = low
                if len(content) > 70:
                    content = content[:67] + "..."
                findings_by_severity["low"].append(content)

    # Build formatted output
    if target:
        result.append(click.style(f"  Scanning {target}", fg="cyan"))
        result.append("")

    if server_info:
        result.append(click.style(f"  Server: {server_info}", fg="bright_black"))

    total_findings = sum(len(f) for f in findings_by_severity.values())
    if total_findings > 0:
        result.append("")

        # Show findings by severity (most severe first)
        for sev in ["critical", "high", "medium", "low"]:
            items = findings_by_severity[sev]
            if items:
                result.append(severity_style(f"  {sev.upper()}: {len(items)}", sev))
                for finding in items[:3]:
                    result.append(severity_style(f"    {finding}", sev))
                if len(items) > 3:
                    result.append(
                        click.style(
                            f"    ... +{len(items) - 3} more", fg="bright_black"
                        )
                    )

        # Info items shown separately
        if findings_by_severity["info"]:
            result.append("")
            for item in findings_by_severity["info"][:3]:
                result.append(click.style(f"    {item}", fg="bright_black"))

    elif scanning_started:
        result.append(click.style("  Scanning in progress...", fg="cyan"))

    return result


# =============================================================================
# NMAP FORMATTER
# =============================================================================
@register_formatter("nmap")
def format_nmap(lines: list, target: str = "") -> list:
    """Format nmap output into friendly summary."""
    result = []
    hosts_found = []
    current_host = None
    open_ports = []
    os_info = None
    scan_type = "port scan"

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue
        if stripped.startswith(("Started:", "Command:", "Timeout:")):
            continue

        # Detect scan type from nmap output
        if "ping scan" in stripped.lower() or "-sn" in stripped:
            scan_type = "host discovery"
        elif "script" in stripped.lower() or "vuln" in stripped.lower():
            scan_type = "vulnerability scan"

        # Host discovery
        if "Nmap scan report for" in stripped:
            # Save previous host if exists
            if current_host and open_ports:
                hosts_found.append(
                    {"host": current_host, "ports": open_ports.copy(), "os": os_info}
                )
                open_ports = []
                os_info = None

            # Extract host
            match = re.search(r"for\s+(\S+)", stripped)
            if match:
                current_host = match.group(1)

        # Open ports
        port_match = re.match(r"(\d+)/(tcp|udp)\s+open\s+(\S+)", stripped)
        if port_match:
            port, proto, service = port_match.groups()
            open_ports.append(f"{port}/{proto} {service}")

        # OS detection
        if stripped.startswith("OS:") or "Running:" in stripped:
            os_match = re.search(r"(?:OS:|Running:)\s*(.+)", stripped)
            if os_match:
                os_info = os_match.group(1).strip()

        # Host is up
        if "Host is up" in stripped:
            if current_host and current_host not in [
                h.get("host") for h in hosts_found
            ]:
                # For ping scans, just note the host is up
                if not open_ports:
                    hosts_found.append(
                        {"host": current_host, "ports": [], "os": None, "up": True}
                    )

    # Don't forget the last host
    if current_host and (
        open_ports or current_host not in [h.get("host") for h in hosts_found]
    ):
        hosts_found.append({"host": current_host, "ports": open_ports, "os": os_info})

    # Build formatted output
    if target:
        result.append(click.style(f"  {scan_type.title()}: {target}", fg="cyan"))
        result.append("")

    if hosts_found:
        for host_info in hosts_found[:5]:  # Limit to 5 hosts in live view
            host = host_info.get("host", "unknown")
            ports = host_info.get("ports", [])
            os_det = host_info.get("os")

            if ports:
                result.append(click.style(f"  {host}", fg="green", bold=True))
                for port in ports[:8]:  # Limit ports shown
                    result.append(click.style(f"      {port}", fg="white"))
                if len(ports) > 8:
                    result.append(
                        click.style(
                            f"      ... +{len(ports) - 8} more ports", fg="bright_black"
                        )
                    )
                if os_det:
                    result.append(
                        click.style(f"      OS: {os_det[:50]}", fg="bright_black")
                    )
            else:
                result.append(click.style(f"  {host} - up", fg="green"))

            result.append("")

        if len(hosts_found) > 5:
            result.append(
                click.style(
                    f"  ... and {len(hosts_found) - 5} more hosts", fg="bright_black"
                )
            )
    else:
        result.append(click.style("  Scanning in progress...", fg="cyan"))

    return result


# =============================================================================
# MSF AUXILIARY (BRUTE FORCE) FORMATTER
# =============================================================================
@register_formatter("msf_auxiliary")
def format_msf_auxiliary(lines: list, target: str = "") -> list:
    """Format MSF auxiliary module output (brute force, enum, etc.)."""
    result = []
    successes = []
    failures = 0
    info_items = []
    module_name = None
    last_attempt = None

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue
        if stripped.startswith(("Started:", "Command:", "Timeout:")):
            continue

        # Detect module name
        if "auxiliary(" in stripped.lower() or "using auxiliary" in stripped.lower():
            match = re.search(r"auxiliary[/(]([^)\s]+)", stripped.lower())
            if match:
                module_name = match.group(1)

        # Success! (credential found, share accessible, etc.)
        if stripped.startswith("[+]"):
            content = stripped[3:].strip()
            # Extract credential if present
            cred_match = re.search(r"['\"]([^'\"]+:[^'\"]+)['\"]", content)
            if cred_match:
                successes.append(f"Valid: {cred_match.group(1)}")
            else:
                successes.append(content[:60])

        # Failed attempt
        elif stripped.startswith("[-]") and "Failed" in stripped:
            failures += 1
            # Track last attempt for context
            cred_match = re.search(r"['\"]([^'\"]+)['\"]", stripped)
            if cred_match:
                last_attempt = cred_match.group(1)

        # Info/status
        elif stripped.startswith("[*]"):
            content = stripped[3:].strip()
            # Skip noisy lines
            if "Scanned" in content or "Starting" in content:
                continue
            if len(content) > 60:
                content = content[:57] + "..."
            info_items.append(content)

        # SMB share enumeration results
        elif "READ" in stripped or "WRITE" in stripped:
            info_items.append(stripped[:60])

    # Build formatted output
    scan_desc = module_name or "auxiliary scan"
    if target:
        result.append(click.style(f"  {scan_desc}: {target}", fg="cyan"))
        result.append("")

    # Show successes prominently
    if successes:
        result.append(click.style("  Successes:", fg="green", bold=True))
        for success in successes[:5]:
            result.append(click.style(f"    {success}", fg="green"))
        result.append("")

    # Show failure count (collapsed)
    if failures > 0:
        msg = f"  {failures} failed attempts"
        if last_attempt:
            msg += f" (last: {last_attempt})"
        result.append(click.style(msg, fg="bright_black"))

    # Show info items
    if info_items:
        for item in info_items[:5]:
            result.append(click.style(f"    {item}", fg="white"))

    if not successes and not info_items and failures == 0:
        result.append(click.style("  Running...", fg="cyan"))

    return result


# =============================================================================
# GOBUSTER FORMATTER
# =============================================================================
@register_formatter("gobuster")
def format_gobuster(lines: list, target: str = "") -> list:
    """Format gobuster output into friendly summary."""
    result = []
    found_paths = []
    status_counts = {"2xx": 0, "3xx": 0, "4xx": 0, "5xx": 0}
    progress = None

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue
        if stripped.startswith(("Started:", "Command:", "Timeout:")):
            continue

        # Found paths (Status: 200, 301, etc.)
        match = re.search(r"(/\S*)\s+\(Status:\s*(\d+)", stripped)
        if match:
            path, status = match.groups()
            status_int = int(status)

            if 200 <= status_int < 300:
                status_counts["2xx"] += 1
                found_paths.append((path, status, "green"))
            elif 300 <= status_int < 400:
                status_counts["3xx"] += 1
                found_paths.append((path, status, "yellow"))
            elif 400 <= status_int < 500:
                status_counts["4xx"] += 1
            else:
                status_counts["5xx"] += 1

        # Progress line
        if "Progress:" in stripped:
            prog_match = re.search(r"Progress:\s*(\d+)\s*/\s*(\d+)", stripped)
            if prog_match:
                current, total = int(prog_match.group(1)), int(prog_match.group(2))
                progress = (current, total)

    # Build formatted output
    if target:
        result.append(click.style(f"  Directory scan: {target}", fg="cyan"))
        result.append("")

    if found_paths:
        result.append(click.style("  Found:", fg="white", bold=True))
        for path, status, color in found_paths[:10]:
            result.append(click.style(f"    [{status}] {path}", fg=color))
        if len(found_paths) > 10:
            result.append(
                click.style(f"    ... +{len(found_paths) - 10} more", fg="bright_black")
            )
        result.append("")

    # Show summary counts
    if any(status_counts.values()):
        summary = []
        if status_counts["2xx"]:
            summary.append(click.style(f"{status_counts['2xx']} OK", fg="green"))
        if status_counts["3xx"]:
            summary.append(
                click.style(f"{status_counts['3xx']} redirects", fg="yellow")
            )
        if summary:
            result.append("  " + " | ".join(summary))

    if progress:
        current, total = progress
        bar = render_progress_bar(current, total, width=25)
        result.append(click.style(f"  {bar}", fg="cyan"))
    elif found_paths:
        # Show running indicator with count
        result.append(
            click.style(f"  Scanning... ({len(found_paths)} found)", fg="cyan")
        )
    else:
        result.append(click.style("  Scanning...", fg="cyan"))

    return result


# =============================================================================
# HYDRA FORMATTER
# =============================================================================
@register_formatter("hydra")
def format_hydra(lines: list, target: str = "") -> list:
    """Format hydra output into friendly summary."""
    result = []
    valid_creds = []
    attempts = 0
    service = None

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue

        # Detect service
        if "Hydra" in stripped and "starting" in stripped.lower():
            match = re.search(r"starting\s+(\S+)", stripped.lower())
            if match:
                service = match.group(1)

        # Valid credentials found
        if stripped.startswith("[") and "host:" in stripped.lower():
            # Extract login and password
            login_match = re.search(r"login:\s*(\S+)", stripped)
            pass_match = re.search(r"password:\s*(\S*)", stripped)
            if login_match:
                login = login_match.group(1)
                password = pass_match.group(1) if pass_match else ""
                valid_creds.append(f"{login}:{password}")

        # Count attempts from DATA lines
        if stripped.startswith("[DATA]"):
            attempts += 16  # Hydra typically does 16 tasks

    # Build formatted output
    desc = f"{service} brute force" if service else "Brute force"
    if target:
        result.append(click.style(f"  {desc}: {target}", fg="cyan"))
        result.append("")

    if valid_creds:
        result.append(click.style("  Valid credentials:", fg="green", bold=True))
        for cred in valid_creds[:5]:
            result.append(click.style(f"    {cred}", fg="green"))
        result.append("")

    if attempts > 0:
        result.append(click.style(f"  ~{attempts} attempts made", fg="bright_black"))

    if not valid_creds and attempts == 0:
        result.append(click.style("  Starting...", fg="cyan"))

    return result


# =============================================================================
# NUCLEI FORMATTER
# =============================================================================
@register_formatter("nuclei")
def format_nuclei(lines: list, target: str = "") -> list:
    """Format nuclei output into friendly summary."""
    result = []
    findings = {"critical": [], "high": [], "medium": [], "low": [], "info": []}
    templates_loaded = 0
    scan_completed = False
    status_msg = None

    for line in lines:
        # Strip ANSI codes for parsing
        stripped = strip_ansi(line.strip())

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue

        # Parse nuclei findings: [template-id] [severity] [protocol] url
        # Example: [CVE-2021-44228] [critical] [http] http://target/path
        # But skip [INF], [WRN], [ERR] status lines
        if not any(stripped.startswith(f"[{s}]") for s in ["INF", "WRN", "ERR"]):
            match = re.search(r"\[([^\]]+)\]\s*\[(\w+)\]", stripped)
            if match:
                template_id, severity = match.groups()
                severity = severity.lower()
                if severity in findings:
                    findings[severity].append(template_id)

        # Templates loaded: "[INF] Templates loaded for current scan: 5739"
        tmpl_match = re.search(
            r"Templates loaded[^:]*:\s*(\d+)", stripped, re.IGNORECASE
        )
        if tmpl_match:
            templates_loaded = int(tmpl_match.group(1))

        # Scan status
        if "Scan completed" in stripped:
            scan_completed = True
            # Extract match count: "0 matches found"
            match_count = re.search(r"(\d+)\s+match", stripped)
            if match_count:
                status_msg = f"{match_count.group(1)} vulnerabilities found"

        # Executing templates
        if "Executing" in stripped and "templates" in stripped:
            exec_match = re.search(r"Executing\s+(\d+)", stripped)
            if exec_match:
                status_msg = f"Running {exec_match.group(1)} templates..."

    # Build formatted output
    if target:
        result.append(click.style(f"  Vulnerability scan: {target}", fg="cyan"))
        result.append("")

    total = sum(len(v) for v in findings.values())
    if total > 0:
        if findings["critical"]:
            result.append(
                severity_style(f"  CRITICAL: {len(findings['critical'])}", "critical")
            )
            for f in findings["critical"][:3]:
                result.append(severity_style(f"    {f}", "critical"))

        if findings["high"]:
            result.append(severity_style(f"  HIGH: {len(findings['high'])}", "high"))
            for f in findings["high"][:3]:
                result.append(severity_style(f"    {f}", "high"))

        if findings["medium"]:
            result.append(
                severity_style(f"  MEDIUM: {len(findings['medium'])}", "medium")
            )
            for f in findings["medium"][:3]:
                result.append(severity_style(f"    {f}", "medium"))

        if findings["low"]:
            result.append(severity_style(f"  LOW: {len(findings['low'])}", "low"))

        if findings["info"]:
            result.append(severity_style(f"  INFO: {len(findings['info'])}", "info"))

        result.append("")

    # Status line
    if scan_completed and status_msg:
        result.append(click.style(f"  {status_msg}", fg="green"))
    elif status_msg:
        result.append(click.style(f"  {status_msg}", fg="cyan"))
    elif templates_loaded > 0:
        result.append(
            click.style(f"  {templates_loaded} templates loaded", fg="bright_black")
        )
    elif total == 0:
        result.append(click.style("  Scanning templates...", fg="cyan"))

    return result


# =============================================================================
# SMBMAP FORMATTER
# =============================================================================
@register_formatter("smbmap")
def format_smbmap(lines: list, target: str = "") -> list:
    """Format smbmap output into friendly summary."""
    result = []
    shares = []
    current_host = None

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue

        # Host header
        if "Working on it" in stripped or "IP:" in stripped:
            match = re.search(r"(\d+\.\d+\.\d+\.\d+)", stripped)
            if match:
                current_host = match.group(1)

        # Share listing
        # Format: SHARENAME    READ/WRITE/NO ACCESS    Comment
        share_match = re.match(
            r"^\s*(\S+)\s+(READ|WRITE|NO ACCESS|READ, WRITE)", stripped
        )
        if share_match:
            share_name, access = share_match.groups()
            if share_name not in ("Disk", "Share"):  # Skip header
                shares.append((share_name, access))

    # Build formatted output
    if target:
        result.append(click.style(f"  SMB enumeration: {target}", fg="cyan"))
        result.append("")

    if shares:
        result.append(click.style("  Shares found:", fg="white", bold=True))
        for share, access in shares:
            if "WRITE" in access:
                color = "green"
                icon = ""
            elif "READ" in access:
                color = "yellow"
                icon = ""
            else:
                color = "bright_black"
                icon = ""
            result.append(click.style(f"    {icon} {share} ({access})", fg=color))
    else:
        result.append(click.style("  Enumerating shares...", fg="cyan"))

    return result


# =============================================================================
# NETEXEC / CRACKMAPEXEC FORMATTER
# =============================================================================
@register_formatter("netexec")
def format_netexec(lines: list, target: str = "") -> list:
    """Format netexec/crackmapexec output."""
    result = []
    successes = []
    info_items = []
    protocol = None

    for line in lines:
        stripped = line.strip()

        # Skip metadata
        if not stripped or stripped.startswith(("===", "Target:", "Args:", "Label:")):
            continue

        # Detect protocol from output
        proto_match = re.match(r"^(SMB|LDAP|WINRM|SSH|MSSQL|RDP)", stripped)
        if proto_match:
            protocol = proto_match.group(1)

        # Success markers (Pwn3d!, +, etc.)
        if "Pwn3d!" in stripped or "(Pwn3d!)" in stripped:
            successes.append("Admin access!")
        elif "[+]" in stripped:
            # Extract the key info
            content = stripped.split("[+]")[-1].strip()
            successes.append(content[:50])
        elif "[*]" in stripped:
            content = stripped.split("[*]")[-1].strip()
            if len(content) < 60:
                info_items.append(content)

    # Build formatted output
    desc = f"{protocol} scan" if protocol else "Network scan"
    if target:
        result.append(click.style(f"  {desc}: {target}", fg="cyan"))
        result.append("")

    if successes:
        for s in successes[:5]:
            result.append(click.style(f"  {s}", fg="green"))

    if info_items:
        for item in info_items[:5]:
            result.append(click.style(f"    {item}", fg="white"))

    if not successes and not info_items:
        result.append(click.style("  Scanning...", fg="cyan"))

    return result


# Alias crackmapexec to netexec
FORMATTERS["crackmapexec"] = format_netexec
