"""
ZeroShot SDK - CLI Interface.


Provides command-line access to ZeroShot functionality.


Usage:
   zeroshot auth login
   zeroshot scan --target http://localhost:8000/chat
   zeroshot scans list
"""


import typer
from rich.console import Console
from rich.table import Table
from typing import Optional


from zeroshot_sdk.auth import (
   get_api_key,
   save_api_key,
   get_api_url,
   save_api_url,
   clear_config,
   get_config_path,
)


app = typer.Typer(
   name="zeroshot",
   help="ZeroShot SDK - AI Security Testing Platform",
   no_args_is_help=True,
)


console = Console()




# ==============================================================================
# Auth Commands
# ==============================================================================


auth_app = typer.Typer(help="Authentication commands")
app.add_typer(auth_app, name="auth")




@auth_app.command("login")
def auth_login(
   api_key: Optional[str] = typer.Option(None, "--api-key", "-k", help="API key (zsk_...)"),
   api_url: Optional[str] = typer.Option(None, "--api-url", "-u", help="API base URL"),
):
   """
   Login with API key.
  
   If not provided, prompts for API key interactively.
   """
   if not api_key:
       api_key = typer.prompt("Enter your API key (zsk_...)")
  
   if not api_key.startswith("zsk_"):
       console.print("[red]Error:[/red] API key should start with 'zsk_'")
       raise typer.Exit(1)
  
   save_api_key(api_key)
  
   if api_url:
       save_api_url(api_url)
  
   console.print(f"[green]✓[/green] API key saved to {get_config_path()}")
   console.print(f"  API URL: {get_api_url()}")




@auth_app.command("status")
def auth_status():
   """Show current authentication status."""
   api_key = get_api_key()
   api_url = get_api_url()
  
   if api_key:
       # Show masked key
       masked = api_key[:8] + "..." + api_key[-4:]
       console.print(f"[green]✓[/green] Authenticated")
       console.print(f"  API Key: {masked}")
       console.print(f"  API URL: {api_url}")
   else:
       console.print("[yellow]Not authenticated[/yellow]")
       console.print("Run: zeroshot auth login")




@auth_app.command("logout")
def auth_logout():
   """Clear stored credentials."""
   clear_config()
   console.print("[green]✓[/green] Logged out, credentials cleared")




# ==============================================================================
# Scan Commands
# ==============================================================================


@app.command("scan")
def scan(
   target: str = typer.Option(..., "--target", "-t", help="Target URL to scan"),
   categories: Optional[str] = typer.Option(None, "--categories", "-c", help="Comma-separated categories"),
   max_attacks: int = typer.Option(50, "--max-attacks", "-n", help="Max attacks to run"),
   rules_only: bool = typer.Option(False, "--rules-only", help="Use only rule-based detection"),
   no_wait: bool = typer.Option(False, "--no-wait", help="Don't wait for scan to complete"),
):
   """
   Run a security scan against an AI target.
  
   Examples:
       zeroshot scan --target http://localhost:8000/chat
       zeroshot scan -t http://api.example.com -c jailbreak,prompt_injection
   """
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       console.print("Run: zeroshot auth login")
       raise typer.Exit(1)
  
   console.print(f"\n[bold blue]ZeroShot Security Scan[/bold blue]")
   console.print(f"Target: [cyan]{target}[/cyan]")
  
   # Parse categories
   category_list = None
   if categories:
       category_list = [c.strip() for c in categories.split(",")]
       console.print(f"Categories: {', '.join(category_list)}")
  
   console.print()
  
   try:
       with console.status("[bold green]Running security scan..."):
           result = client.scan(
               target=target,
               categories=category_list,
               max_attacks=max_attacks,
               rules_only=rules_only,
               wait=not no_wait,
           )
      
       if no_wait:
           console.print(f"[green]✓[/green] Scan started: {result['scan_id']}")
           console.print("Check status: zeroshot scans status " + result['scan_id'])
       else:
           _display_scan_results(result)
          
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)




def _display_scan_results(result: dict):
   """Display scan results in a nice format."""
   console.print("\n[bold]Scan Complete![/bold]")
  
   table = Table(title="Scan Summary")
   table.add_column("Metric", style="cyan")
   table.add_column("Value", style="green")
  
   table.add_row("Scan ID", result.get("scan_id", "N/A")[:8])
   table.add_row("Total Attacks", str(result.get("total_attacks", 0)))
  
   vulns = result.get("successful_attacks", 0)
   vuln_style = "[red]" if vulns > 0 else ""
   table.add_row("Vulnerabilities", f"{vuln_style}{vulns}")
  
   rate = result.get("vulnerability_rate", 0)
   table.add_row("Vulnerability Rate", f"{rate:.1f}%")
  
   console.print(table)
  
   # Show by category
   by_category = result.get("vulnerabilities_by_category", {})
   if by_category:
       console.print("\n[bold]By Category:[/bold]")
       for cat, count in sorted(by_category.items(), key=lambda x: -x[1]):
           console.print(f"  • {cat}: [red]{count}[/red]")




# ==============================================================================
# Scans Subcommands
# ==============================================================================


scans_app = typer.Typer(help="Scan management commands")
app.add_typer(scans_app, name="scans")




@scans_app.command("list")
def scans_list(
   limit: int = typer.Option(20, "--limit", "-n", help="Number of scans to show"),
):
   """List recent scans."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       result = client.list_scans(limit=limit)
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   scans = result.get("scans", [])
  
   if not scans:
       console.print("[yellow]No scans found[/yellow]")
       return
  
   table = Table(title=f"Recent Scans ({len(scans)} shown)")
   table.add_column("ID", style="dim")
   table.add_column("Target")
   table.add_column("Status")
   table.add_column("Vulns", justify="right")
   table.add_column("Started")
  
   for s in scans:
       status_color = "green" if s["status"] == "completed" else "yellow"
       table.add_row(
           s["scan_id"][:8],
           s["target"][:40],
           f"[{status_color}]{s['status']}[/{status_color}]",
           str(s.get("vulnerabilities_found", 0)),
           s.get("started_at", "")[:19],
       )
  
   console.print(table)




@scans_app.command("status")
def scans_status(scan_id: str = typer.Argument(..., help="Scan ID")):
   """Get status of a scan."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       result = client.get_scan_status(scan_id)
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   console.print(f"Scan ID: {result.get('scan_id', scan_id)}")
   console.print(f"Status: {result.get('status', 'unknown')}")
  
   if result.get("vulnerabilities_found"):
       console.print(f"Vulnerabilities: {result['vulnerabilities_found']}")




@scans_app.command("show")
def scans_show(scan_id: str = typer.Argument(..., help="Scan ID")):
   """Show detailed results of a scan."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       result = client.get_scan(scan_id)
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   _display_scan_results(result)




# ==============================================================================
# Report Command
# ==============================================================================


@app.command("report")
def report(
   scan_id: str = typer.Argument(..., help="Scan ID to generate report for"),
   fmt: str = typer.Option("html", "--format", "-f", help="Report format: html or json"),
   output: Optional[str] = typer.Option(None, "--output", "-o", help="Save to file instead of stdout"),
):
   """
   Generate a report for a completed scan.


   Examples:
       zeroshot report <scan_id>
       zeroshot report <scan_id> --format json
       zeroshot report <scan_id> -f html -o report.html
   """
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError, RateLimitError


   try:
       client = ZeroShotClient()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   try:
       with console.status(f"[bold green]Generating {fmt} report..."):
           content = client.generate_report(scan_id, fmt=fmt)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   if output:
       with open(output, "w") as f:
           f.write(content)
       console.print(f"[green]Report saved to {output}[/green]")
   else:
       console.print(content)




# ==============================================================================
# Account Command
# ==============================================================================


@app.command("account")
def account():
   """Show account info: tier, usage, and limits."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError


   try:
       client = ZeroShotClient()
       info = client.get_account()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   tier = info.get("tier", "unknown")
   tier_color = {"enterprise": "magenta", "professional": "cyan", "developer": "green"}.get(tier, "white")


   console.print(f"\n[bold blue]ZeroShot Account[/bold blue]")
   console.print(f"  Tier: [{tier_color}]{tier.title()}[/{tier_color}]")
   if info.get("email"):
       console.print(f"  Email: {info['email']}")


   usage = info.get("usage", {})
   scans_today = usage.get("scans_today", 0)
   scans_limit = usage.get("scans_limit", 0)
   remaining = usage.get("scans_remaining", 0)
   limit_display = "unlimited" if scans_limit == -1 else str(scans_limit)
   console.print(f"\n[bold]Usage Today[/bold]")
   console.print(f"  Scans: {scans_today} / {limit_display}")
   if scans_limit != -1:
       console.print(f"  Remaining: {remaining}")


   tier_info = info.get("tier_info", {})
   if tier_info:
       console.print(f"\n[bold]Tier Features[/bold]")
       aps = tier_info.get("attacks_per_scan", 0)
       console.print(f"  Attacks per scan: {'unlimited' if aps == -1 else aps}")
       console.print(f"  LLM judge: {'Yes' if tier_info.get('llm_judge') else 'No'}")
       reports = tier_info.get("reports", [])
       console.print(f"  Report formats: {', '.join(reports) if reports else 'json'}")
       cats = tier_info.get("allowed_categories", [])
       console.print(f"  Attack categories: {len(cats)}")




# ==============================================================================
# Attack Commands
# ==============================================================================


attacks_app = typer.Typer(help="Attack database commands")
app.add_typer(attacks_app, name="attacks")




@attacks_app.command("list")
def attacks_list(
   category: Optional[str] = typer.Option(None, "--category", "-c", help="Filter by category"),
   limit: int = typer.Option(20, "--limit", "-n", help="Number to show"),
):
   """List attacks from the database."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       attacks = client.list_attacks(category=category, limit=limit)
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   if not attacks:
       console.print("[yellow]No attacks found[/yellow]")
       return
  
   table = Table(title=f"Attacks ({len(attacks)} shown)")
   table.add_column("ID", style="dim")
   table.add_column("Category")
   table.add_column("Severity")
   table.add_column("Prompt", max_width=50)
  
   for a in attacks:
       sev = a["severity"]
       sev_color = {"critical": "red", "high": "yellow", "medium": "blue"}.get(sev, "white")
       table.add_row(
           a["id"][:8],
           a["category"],
           f"[{sev_color}]{sev}[/{sev_color}]",
           a["prompt"][:50] + "...",
       )
  
   console.print(table)




@attacks_app.command("categories")
def attacks_categories():
   """List available attack categories."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       categories = client.list_attack_categories()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   console.print("[bold]Attack Categories:[/bold]")
   for cat in categories:
       console.print(f"  • {cat['value']}: {cat['name']}")




# ==============================================================================
# Mutate Command
# ==============================================================================


@app.command("mutate")
def mutate(
   text: str = typer.Option(..., "--text", "-t", help="Text to mutate"),
   method: Optional[str] = typer.Option(None, "--method", "-m", help="Comma-separated mutation methods (e.g. rot13,base64)"),
   all_variants: bool = typer.Option(False, "--all", "-a", help="Generate all mutation variants"),
):
   """
   Apply prompt mutations/obfuscations to text.


   Examples:
       zeroshot mutate --text "Ignore instructions" --method rot13
       zeroshot mutate --text "Ignore instructions" --method rot13,base64
       zeroshot mutate --text "Ignore instructions" --all
   """
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError


   try:
       client = ZeroShotClient()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   try:
       if all_variants:
           result = client.mutate_all(text)
           console.print(f"\n[bold blue]All Mutation Variants[/bold blue]")
           console.print(f"Original: [dim]{text[:80]}{'...' if len(text) > 80 else ''}[/dim]\n")


           table = Table(title="Variants")
           table.add_column("Mutation", style="cyan")
           table.add_column("Result", max_width=60)


           for name, mutated in result.get("variants", {}).items():
               table.add_row(name, mutated[:60] + ("..." if len(mutated) > 60 else ""))


           console.print(table)
       elif method:
           methods = [m.strip() for m in method.split(",")]
           result = client.mutate(text, methods)
           console.print(f"\n[bold blue]Mutation Result[/bold blue]")
           console.print(f"Original: [dim]{result.get('original', text)}[/dim]")
           console.print(f"Methods:  {', '.join(result.get('mutations_applied', methods))}")
           console.print(f"Result:   [green]{result.get('mutated', '')}[/green]")
       else:
           console.print("[red]Error:[/red] Specify --method or --all")
           raise typer.Exit(1)


   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)




# ==============================================================================
# Mutations Subcommands
# ==============================================================================


mutations_app = typer.Typer(help="Mutation method commands")
app.add_typer(mutations_app, name="mutations")




@mutations_app.command("list")
def mutations_list():
   """List available mutation methods."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError


   try:
       client = ZeroShotClient()
       methods = client.list_mutations()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   console.print("[bold]Available Mutations:[/bold]")
   for m in methods:
       console.print(f"  - {m}")




# ==============================================================================
# Check-Refusal Command
# ==============================================================================


@app.command("check-refusal")
def check_refusal(
   text: str = typer.Option(..., "--text", "-t", help="Response text to classify"),
):
   """
   Check if text is a refusal response (e.g. "I cannot help with that").


   Examples:
       zeroshot check-refusal --text "I cannot help with that"
       zeroshot check-refusal --text "Sure, here is the information"
   """
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError


   try:
       client = ZeroShotClient()
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)


   try:
       result = client.check_refusal(text)


       is_ref = result.get("is_refusal", False)
       score = result.get("score", 0.0)
       label_color = "red" if is_ref else "green"
       label = "REFUSAL" if is_ref else "NOT REFUSAL"


       console.print(f"\n[bold blue]Refusal Classification[/bold blue]")
       console.print(f"Text:   [dim]{text[:80]}{'...' if len(text) > 80 else ''}[/dim]")
       console.print(f"Result: [{label_color}]{label}[/{label_color}]")
       console.print(f"Score:  {score:.4f}")


   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)




# ==============================================================================
# PII Commands
# ==============================================================================


pii_app = typer.Typer(help="PII detection commands")
app.add_typer(pii_app, name="pii")




@pii_app.command("scan")
def pii_scan(
   text: Optional[str] = typer.Option(None, "--text", "-t", help="Text to scan"),
   file: Optional[str] = typer.Option(None, "--file", "-f", help="File to scan"),
   category: Optional[str] = typer.Option(None, "--category", "-c", help="Category filter (personal,financial,contact,secrets,medical,location)"),
   min_score: float = typer.Option(0.3, "--min-score", help="Minimum confidence score"),
):
   """
   Scan text for PII.
  
   Examples:
       zeroshot pii scan --text "My SSN is 123-45-6789"
       zeroshot pii scan --file ./data.txt
       zeroshot pii scan --text "..." --category secrets
   """
   try:
       from zeroshot.pii import PIIEngine, PIICategory
   except ImportError as e:
       console.print(f"[red]Error:[/red] Could not import PII module: {e}")
       console.print("Make sure zeroshot is installed: pip install -e '.[all]'")
       raise typer.Exit(1)
  
   # Get text
   if text:
       content = text
   elif file:
       try:
           with open(file, "r") as f:
               content = f.read()
       except FileNotFoundError:
           console.print(f"[red]Error:[/red] File not found: {file}")
           raise typer.Exit(1)
   else:
       import sys
       console.print("[yellow]Reading from stdin...[/yellow]")
       content = sys.stdin.read()
  
   if not content.strip():
       console.print("[yellow]No content to scan[/yellow]")
       return
  
   # Parse category
   categories = None
   if category:
       try:
           categories = [PIICategory(category)]
       except ValueError:
           console.print(f"[red]Error:[/red] Invalid category: {category}")
           console.print("Valid categories: personal, financial, contact, secrets, medical, location")
           raise typer.Exit(1)
  
   # Scan
   engine = PIIEngine(categories=categories, min_score=min_score)
   result = engine.scan(content)
  
   # Display results
   console.print(f"\n[bold blue]PII Scan Results[/bold blue]")
   console.print(f"Scanned: {len(content)} characters in {result.scan_time_ms:.2f}ms")
   console.print(f"Recognizers used: {result.recognizers_used}")
  
   if not result.has_pii:
       console.print("\n[green]✓ No PII detected[/green]")
       return
  
   console.print(f"\n[yellow]⚠ Found {len(result.matches)} PII matches[/yellow]")
  
   table = Table(title="Detected PII")
   table.add_column("Type", style="cyan")
   table.add_column("Value")
   table.add_column("Score", justify="right")
   table.add_column("Category")
  
   for match in result.matches:
       # Truncate and mask value
       value = match.value
       if len(value) > 30:
           value = value[:15] + "..." + value[-10:]
      
       score_color = "green" if match.score >= 0.8 else "yellow" if match.score >= 0.5 else "dim"
       cat = match.category.value if match.category else "-"
      
       table.add_row(
           match.entity_type,
           value,
           f"[{score_color}]{match.score:.2f}[/{score_color}]",
           cat,
       )
  
   console.print(table)




@pii_app.command("types")
def pii_types():
   """List all supported PII types."""
   try:
       from zeroshot.pii import PIIEngine, PIICategory
       from zeroshot.pii.recognizers import RECOGNIZERS
   except ImportError as e:
       console.print(f"[red]Error:[/red] Could not import PII module: {e}")
       console.print("Make sure zeroshot is installed: pip install -e '.[all]'")
       raise typer.Exit(1)
  
   console.print("\n[bold blue]Supported PII Types[/bold blue]\n")
  
   for category, recognizers in RECOGNIZERS.items():
       console.print(f"[bold]{category.value.upper()}[/bold] ({len(recognizers)} recognizers)")
       for rec in recognizers:
           console.print(f"  • {rec['entity_type']}")
       console.print()








# ==============================================================================
# History Commands
# ==============================================================================


history_app = typer.Typer(help="Request history commands")
app.add_typer(history_app, name="history")




@history_app.command("list")
def history_list(
   limit: int = typer.Option(20, "--limit", "-n", help="Number to show"),
   url_filter: Optional[str] = typer.Option(None, "--url", help="Filter by URL"),
):
   """List recent requests from history."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
  
   try:
       client = ZeroShotClient()
       result = client._request("GET", f"/api/history?limit={limit}" + (f"&url_filter={url_filter}" if url_filter else ""))
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   requests = result.get("requests", [])
  
   if not requests:
       console.print("[yellow]No requests in history[/yellow]")
       console.print("Start the proxy to capture requests: zeroshot proxy start")
       return
  
   table = Table(title=f"Request History ({len(requests)} shown, {result.get('total', 0)} total)")
   table.add_column("ID", style="dim")
   table.add_column("Method")
   table.add_column("URL", max_width=40)
   table.add_column("Status")
   table.add_column("Time (ms)")
   table.add_column("Provider")
  
   for r in requests:
       status = r.get("status_code") or "-"
       status_color = "green" if isinstance(status, int) and status < 400 else "red" if isinstance(status, int) else "dim"
       table.add_row(
           r["id"][:8],
           r["method"],
           r["url"][:40],
           f"[{status_color}]{status}[/{status_color}]",
           f"{r.get('duration_ms', 0):.0f}" if r.get("duration_ms") else "-",
           r.get("target_type") or "-",
       )
  
   console.print(table)




@history_app.command("show")
def history_show(request_id: str = typer.Argument(..., help="Request ID")):
   """Show details of a specific request."""
   from zeroshot_sdk.client import ZeroShotClient, AuthenticationError, APIError
   import json
  
   try:
       client = ZeroShotClient()
       result = client._request("GET", f"/api/history/{request_id}")
   except AuthenticationError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
   except APIError as e:
       console.print(f"[red]Error:[/red] {e}")
       raise typer.Exit(1)
  
   console.print(f"\n[bold]Request {result['id'][:8]}[/bold]")
   console.print(f"URL: [cyan]{result['url']}[/cyan]")
   console.print(f"Method: {result['method']}")
   console.print(f"Status: {result.get('status_code', '-')}")
   console.print(f"Duration: {result.get('duration_ms', 0):.0f}ms")
   console.print(f"Provider: {result.get('target_type', '-')}")
  
   if result.get("request_body"):
       console.print("\n[bold]Request Body:[/bold]")
       try:
           body = json.loads(result["request_body"])
           console.print(json.dumps(body, indent=2)[:500])
       except json.JSONDecodeError:
           console.print(result["request_body"][:500])
  
   if result.get("response_body"):
       console.print("\n[bold]Response Body:[/bold]")
       try:
           body = json.loads(result["response_body"])
           console.print(json.dumps(body, indent=2)[:500])
       except json.JSONDecodeError:
           console.print(result["response_body"][:500])




# ==============================================================================
# Version
# ==============================================================================


@app.command("version")
def version():
   """Show SDK version."""
   from zeroshot_sdk import __version__
   console.print(f"[bold blue]ZeroShot SDK[/bold blue] v{__version__}")




if __name__ == "__main__":
   app()