from __future__ import annotations

import sys
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import List

import click

from bastion import __version__
from bastion.artifacts import (
    append_event,
    write_aibom,
    write_capability_surface,
    write_drift_summary,
    write_security_state,
)
from bastion.config import get_bastion_dir, get_cloud_url, is_initialized, save_config
from bastion.policy import load_policy, add_acknowledgment, remove_acknowledgment, is_acknowledged
from bastion.drift import compute_drift, load_baseline, save_baseline
from bastion.models import (
    AIBOM,
    CapabilitySurface,
    ScanResult,
    SecurityState,
)
from bastion.reporting import print_summary


@click.group()
@click.version_option(version=__version__, prog_name="bastion")
def bastion() -> None:
    """Bastion AI — Supply-chain security for AI agent plugins/skills."""
    pass


@bastion.command()
def init() -> None:
    """Initialize a Bastion project in the current directory."""
    if is_initialized():
        click.echo(click.style("Bastion is already initialized in this directory.", fg="yellow"))
        return

    bastion_dir = get_bastion_dir()
    bastion_dir.mkdir(parents=True, exist_ok=True)

    project_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc).isoformat()

    config = {
        "project_id": project_id,
        "created_at": now,
        "version": __version__,
    }
    save_config(config)

    click.echo(click.style("✔ Bastion initialized successfully!", fg="green", bold=True))
    click.echo(f"  Project ID: {project_id}")
    click.echo(f"  Config: {bastion_dir / 'config.yaml'}")
    click.echo("")
    click.echo(f"  Run {click.style('bastion scan', bold=True)} to analyze your plugins for security risks.")
    click.echo(f"  Run {click.style('bastion connect', bold=True)} to link this project to the cloud dashboard")
    click.echo("  for monitoring, trust scores, and shareable verification pages.")


@bastion.command()
@click.option("--path", default=".", help="Path to scan for plugins/skills.")
@click.option("--ci", is_flag=True, default=False, help="Run in CI mode (deterministic output).")
@click.option("--openclaw", is_flag=True, default=False, help="Enable OpenClaw detection pack.")
def scan(path: str, ci: bool, openclaw: bool) -> None:
    """Scan for plugin/skill supply-chain risks."""
    if not is_initialized():
        click.echo(click.style("Error: Bastion is not initialized. Run 'bastion init' first.", fg="red"))
        sys.exit(1)

    from bastion.config import load_config

    config = load_config()
    project_id = config["project_id"]
    bastion_dir = get_bastion_dir()
    now = datetime.now(timezone.utc).isoformat()
    scan_id = str(uuid.uuid4())

    click.echo(click.style("Scanning...", fg="blue"))

    discovered_plugins = _discover_plugins(path, openclaw)

    scan_results: List[ScanResult] = []
    plugin_map = {}
    for plugin in discovered_plugins:
        result = _analyze_plugin(plugin, now)
        scan_results.append(result)
        plugin_map[plugin.id] = plugin

    baseline = load_baseline(bastion_dir)
    is_first_scan = baseline is None
    drift_events = []

    if is_first_scan:
        save_baseline(scan_results, bastion_dir, label="Initial Baseline (Unverified)")
        click.echo(click.style("  Created initial baseline (unverified).", dim=True))
    else:
        drift_events = compute_drift(scan_results, baseline)
        save_baseline(scan_results, bastion_dir)

    risk_summary = {"Critical": 0, "High": 0, "Medium": 0, "Low": 0}
    top_findings = []
    for r in scan_results:
        if r.risk_level in risk_summary:
            risk_summary[r.risk_level] += 1
        if r.risk_level in ("Critical", "High"):
            top_findings.append(r.to_dict())

    security_state = SecurityState(
        project_id=project_id,
        scan_id=scan_id,
        timestamp=now,
        total_plugins=len(scan_results),
        risk_summary=risk_summary,
        top_findings=top_findings[:10],
    )

    plugin_capabilities = []
    aibom_components = []
    for r in scan_results:
        plugin_capabilities.append({
            "plugin_id": r.plugin_id,
            "capabilities": [c.to_dict() for c in r.capabilities],
            "risk_level": r.risk_level,
        })
        aibom_components.append({
            "id": r.plugin_id,
            "capabilities": [c.category for c in r.capabilities],
            "risk_level": r.risk_level,
        })

    capability_surface = CapabilitySurface(
        project_id=project_id,
        plugins=plugin_capabilities,
    )

    aibom = AIBOM(
        version="1.0",
        project_id=project_id,
        timestamp=now,
        components=aibom_components,
    )

    write_security_state(security_state, bastion_dir)
    write_capability_surface(capability_surface, bastion_dir)
    write_drift_summary(drift_events, bastion_dir)
    write_aibom(aibom, bastion_dir)

    event = {
        "type": "scan",
        "scan_id": scan_id,
        "timestamp": now,
        "total_plugins": len(scan_results),
        "risk_summary": risk_summary,
        "is_first_scan": is_first_scan,
    }
    append_event(event, bastion_dir)

    policy = load_policy(bastion_dir)
    print_summary(scan_results, None if is_first_scan else drift_events, policy=policy)

    if config.get("cloud_connected"):
        try:
            import requests as req

            cloud_server = config.get("cloud_server", "").rstrip("/")
            cloud_project_id = config.get("cloud_project_id", "")
            cloud_api_key = config.get("cloud_api_key", "")

            if not cloud_server or not cloud_project_id or not cloud_api_key:
                click.echo(click.style("⚠ Cloud config incomplete. Run 'bastion connect' again.", fg="yellow"))
            else:
                cloud_risk_summary = {k.lower(): v for k, v in risk_summary.items()}

                plugins_payload = []
                for r in scan_results:
                    plugin = plugin_map.get(r.plugin_id)
                    plugins_payload.append({
                        "name": plugin.name if plugin else r.plugin_id,
                        "path": plugin.path if plugin else None,
                        "risk_level": r.risk_level,
                        "capabilities": [c.to_dict() for c in r.capabilities],
                        "hash": plugin.hash if plugin else None,
                        "language": plugin.language if plugin else None,
                        "type": plugin.type if plugin else None,
                        "dependencies": r.dependencies,
                        "skills": r.skills,
                        "description": r.description,
                        "file_count": r.file_count,
                        "total_lines": r.total_lines,
                    })

                drift_payload = []
                for d in drift_events:
                    drift_payload.append({
                        "event_type": d.event_type,
                        "plugin_name": d.plugin_id,
                        "severity": d.severity,
                        "details": {"capability_changes": d.capability_changes},
                    })

                upload_data = {
                    "scan_data": security_state.to_dict(),
                    "risk_summary": cloud_risk_summary,
                    "total_plugins": len(scan_results),
                    "plugins": plugins_payload,
                    "drift_events": drift_payload,
                }

                resp = req.post(
                    f"{cloud_server}/api/projects/{cloud_project_id}/scans",
                    json=upload_data,
                    headers={"X-API-Key": cloud_api_key},
                    timeout=30,
                )

                if resp.status_code == 201:
                    upload_result = resp.json()
                    trust_score = upload_result.get('trust_score')
                    status = upload_result.get('verification_status', '')
                    public_id = upload_result.get('public_id', config.get('cloud_public_id', ''))
                    click.echo(click.style("✔ Scan uploaded to Bastion Cloud.", fg="green"))
                    if trust_score is not None:
                        click.echo(f"  Trust Score: {trust_score}/100")
                    click.echo(f"  Status: {status.replace('_', ' ').title()}")
                    click.echo(f"  Dashboard: {cloud_server}/dashboard")
                    if public_id:
                        verify_url = f"{cloud_server}/verify/{public_id}"
                        click.echo(f"  Verification: {verify_url}")
                        click.echo(click.style("  ^ Share this link with partners to verify your trust score", dim=True))
                    claim_url = f"{cloud_server}/claim?key={cloud_api_key}"
                    click.echo(f"  Claim Project: {claim_url}")
                    click.echo(click.style("  ^ Visit this link and sign in to see the project on your dashboard", dim=True))
                else:
                    error_detail = ""
                    try:
                        error_detail = resp.text[:200]
                    except Exception:
                        pass
                    click.echo(click.style(f"⚠ Cloud upload failed: {resp.status_code} {error_detail}", fg="yellow"))
        except Exception as e:
            click.echo(click.style(f"⚠ Cloud upload error: {e}", fg="yellow"))

    has_critical = risk_summary.get("Critical", 0) > 0 or risk_summary.get("High", 0) > 0
    if has_critical and policy:
        from bastion.policy import filter_acknowledged
        _, ack_count = filter_acknowledged(scan_results, policy)
        # Count total critical/high capabilities
        total_critical_high = sum(
            1 for r in scan_results for c in r.capabilities
            if c.severity in ("Critical", "High")
        )
        if ack_count >= total_critical_high and total_critical_high > 0:
            click.echo(click.style(f"✔ All {total_critical_high} high-risk findings are acknowledged in policy.", fg="green"))
        else:
            unack = total_critical_high - ack_count
            click.echo(click.style(f"⚠ {unack} unacknowledged High/Critical findings detected.", fg="red", bold=True))
            sys.exit(1)
    elif has_critical:
        click.echo(click.style("⚠ High/Critical findings detected.", fg="red", bold=True))
        sys.exit(1)
    else:
        click.echo(click.style("✔ Scan completed. No high-risk findings.", fg="green"))


def _discover_plugins(path: str, openclaw: bool) -> list:
    try:
        from bastion.discovery import discover_plugins
        return discover_plugins(path, openclaw_mode=openclaw)
    except (ImportError, AttributeError):
        return []


def _analyze_plugin(plugin: dict, timestamp: str) -> ScanResult:
    try:
        from bastion.analyzers import analyze_plugin
        return analyze_plugin(plugin, timestamp)
    except (ImportError, AttributeError):
        from bastion.models import ScanResult
        return ScanResult(
            plugin_id=plugin.get("id", plugin.get("name", "unknown")),
            capabilities=[],
            risk_level="Low",
            timestamp=timestamp,
        )


@bastion.command()
@click.option("--project-name", default=None, help="Project name for cloud registration.")
def connect(project_name: str | None) -> None:
    """Connect this project to Bastion Cloud for monitoring and trust scores."""
    if not is_initialized():
        click.echo(click.style("Error: Bastion is not initialized. Run 'bastion init' first.", fg="red"))
        sys.exit(1)

    from bastion.config import load_config

    config = load_config()
    project_id = config["project_id"]

    if config.get("cloud_connected"):
        click.echo(click.style("This project is already connected to Bastion Cloud.", fg="yellow"))
        click.echo(f"  Cloud Project ID: {config.get('cloud_project_id')}")
        click.echo(f"  Server: {config.get('cloud_server')}")
        return

    server = get_cloud_url()

    if not project_name:
        project_name = click.prompt("Project name", default=Path.cwd().name)

    click.echo(click.style("Connecting to Bastion Cloud...", fg="blue"))

    import requests as req

    try:
        resp = req.post(
            f"{server.rstrip('/')}/api/connect",
            json={"project_id": project_id, "name": project_name},
            timeout=30,
        )
        if resp.status_code == 200:
            data = resp.json()
            config["cloud_connected"] = True
            config["cloud_server"] = server
            config["cloud_project_id"] = data.get("cloud_project_id", "")
            config["cloud_api_key"] = data.get("api_key", "")
            config["cloud_public_id"] = data.get("public_id", project_id)
            save_config(config)

            claim_url = f"{server.rstrip('/')}/claim?key={data.get('api_key', '')}"
            click.echo(click.style("✔ Connected to Bastion Cloud!", fg="green", bold=True))
            click.echo(f"  Cloud Project ID: {config['cloud_project_id']}")
            click.echo("")
            click.echo("  Next steps:")
            click.echo(f"    1. Run {click.style('bastion scan', bold=True)} to upload your first scan")
            click.echo(f"    2. Claim your project by visiting this URL:")
            click.echo(click.style(f"       {claim_url}", fg="cyan"))
            click.echo("")
            click.echo(click.style("  Important:", fg="yellow", bold=True) + " Your project won't appear on your dashboard")
            click.echo("  until you claim it. Visit the link above, sign in with Replit,")
            click.echo("  and the project will be linked to your account.")
        else:
            click.echo(click.style(f"Connection failed: {resp.status_code} - {resp.text}", fg="red"))
            sys.exit(1)
    except req.exceptions.ConnectionError:
        click.echo(click.style(f"Could not reach server at {server}", fg="red"))
        sys.exit(1)
    except Exception as e:
        click.echo(click.style(f"Connection error: {e}", fg="red"))
        sys.exit(1)


@bastion.command()
def status() -> None:
    """Show current project status and configuration."""
    if not is_initialized():
        click.echo(click.style("Bastion is not initialized in this directory.", fg="yellow"))
        click.echo("Run 'bastion init' to get started.")
        return

    from bastion.config import load_config

    config = load_config()
    bastion_dir = get_bastion_dir()

    click.echo(click.style("Bastion AI Project Status", fg="blue", bold=True))
    click.echo(f"  Project ID:   {config['project_id']}")
    click.echo(f"  Version:      {config.get('version', 'unknown')}")
    click.echo(f"  Created:      {config.get('created_at', 'unknown')}")

    baseline_path = bastion_dir / "baseline.json"
    if baseline_path.exists():
        click.echo(click.style("  Baseline:     ✔ exists", fg="green"))
    else:
        click.echo(click.style("  Baseline:     ✖ no baseline (run 'bastion scan')", fg="yellow"))

    if config.get("cloud_connected"):
        click.echo(click.style("  Cloud:        ✔ connected", fg="green"))
        click.echo(f"  Cloud Server: {config.get('cloud_server')}")
        click.echo(f"  Dashboard:    {get_cloud_url()}/dashboard")
        public_id = config.get("cloud_public_id", config.get("project_id", ""))
        if public_id:
            click.echo(f"  Verification: {get_cloud_url()}/verify/{public_id}")
            click.echo(click.style("  ^ Share this link with partners to verify your trust score", dim=True))
    else:
        click.echo(click.style("  Cloud:        ✖ not connected (run 'bastion connect')", dim=True))


@bastion.command(name="mcp-serve")
def mcp_serve():
    """Start Bastion as an MCP server for AI agent integration."""
    try:
        from bastion.mcp_server import run_server
    except ImportError:
        click.echo(click.style("Error: MCP support not installed. Install with:", fg="red"))
        click.echo("  pip install bastion-ai-security[mcp]")
        sys.exit(1)
    run_server()


@bastion.group()
def policy():
    """Manage the security policy (allow-list) for acknowledged findings."""
    pass

@policy.command(name="add")
@click.argument("plugin_id")
@click.argument("category")
@click.option("--reason", "-r", required=True, help="Why this capability is acknowledged.")
def policy_add(plugin_id: str, category: str, reason: str):
    """Acknowledge a capability for a plugin in the security policy."""
    if not is_initialized():
        click.echo(click.style("Error: Bastion is not initialized.", fg="red"))
        sys.exit(1)
    bastion_dir = get_bastion_dir()
    valid_categories = [
        "system_execution", "network_access", "secret_access", "dynamic_code",
        "obfuscation", "file_system_access", "sandbox_bypass", "install_risk",
        "hardcoded_secret", "privilege_escalation", "dependency_risk",
    ]
    if category not in valid_categories:
        click.echo(click.style(f"Error: Unknown category '{category}'.", fg="red"))
        click.echo(f"  Valid categories: {', '.join(valid_categories)}")
        sys.exit(1)
    add_acknowledgment(bastion_dir, plugin_id, category, reason)
    click.echo(click.style(f"✔ Acknowledged {category} for {plugin_id}", fg="green"))
    click.echo(click.style(f"  Reason: {reason}", dim=True))

@policy.command(name="remove")
@click.argument("plugin_id")
@click.argument("category")
def policy_remove(plugin_id: str, category: str):
    """Remove an acknowledgment from the security policy."""
    if not is_initialized():
        click.echo(click.style("Error: Bastion is not initialized.", fg="red"))
        sys.exit(1)
    bastion_dir = get_bastion_dir()
    removed = remove_acknowledgment(bastion_dir, plugin_id, category)
    if removed:
        click.echo(click.style(f"✔ Removed acknowledgment for {category} on {plugin_id}", fg="green"))
    else:
        click.echo(click.style(f"No acknowledgment found for {category} on {plugin_id}", fg="yellow"))

@policy.command(name="list")
def policy_list():
    """List all acknowledged capabilities in the security policy."""
    if not is_initialized():
        click.echo(click.style("Error: Bastion is not initialized.", fg="red"))
        sys.exit(1)
    bastion_dir = get_bastion_dir()
    pol = load_policy(bastion_dir)
    acknowledged = pol.get("acknowledged", {})
    if not acknowledged:
        click.echo(click.style("No policy acknowledgments configured.", dim=True))
        click.echo("  Use 'bastion policy add <plugin_id> <category> -r \"reason\"' to add one.")
        return
    click.echo(click.style("Security Policy — Acknowledged Capabilities", fg="blue", bold=True))
    click.echo("")
    for plugin_id, entries in acknowledged.items():
        click.echo(click.style(f"  Plugin: {plugin_id}", bold=True))
        for entry in entries:
            click.echo(f"    • {entry['category']}")
            click.echo(click.style(f"      Reason: {entry.get('reason', 'N/A')}", dim=True))
            click.echo(click.style(f"      Added: {entry.get('acknowledged_at', 'N/A')}", dim=True))
        click.echo("")


def main() -> None:
    bastion()


if __name__ == "__main__":
    main()
