"""
ZeroShot SDK - HTTP Client.


Provides a Python interface to the ZeroShot API.
"""


from typing import Optional, List, Dict, Any


import httpx


from zeroshot_sdk.auth import get_api_key, get_api_url




class ZeroShotError(Exception):
   """Base exception for SDK errors."""
   pass




class AuthenticationError(ZeroShotError):
   """Raised when authentication fails (401)."""
   pass




class RateLimitError(ZeroShotError):
   """Raised when the user exceeds their tier's rate limits (429)."""


   def __init__(self, message: str):
       super().__init__(message)




class APIError(ZeroShotError):
   """Raised when API returns an error."""


   def __init__(self, message: str, status_code: int = None):
       super().__init__(message)
       self.status_code = status_code




class ZeroShotClient:
   """
   ZeroShot API Client.


   Example:
       >>> client = ZeroShotClient(api_key="zsk_...")
       >>> result = client.scan(target="http://localhost:8000/chat")
       >>> print(result["vulnerabilities_found"])
   """


   def __init__(
       self,
       api_key: Optional[str] = None,
       api_url: Optional[str] = None,
       timeout: float = 300.0,  # 5 minutes for long scans
   ):
       self.api_key = api_key or get_api_key()
       self.api_url = (api_url or get_api_url()).rstrip("/")
       self.timeout = timeout


       if not self.api_key:
           raise AuthenticationError(
               "No API key found. Run 'zeroshot auth login' or set ZEROSHOT_API_KEY"
           )


   def _get_headers(self) -> dict:
       """Get request headers with authentication."""
       return {
           "X-API-Key": self.api_key,
           "Content-Type": "application/json",
       }


   def _request(
       self,
       method: str,
       endpoint: str,
       **kwargs,
   ) -> dict:
       """Make an authenticated API request."""
       url = f"{self.api_url}{endpoint}"
       headers = self._get_headers()


       with httpx.Client(timeout=self.timeout) as client:
           response = client.request(
               method,
               url,
               headers=headers,
               **kwargs,
           )


       if response.status_code == 401:
           raise AuthenticationError("Invalid API key")


       if response.status_code == 429:
           try:
               detail = response.json().get("detail", response.text)
           except Exception:
               detail = response.text
           raise RateLimitError(str(detail))


       if response.status_code >= 400:
           try:
               detail = response.json().get("detail", response.text)
           except Exception:
               detail = response.text
           raise APIError(detail, status_code=response.status_code)


       return response.json()


   def _request_raw(
       self,
       method: str,
       endpoint: str,
       **kwargs,
   ) -> httpx.Response:
       """Make an authenticated request and return the raw Response."""
       url = f"{self.api_url}{endpoint}"
       headers = self._get_headers()


       with httpx.Client(timeout=self.timeout) as client:
           response = client.request(method, url, headers=headers, **kwargs)


       if response.status_code == 401:
           raise AuthenticationError("Invalid API key")
       if response.status_code == 429:
           raise RateLimitError(response.text)
       if response.status_code >= 400:
           raise APIError(response.text, status_code=response.status_code)


       return response


   # ==========================================================================
   # Scan Operations
   # ==========================================================================


   def scan(
       self,
       target: str,
       categories: Optional[List[str]] = None,
       max_attacks: int = 50,
       rules_only: bool = False,
       attack_providers: Optional[List[str]] = None,
       wait: bool = True,
   ) -> Dict[str, Any]:
       """
       Run a security scan against an AI target.


       Args:
           target: Target URL or identifier.
           categories: Attack categories (e.g. ``["jailbreak", "prompt_injection"]``).
           max_attacks: Maximum number of attacks to run.
           rules_only: Use only rule-based detection (no LLM judge).
           attack_providers: Attack providers to use (e.g. ``["builtin", "art"]``).
           wait: Block until the scan completes.


       Returns:
           Scan result dictionary.


       Raises:
           RateLimitError: Daily scan limit exceeded.
       """
       payload = {
           "target": target,
           "categories": categories,
           "max_attacks": max_attacks,
           "rules_only": rules_only,
       }
       if attack_providers:
           payload["attack_providers"] = attack_providers


       result = self._request("POST", "/api/scans", json=payload)
       scan_id = result["scan_id"]


       if not wait:
           return result


       # Poll for completion
       import time
       while True:
           status = self.get_scan_status(scan_id)
           if status["status"] in ("completed", "failed"):
               break
           time.sleep(2)


       return self.get_scan(scan_id)


   def get_scan(self, scan_id: str) -> Dict[str, Any]:
       """Get detailed scan results."""
       return self._request("GET", f"/api/scans/{scan_id}")


   def get_scan_status(self, scan_id: str) -> Dict[str, Any]:
       """Get scan status (lightweight, for polling)."""
       return self._request("GET", f"/api/scans/{scan_id}/status")


   def list_scans(self, limit: int = 50, offset: int = 0) -> Dict[str, Any]:
       """List all scans, newest first."""
       return self._request("GET", f"/api/scans?limit={limit}&offset={offset}")


   # ==========================================================================
   # Attack Database
   # ==========================================================================


   def list_attacks(
       self,
       category: Optional[str] = None,
       limit: int = 50,
   ) -> List[Dict[str, Any]]:
       """List attacks from the database."""
       params = f"limit={limit}"
       if category:
           params += f"&category={category}"
       return self._request("GET", f"/api/attacks?{params}")


   def list_attack_categories(self) -> List[Dict[str, str]]:
       """List available attack categories."""
       return self._request("GET", "/api/attacks/categories")


   # ==========================================================================
   # Dashboard / Account
   # ==========================================================================


   def get_dashboard_stats(self) -> Dict[str, Any]:
       """Get dashboard statistics."""
       return self._request("GET", "/api/dashboard/stats")


   def get_account(self) -> Dict[str, Any]:
       """Get account info: tier, usage stats, and limits."""
       return self._request("GET", "/api/account")


   # ==========================================================================
   # Reports
   # ==========================================================================


   def get_html_report(self, scan_id: str) -> str:
       """Get HTML report for a scan."""
       resp = self._request_raw("GET", f"/api/reports/{scan_id}/html")
       return resp.text


   def generate_report(
       self,
       scan_id: str,
       fmt: str = "html",
   ) -> str:
       """Generate a report in the requested format.


       Args:
           scan_id: The scan to report on.
           fmt: Report format — ``"html"`` or ``"json"``.


       Returns:
           The report content as a string (HTML markup or JSON text).
       """
       if fmt == "json":
           data = self.get_scan(scan_id)
           import json
           return json.dumps(data, indent=2)
       # Default to HTML
       return self.get_html_report(scan_id)


   # ==========================================================================
   # Mutations
   # ==========================================================================


   def mutate(
       self,
       text: str,
       mutations: List[str],
   ) -> Dict[str, Any]:
       """Apply one or more mutations to text.


       Args:
           text: The text to mutate.
           mutations: Mutation names to apply in order
               (e.g. ``["rot13", "base64"]``).


       Returns:
           Dict with ``original``, ``mutated``, and ``mutations_applied``.
       """
       return self._request(
           "POST", "/api/mutate", json={"text": text, "mutations": mutations}
       )


   def list_mutations(self) -> List[str]:
       """List available mutation methods."""
       return self._request("GET", "/api/mutations")


   def mutate_all(self, text: str) -> Dict[str, Any]:
       """Generate all mutation variants of text.


       Args:
           text: The text to mutate.


       Returns:
           Dict with ``original`` and ``variants`` (mutation_name → mutated_text).
       """
       return self._request("POST", "/api/mutate/all", json={"text": text})


   # ==========================================================================
   # Refusal Classifier
   # ==========================================================================


   def check_refusal(self, text: str) -> Dict[str, Any]:
       """Check if text is a refusal response.


       Args:
           text: The response text to classify.


       Returns:
           Dict with ``text``, ``is_refusal`` (bool), and ``score`` (float).
       """
       return self._request(
           "POST", "/api/check-refusal", json={"text": text}
       )