#!/usr/bin/env python3
"""OpenClaw Discovery Tool - Detect OpenClaw AI agent installations via EDR."""
import argparse
import json
import logging
import webbrowser
from datetime import datetime, timezone
from pathlib import Path

from astrix_openclaw_scanner.crowdstrike.aggregator import aggregate_events as aggregate_crowdstrike
from astrix_openclaw_scanner.crowdstrike.client import CrowdStrikeClient
from astrix_openclaw_scanner.crowdstrike.collectors.processes import CrowdStrikeProcessCollector
from astrix_openclaw_scanner.microsoft_defender.aggregator import aggregate_events as aggregate_defender
from astrix_openclaw_scanner.microsoft_defender.client import DefenderClient
from astrix_openclaw_scanner.microsoft_defender.collectors.processes import DefenderProcessCollector
from astrix_openclaw_scanner.report.generator import render_report

BLUE = "\033[94m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
DIM = "\033[2m"
RESET = "\033[0m"
BOLD = "\033[1m"


class _Formatter(logging.Formatter):
    """Colored, clean log formatter."""

    _LEVEL_STYLES = {
        logging.DEBUG: DIM,
        logging.INFO: GREEN,
        logging.WARNING: YELLOW,
        logging.ERROR: RED,
        logging.CRITICAL: RED + BOLD,
    }

    def format(self, record: logging.LogRecord) -> str:
        style = self._LEVEL_STYLES.get(record.levelno, "")
        ts = self.formatTime(record, "%H:%M:%S")
        level = record.levelname
        msg = record.getMessage()
        return f"{DIM}{ts}{RESET} {style}{level}{RESET} {msg}"


def _setup_logging() -> None:
    handler = logging.StreamHandler()
    handler.setFormatter(_Formatter())
    logging.root.addHandler(handler)
    logging.root.setLevel(logging.INFO)


_setup_logging()
logger = logging.getLogger(__name__)

EDR_LABELS = {
    "crowdstrike": "CrowdStrike",
    "microsoft_defender": "Microsoft Defender",
}


def save_raw_events(events: list[str], output_path: Path) -> None:
    """Save raw events to JSON file."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    data = {
        "collected_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "event_count": len(events),
        "events": [json.loads(e) for e in events],
    }
    with open(output_path, "w") as f:
        json.dump(data, f, indent=2)
    logger.info(f"Saved {len(events)} raw events to {output_path}")


def collect_crowdstrike(client_id: str, client_secret: str, days_back: int) -> list[str]:
    """Collect events from CrowdStrike."""
    client = CrowdStrikeClient(client_id=client_id, client_secret=client_secret)
    collector = CrowdStrikeProcessCollector(client=client, days_back=days_back)
    return collector.run()


def collect_defender(client_id: str, client_secret: str, tenant_id: str, days_back: int) -> list[str]:
    """Collect events from Microsoft Defender."""
    client = DefenderClient(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)
    collector = DefenderProcessCollector(client=client, days_back=days_back)
    return collector.run()


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Detect OpenClaw AI agent installations via EDR telemetry"
    )
    parser.add_argument(
        "edr",
        choices=list(EDR_LABELS.keys()),
        help="EDR platform to query",
    )
    parser.add_argument("--client-id", required=True, help="API client ID")
    parser.add_argument("--client-secret", required=True, help="API client secret")
    parser.add_argument("--tenant-id", help="Azure tenant ID (required for microsoft_defender)")
    parser.add_argument("--days-back", type=int, default=14, help="Days to look back (default: 14)")
    parser.add_argument(
        "--no-filter",
        action="store_true",
        default=False,
        help="Disable post-collection noise filtering (show all matched events)",
    )
    parser.add_argument(
        "--save-raw-events",
        type=Path,
        default=None,
        metavar="PATH",
        help="Save raw EDR events to this JSON file (omit to not save)",
    )
    parser.add_argument(
        "--report-path",
        type=Path,
        default=None,
        metavar="DIR",
        help="Directory to write the HTML report (default: system temp dir)",
    )

    args = parser.parse_args()

    if args.edr == "microsoft_defender" and not args.tenant_id:
        parser.error("--tenant-id is required for microsoft_defender")

    return args


def main() -> None:
    """Main entry point."""
    args = parse_args()

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")

    # Collect events
    logger.info(f"Collecting from {EDR_LABELS[args.edr]} (days_back={args.days_back})")

    if args.edr == "crowdstrike":
        events = collect_crowdstrike(args.client_id, args.client_secret, args.days_back)
        endpoints = aggregate_crowdstrike(events, skip_filter=args.no_filter)
    else:  # microsoft_defender
        events = collect_defender(args.client_id, args.client_secret, args.tenant_id, args.days_back)
        endpoints = aggregate_defender(events, skip_filter=args.no_filter)

    # Save raw events only if a path was given
    if args.save_raw_events is not None:
        save_raw_events(events, args.save_raw_events)

    # Generate and open report (temp dir unless --report-path given)
    if endpoints:
        if args.report_path is not None:
            report_dir = args.report_path
        else:
            report_dir = Path.cwd() / "astrix_openclaw_scanner_results" / args.edr / f"report_{timestamp}"
        render_report(endpoints, report_dir, edr_source=EDR_LABELS[args.edr])

        report_path = report_dir / "index.html"
        logger.info(f"Opening report: {report_path}")
        webbrowser.open(report_path.as_uri())

    logger.info("Collection complete!")


if __name__ == "__main__":
    main()
