"""
PRAQTOR MCP-S — CLI Scanner
=============================
Command-line interface for scanning MCP servers and configurations.

Usage:
  praqtor-mcp scan-url https://mcp.example.com/sse --api-url https://praqtormcp.fly.dev --api-key KEY
  praqtor-mcp scan config.json --api-url https://praqtormcp.fly.dev --api-key KEY
  praqtor-mcp health --api-url https://praqtormcp.fly.dev --api-key KEY

Env vars (alternative to flags):
  PRAQTOR_API_URL=https://praqtormcp.fly.dev
  PRAQTOR_API_KEY=your-api-key

Output formats: text (default), json, sarif
"""

import json
import sys
import os
import argparse
from datetime import datetime

# ANSI colors
RED = "\033[91m"
ORANGE = "\033[38;5;208m"
YELLOW = "\033[93m"
GREEN = "\033[92m"
CYAN = "\033[96m"
WHITE = "\033[97m"
DIM = "\033[90m"
BOLD = "\033[1m"
RESET = "\033[0m"
BG_RED = "\033[41m"
BG_GREEN = "\033[42m"
MAGENTA = "\033[95m"

SEVERITY_COLORS = {"critical": RED, "high": ORANGE, "medium": YELLOW, "low": GREEN}
SEVERITY_ICONS = {"critical": "●", "high": "●", "medium": "●", "low": "●"}

try:
    from praqtor_mcp import __version__ as VERSION
except ImportError:
    VERSION = "1.0.0"

DEFAULT_API_URL = "https://praqtormcp.fly.dev"


def print_banner():
    print(f"""{RED}
  ╔═══════════════════════════════════════════════════╗
  ║  {WHITE}PRAQTOR MCP-S{RED} — AI Agent Security Scanner       ║
  ║  {DIM}8 engines · attack chains · zero false positives{RED} ║
  ║  {DIM}v{VERSION}{RED}                                          ║
  ╚═══════════════════════════════════════════════════╝{RESET}
""")


# ===== Remote API =====

def _api_request(url: str, api_key: str, method: str = "GET", data: dict = None, timeout: int = 120) -> dict:
    """Make authenticated API request."""
    import urllib.request
    import urllib.error

    body = json.dumps(data).encode("utf-8") if data else None
    req = urllib.request.Request(url, data=body, method=method)
    req.add_header("X-API-Key", api_key)
    req.add_header("X-Scan-Source", "cli")
    if body:
        req.add_header("Content-Type", "application/json")

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        body_text = e.read().decode()
        try:
            err = json.loads(body_text)
            detail = err.get("detail", {})
            msg = detail.get("message", body_text) if isinstance(detail, dict) else str(detail)
        except Exception:
            msg = body_text
        print(f"\n  {RED}API Error ({e.code}): {msg}{RESET}")
        sys.exit(1)
    except urllib.error.URLError as e:
        print(f"\n  {RED}Connection Error: {e.reason}{RESET}")
        print(f"  {DIM}Check API URL: {url}{RESET}")
        sys.exit(1)


def remote_scan(config: dict, api_url: str, api_key: str) -> dict:
    """Run config scan via PRAQTOR MCP API."""
    url = api_url.rstrip("/") + "/api/v1/scan"
    return _api_request(url, api_key, method="POST", data=config)


def remote_scan_url(mcp_url: str, api_url: str, api_key: str) -> dict:
    """Run URL scan via PRAQTOR MCP API — connects to live MCP server."""
    url = api_url.rstrip("/") + "/api/v1/scan/url"
    return _api_request(url, api_key, method="POST", data={"url": mcp_url})


def remote_health(api_url: str, api_key: str) -> dict:
    """Check API health."""
    url = api_url.rstrip("/") + "/api/v1/health"
    return _api_request(url, api_key, timeout=10)


# ===== Output Formatters =====

def print_summary(data: dict):
    s = data["summary"]
    sev = s["by_severity"]
    score = data.get("security_score", 0)

    if score >= 80:
        score_clr = GREEN
    elif score >= 50:
        score_clr = YELLOW
    else:
        score_clr = RED

    print(f"  {BOLD}Scan ID:{RESET}        {data['scan_id']}")
    print(f"  {BOLD}Security Score:{RESET}  {score_clr}{BOLD}{score}/100{RESET}")
    print(f"  {BOLD}Servers:{RESET}         {s['total_servers']}  {BOLD}Tools:{RESET}  {s['total_tools']}")
    print()

    print(f"  {BOLD}FINDINGS{RESET}  {s['total_findings']} total")
    print(f"    {RED}{SEVERITY_ICONS['critical']} Critical: {sev.get('critical',0):3d}{RESET}   "
          f"{ORANGE}{SEVERITY_ICONS['high']} High: {sev.get('high',0):3d}{RESET}   "
          f"{YELLOW}{SEVERITY_ICONS['medium']} Medium: {sev.get('medium',0):3d}{RESET}   "
          f"{GREEN}{SEVERITY_ICONS['low']} Low: {sev.get('low',0):3d}{RESET}")
    print()

    chains = s.get("total_attack_chains", 0)
    if chains > 0:
        print(f"  {RED}{BOLD}⚠ ATTACK CHAINS: {chains}{RESET}  "
              f"({s.get('enables_edges', 0)} ENABLES edges)")
    else:
        print(f"  {GREEN}✓ No attack chains detected{RESET}")
    print()

    print(f"  {BOLD}ENGINE COVERAGE{RESET}")
    for eng, count in sorted(s.get("by_engine", {}).items()):
        bar = "█" * min(count, 30) + "░" * max(0, 30 - count)
        print(f"    {eng:4s}  {bar} {count}")
    print()


def print_findings(data: dict, severity_filter: str = None, limit: int = 50):
    findings = data.get("findings", [])
    if severity_filter:
        findings = [f for f in findings if f["severity"] == severity_filter]

    findings.sort(key=lambda f: (
        {"critical": 0, "high": 1, "medium": 2, "low": 3}.get(f["severity"], 9),
        -f.get("confidence", 0)
    ))

    findings = findings[:limit]
    print(f"  {BOLD}TOP FINDINGS ({len(findings)} shown){RESET}")
    print(f"  {'─' * 70}")

    for f in findings:
        color = SEVERITY_COLORS.get(f["severity"], WHITE)
        icon = SEVERITY_ICONS.get(f["severity"], "?")
        code = f.get("code", f.get("id", ""))
        module = f.get("module", "")
        print(f"  {color}{icon} [{f['severity'].upper():8s}]{RESET} {code} | {module} | "
              f"conf={f.get('confidence',0):.2f}")
        print(f"    {WHITE}{f['title']}{RESET}")
        if f.get("server_id"):
            print(f"    {DIM}{f.get('server_id','')} → {f.get('tool_id','')}{RESET}")
        print()


def print_chains(data: dict):
    chains = data.get("attack_chains", [])
    if not chains:
        return

    print(f"  {RED}{BOLD}ATTACK CHAINS{RESET}")
    print(f"  {'─' * 70}")

    for c in chains:
        sev = c.get("max_severity", "medium")
        color = SEVERITY_COLORS.get(sev, WHITE)
        print(f"\n  {color}{BOLD}{c['id']}{RESET} | {color}{sev.upper()}{RESET} | "
              f"depth={c.get('chain_depth',0)} | {c.get('finding_count',0)} findings")
        print(f"    {WHITE}{c['title']}{RESET}")
        if c.get("description"):
            print(f"    {DIM}{c['description'][:120]}{RESET}")
        print(f"    {DIM}Modules: {', '.join(c.get('modules_involved', []))}{RESET}")
        print(f"    {DIM}Servers: {', '.join(c.get('servers_involved', []))}{RESET}")
        if c.get("is_cross_server"):
            print(f"    {RED}⚠ CROSS-SERVER{RESET}")
    print()


def print_health(data: dict):
    print(f"  {GREEN}✓ PRAQTOR MCP-S API is healthy{RESET}")
    print(f"  {BOLD}Version:{RESET}   {data.get('version', '?')}")
    print(f"  {BOLD}Engines:{RESET}   {data.get('engines', {}).get('count', '?')}")
    engines = data.get("engines", {}).get("names", {})
    for code, name in engines.items():
        print(f"    {CYAN}{code}{RESET}  {name}")
    print()


def format_sarif(data: dict) -> dict:
    """Convert to SARIF v2.1.0 for GitHub Advanced Security."""
    rules = []
    results_list = []
    rule_ids = set()

    for f in data.get("findings", []):
        rule_id = f.get("code", f.get("id", ""))
        if rule_id not in rule_ids:
            rule_ids.add(rule_id)
            rules.append({
                "id": rule_id,
                "name": f.get("module", ""),
                "shortDescription": {"text": f["title"][:200]},
                "defaultConfiguration": {
                    "level": {"critical": "error", "high": "error", "medium": "warning", "low": "note"}.get(f["severity"], "note")
                }
            })

        results_list.append({
            "ruleId": rule_id,
            "message": {"text": f["title"]},
            "level": {"critical": "error", "high": "error", "medium": "warning", "low": "note"}.get(f["severity"], "note"),
            "properties": {
                "module": f.get("module", ""),
                "severity": f["severity"],
                "confidence": f.get("confidence", 0),
                "server_id": f.get("server_id", ""),
                "tool_id": f.get("tool_id", ""),
            }
        })

    return {
        "$schema": "https://raw.githubusercontent.com/oasis-tcs/sarif-spec/main/sarif-2.1/schema/sarif-schema-2.1.0.json",
        "version": "2.1.0",
        "runs": [{
            "tool": {
                "driver": {
                    "name": "PRAQTOR MCP-S",
                    "version": VERSION,
                    "informationUri": "https://praqtormcp.fly.dev",
                    "rules": rules
                }
            },
            "results": results_list
        }]
    }


def _exit_on_severity(data: dict):
    """Print pass/fail and exit with appropriate code."""
    sev = data.get("summary", {}).get("by_severity", {})
    if sev.get("critical", 0) > 0:
        print(f"  {RED}{BOLD}✗ FAILED — {sev['critical']} critical findings{RESET}\n")
        sys.exit(2)
    elif sev.get("high", 0) > 0:
        print(f"  {ORANGE}{BOLD}⚠ WARNING — {sev['high']} high findings{RESET}\n")
        sys.exit(1)
    else:
        print(f"  {GREEN}{BOLD}✓ PASSED — no critical or high findings{RESET}\n")
        sys.exit(0)


def _output_results(data: dict, fmt: str, severity_filter: str = None, limit: int = 50):
    """Output results in requested format."""
    if fmt == "json":
        print(json.dumps(data, indent=2))
    elif fmt == "sarif":
        print(json.dumps(format_sarif(data), indent=2))
    else:
        print_summary(data)
        print_findings(data, severity_filter=severity_filter, limit=limit)
        print_chains(data)
        _exit_on_severity(data)


# ===== Main CLI =====

def main():
    parser = argparse.ArgumentParser(
        prog="praqtor-mcp",
        description="PRAQTOR MCP-S — AI Agent Security Scanner",
        epilog="Examples:\n"
               "  praqtor-mcp scan-url https://mcp.example.com/sse\n"
               "  praqtor-mcp scan config.json\n"
               "  praqtor-mcp scan config.json --format sarif > results.sarif\n"
               "  praqtor-mcp health\n\n"
               "Set PRAQTOR_API_URL and PRAQTOR_API_KEY env vars to skip --api-url/--api-key flags.\n",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    sub = parser.add_subparsers(dest="command", help="Command to run")

    # Common API args
    def add_api_args(p):
        p.add_argument("--api-url", help=f"API URL (env: PRAQTOR_API_URL, default: {DEFAULT_API_URL})")
        p.add_argument("--api-key", help="API key (env: PRAQTOR_API_KEY)")

    def add_output_args(p):
        p.add_argument("--format", choices=["text", "json", "sarif"], default="text",
                        help="Output format (default: text)")
        p.add_argument("--severity", choices=["critical", "high", "medium", "low"],
                        help="Filter findings by severity")
        p.add_argument("--limit", type=int, default=50, help="Max findings to show (default: 50)")

    # scan-url command
    url_parser = sub.add_parser("scan-url", help="Scan a live MCP server by URL (SSE or Streamable HTTP)")
    url_parser.add_argument("url", help="MCP server URL (e.g. https://mcp.example.com/sse)")
    add_output_args(url_parser)
    add_api_args(url_parser)

    # scan command
    scan_parser = sub.add_parser("scan", help="Scan an MCP configuration file for security issues")
    scan_parser.add_argument("config", help="Path to MCP config JSON file")
    add_output_args(scan_parser)
    add_api_args(scan_parser)

    # health command
    health_parser = sub.add_parser("health", help="Check API health and engine status")
    add_api_args(health_parser)

    # version command
    sub.add_parser("version", help="Show version")

    args = parser.parse_args()

    if not args.command:
        print_banner()
        parser.print_help()
        sys.exit(1)

    if args.command == "version":
        print(f"praqtor-mcp v{VERSION}")
        sys.exit(0)

    # Resolve API credentials
    api_url = getattr(args, "api_url", None) or os.environ.get("PRAQTOR_API_URL", DEFAULT_API_URL)
    api_key = getattr(args, "api_key", None) or os.environ.get("PRAQTOR_API_KEY", "")

    if not api_key:
        print(f"  {RED}Error: API key required.{RESET}")
        print(f"  {DIM}Set PRAQTOR_API_KEY env var or use --api-key flag.{RESET}")
        print(f"  {DIM}Get your key at: https://praqtormcp.fly.dev/api{RESET}")
        sys.exit(1)

    # --- health ---
    if args.command == "health":
        print_banner()
        data = remote_health(api_url, api_key)
        print_health(data)
        sys.exit(0)

    # --- scan-url ---
    if args.command == "scan-url":
        if args.format == "text":
            print_banner()
            print(f"  {BOLD}Target:{RESET}  {args.url}")
            print(f"  {DIM}API:     {api_url}{RESET}")
            print(f"  {DIM}Connecting to live MCP server...{RESET}\n")

        data = remote_scan_url(args.url, api_url, api_key)

        if args.format == "text":
            disc = data.get("discovery", {})
            if disc:
                print(f"  {GREEN}✓ Connected{RESET} via {BOLD}{disc.get('transport','?').upper()}{RESET}")
                sinfo = disc.get("server_info", {})
                if sinfo:
                    print(f"  {DIM}Server: {sinfo.get('name','?')} v{sinfo.get('version','?')}{RESET}")
                print(f"  {DIM}Tools discovered: {disc.get('tool_count',0)}{RESET}\n")

        _output_results(data, args.format, args.severity, args.limit)

    # --- scan ---
    if args.command == "scan":
        try:
            with open(args.config) as f:
                config = json.load(f)
        except FileNotFoundError:
            print(f"  {RED}Error: File not found: {args.config}{RESET}")
            sys.exit(1)
        except json.JSONDecodeError as e:
            print(f"  {RED}Error: Invalid JSON in {args.config}: {e}{RESET}")
            sys.exit(1)

        if args.format == "text":
            print_banner()
            print(f"  {BOLD}Config:{RESET}  {args.config}")
            print(f"  {DIM}API:     {api_url}{RESET}")
            print(f"  {DIM}Scanning...{RESET}\n")

        data = remote_scan(config, api_url, api_key)
        _output_results(data, args.format, args.severity, args.limit)


if __name__ == "__main__":
    main()
