"""Core HIPAA Agent client — wraps every MCP tool as a typed Python method.

All methods call ``POST /api/mcp`` using JSON-RPC 2.0 (``method: "tools/call"``).
Retry logic: 3x exponential backoff (1 s, 2 s, 4 s) on 429 / 500 / 502 / 503.
Rate-limit handling: respects the ``Retry-After`` header when present.
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Dict, List, Optional, Sequence, TypedDict, Union

import httpx


# ---------------------------------------------------------------------------
# Response type definitions
# ---------------------------------------------------------------------------

class ScanResult(TypedDict, total=False):
    scan_id: str
    npi: str
    domain: str
    status: str
    findings_count: int
    grade: str
    score: float
    created_at: str


class ScanStatus(TypedDict, total=False):
    scan_id: str
    status: str
    progress: float
    findings_count: int
    started_at: str
    completed_at: Optional[str]


class ComplianceScore(TypedDict, total=False):
    npi: str
    domain: str
    grade: str
    score: float
    categories: Dict[str, Any]
    last_scanned: str


class Report(TypedDict, total=False):
    report_id: str
    npi: str
    pdf_url: str
    generated_at: str


class AuditLogEntry(TypedDict, total=False):
    entry_id: str
    action: str
    actor: str
    timestamp: str
    details: Dict[str, Any]
    hash: str


class EvidencePackage(TypedDict, total=False):
    package_id: str
    npi: str
    components: List[Dict[str, Any]]
    generated_at: str


class BAADocument(TypedDict, total=False):
    baa_id: str
    vendor_name: str
    document_url: str
    status: str
    generated_at: str


class Policy(TypedDict, total=False):
    policy_id: str
    title: str
    category: str
    content: str
    version: str
    last_updated: str


class SRAResult(TypedDict, total=False):
    sra_id: str
    npi: str
    risk_level: str
    findings: List[Dict[str, Any]]
    generated_at: str


class BreachRecord(TypedDict, total=False):
    breach_id: str
    entity_name: str
    state: str
    individuals_affected: int
    breach_type: str
    severity: int
    reported_date: str
    summary: str


class BreachScore(TypedDict, total=False):
    npi: str
    breach_probability: float
    risk_factors: List[str]
    historical_breaches: int


class InternalScanResult(TypedDict, total=False):
    scan_id: str
    agent_id: str
    status: str
    started_at: str


class InternalScanStatus(TypedDict, total=False):
    scan_id: str
    agent_id: str
    status: str
    progress: float
    findings_count: int
    completed_at: Optional[str]


class InternalFinding(TypedDict, total=False):
    finding_id: str
    category: str
    severity: str
    title: str
    description: str
    remediation: str


class PracticeLookup(TypedDict, total=False):
    npi: str
    name: str
    domain: str
    grade: str
    score: float
    findings_count: int
    last_scanned: str


class OutreachStatus(TypedDict, total=False):
    npi: str
    drip_stage: int
    last_sent: str
    opens: int
    clicks: int
    status: str


class PracticeSummary(TypedDict, total=False):
    npi: str
    name: str
    address: str
    phone: str
    taxonomy: str
    grade: str
    score: float
    findings: List[Dict[str, Any]]


class TrainingStatus(TypedDict, total=False):
    npi: str
    modules_completed: int
    modules_total: int
    completion_percentage: float
    last_activity: str


class VendorBAA(TypedDict, total=False):
    vendor_name: str
    baa_status: str
    signed_date: Optional[str]
    expiry_date: Optional[str]


class Incident(TypedDict, total=False):
    incident_id: str
    npi: str
    title: str
    severity: str
    status: str
    reported_at: str
    details: Dict[str, Any]


class Control(TypedDict, total=False):
    control_id: str
    title: str
    category: str
    status: str
    implementation_notes: str


class WorkflowValidation(TypedDict, total=False):
    valid: bool
    errors: List[str]
    warnings: List[str]
    workflow_type: str


class BreachProbability(TypedDict, total=False):
    npi: str
    probability: float
    risk_level: str
    contributing_factors: List[Dict[str, Any]]


class WebhookSubscription(TypedDict, total=False):
    webhook_id: str
    url: str
    events: List[str]
    status: str
    created_at: str


class ComplianceState(TypedDict, total=False):
    npi: str
    overall_status: str
    categories: Dict[str, Any]
    last_evaluated: str


class ComplianceDelta(TypedDict, total=False):
    npi: str
    changes: List[Dict[str, Any]]
    period_start: str
    period_end: str
    score_change: float


class VendorCheck(TypedDict, total=False):
    vendor_name: str
    domain: str
    hipaa_compliant: bool
    has_baa: bool
    risk_level: str
    details: Dict[str, Any]


class BatchScanResult(TypedDict, total=False):
    batch_id: str
    total: int
    queued: int
    status: str
    created_at: str


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------

class HIPAAAgentError(Exception):
    """Base exception for all SDK errors."""

    def __init__(self, message: str, code: Optional[int] = None, data: Any = None) -> None:
        super().__init__(message)
        self.code = code
        self.data = data


class AuthenticationError(HIPAAAgentError):
    """Raised on 401 / 403 responses."""


class RateLimitError(HIPAAAgentError):
    """Raised when retries are exhausted on 429."""


class ServerError(HIPAAAgentError):
    """Raised on persistent 5xx responses after retries."""


class ValidationError(HIPAAAgentError):
    """Raised when the API returns a JSON-RPC validation error."""


# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------

_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503})
_MAX_RETRIES = 3
_BACKOFF_SECONDS = (1.0, 2.0, 4.0)
_DEFAULT_TIMEOUT = 120.0  # seconds — scans can take a while


class HIPAAAgent:
    """Synchronous Python client for the HIPAA Agent MCP API.

    Parameters
    ----------
    api_key:
        Your ``x-api-key`` credential.
    base_url:
        Root URL of the HIPAA Agent instance.  Defaults to production.
    timeout:
        HTTP timeout in seconds.  Defaults to 120.
    max_retries:
        Maximum retry attempts for transient failures.  Defaults to 3.
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://hipaaagent.ai",
        timeout: float = _DEFAULT_TIMEOUT,
        max_retries: int = _MAX_RETRIES,
    ) -> None:
        if not api_key:
            raise AuthenticationError("api_key must be a non-empty string")
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._max_retries = max_retries
        self._client = httpx.Client(
            base_url=self._base_url,
            headers={
                "x-api-key": self._api_key,
                "Content-Type": "application/json",
                "User-Agent": "hipaa-agent-python/1.0.0",
            },
            timeout=httpx.Timeout(self._timeout),
        )

    # ------------------------------------------------------------------
    # Transport
    # ------------------------------------------------------------------

    def _call(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """Execute a JSON-RPC ``tools/call`` request with retry logic.

        Returns the parsed ``result`` value from the JSON-RPC response, or
        raises a typed exception on failure.
        """
        payload: Dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": str(uuid.uuid4()),
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments,
            },
        }

        last_exc: Optional[Exception] = None

        for attempt in range(self._max_retries):
            try:
                response = self._client.post("/api/mcp", json=payload)
            except httpx.TransportError as exc:
                last_exc = exc
                if attempt < self._max_retries - 1:
                    time.sleep(_BACKOFF_SECONDS[attempt])
                    continue
                raise ServerError(
                    f"Transport error after {self._max_retries} attempts: {exc}"
                ) from exc

            # --- Auth errors (never retry) ---
            if response.status_code in (401, 403):
                raise AuthenticationError(
                    f"Authentication failed ({response.status_code}): {response.text}",
                    code=response.status_code,
                )

            # --- Retryable status codes ---
            if response.status_code in _RETRYABLE_STATUS_CODES:
                last_exc = ServerError(
                    f"HTTP {response.status_code}: {response.text}",
                    code=response.status_code,
                )
                if attempt < self._max_retries - 1:
                    retry_after = response.headers.get("Retry-After")
                    if retry_after is not None:
                        try:
                            wait = float(retry_after)
                        except ValueError:
                            wait = _BACKOFF_SECONDS[attempt]
                    else:
                        wait = _BACKOFF_SECONDS[attempt]
                    time.sleep(wait)
                    continue

                # Exhausted retries
                if response.status_code == 429:
                    raise RateLimitError(
                        f"Rate limited after {self._max_retries} retries",
                        code=429,
                    )
                raise last_exc

            # --- Other non-2xx ---
            if response.status_code >= 400:
                raise HIPAAAgentError(
                    f"HTTP {response.status_code}: {response.text}",
                    code=response.status_code,
                )

            # --- Success path ---
            body = response.json()

            if "error" in body:
                err = body["error"]
                code = err.get("code", -1)
                message = err.get("message", "Unknown JSON-RPC error")
                data = err.get("data")
                if code == -32602:
                    raise ValidationError(message, code=code, data=data)
                raise HIPAAAgentError(message, code=code, data=data)

            result = body.get("result", {})

            # MCP tools/call returns content array — unwrap first text item
            if isinstance(result, dict) and "content" in result:
                content_items: List[Any] = result["content"]
                for item in content_items:
                    if isinstance(item, dict) and item.get("type") == "text":
                        import json as _json

                        try:
                            return _json.loads(item["text"])
                        except (ValueError, KeyError):
                            return item.get("text", result)
                return result

            return result

        # Should not reach here, but just in case
        raise ServerError(
            f"Request failed after {self._max_retries} attempts",
        ) from last_exc

    # ------------------------------------------------------------------
    # Scanning
    # ------------------------------------------------------------------

    def scan(
        self,
        npi: str,
        domain: Optional[str] = None,
        *,
        callback_url: Optional[str] = None,
    ) -> ScanResult:
        """Launch a comprehensive HIPAA compliance scan for a practice.

        Parameters
        ----------
        npi:
            10-digit National Provider Identifier.
        domain:
            Practice website domain. Auto-discovered from NPI if omitted.
        callback_url:
            Optional webhook URL to notify on completion.
        """
        args: Dict[str, Any] = {"npi": npi}
        if domain is not None:
            args["domain"] = domain
        if callback_url is not None:
            args["callback_url"] = callback_url
        return self._call("hipaa_scan", args)

    def get_scan_status(self, scan_id: str) -> ScanStatus:
        """Poll the status of a running or completed scan."""
        return self._call("hipaa_scan_status", {"scan_id": scan_id})

    def batch_scan(
        self,
        targets: Sequence[Dict[str, str]],
        *,
        callback_url: Optional[str] = None,
    ) -> BatchScanResult:
        """Queue a batch of practices for scanning.

        Parameters
        ----------
        targets:
            List of dicts, each with ``npi`` and optionally ``domain``.
        callback_url:
            Webhook URL notified when the entire batch completes.
        """
        args: Dict[str, Any] = {"targets": list(targets)}
        if callback_url is not None:
            args["callback_url"] = callback_url
        return self._call("hipaa_batch_scan", args)

    # ------------------------------------------------------------------
    # Compliance scoring & grading
    # ------------------------------------------------------------------

    def get_compliance_score(self, npi: str) -> ComplianceScore:
        """Retrieve the HIPAA Agent Compliance Score for a practice."""
        return self._call("hipaa_grade", {"npi": npi})

    def get_compliance_state(self, npi: str) -> ComplianceState:
        """Get the full current compliance state for a practice."""
        return self._call("hipaa_compliance_state", {"npi": npi})

    def get_compliance_delta(
        self,
        npi: str,
        *,
        since: Optional[str] = None,
        until: Optional[str] = None,
    ) -> ComplianceDelta:
        """Get compliance changes over a time period.

        Parameters
        ----------
        since:
            ISO-8601 start date. Defaults to 30 days ago.
        until:
            ISO-8601 end date. Defaults to now.
        """
        args: Dict[str, Any] = {"npi": npi}
        if since is not None:
            args["since"] = since
        if until is not None:
            args["until"] = until
        return self._call("hipaa_compliance_delta", args)

    # ------------------------------------------------------------------
    # Reports & evidence
    # ------------------------------------------------------------------

    def get_report(self, npi: str) -> Report:
        """Generate or retrieve the PDF Risk Posture Report."""
        return self._call("hipaa_report", {"npi": npi})

    def get_audit_log(
        self,
        npi: str,
        *,
        limit: int = 100,
        offset: int = 0,
    ) -> List[AuditLogEntry]:
        """Retrieve the SHA-256 hash-chain audit trail for a practice."""
        return self._call(
            "hipaa_audit_log",
            {"npi": npi, "limit": limit, "offset": offset},
        )

    def get_evidence_package(self, npi: str) -> EvidencePackage:
        """Generate a 10-component evidence package for a practice."""
        return self._call("hipaa_evidence_package", {"npi": npi})

    # ------------------------------------------------------------------
    # BAA management
    # ------------------------------------------------------------------

    def generate_baa(
        self,
        vendor_name: str,
        *,
        npi: Optional[str] = None,
        vendor_domain: Optional[str] = None,
    ) -> BAADocument:
        """Generate a Business Associate Agreement document."""
        args: Dict[str, Any] = {"vendor_name": vendor_name}
        if npi is not None:
            args["npi"] = npi
        if vendor_domain is not None:
            args["vendor_domain"] = vendor_domain
        return self._call("hipaa_generate_baa", args)

    def get_vendor_baa_list(self, npi: str) -> List[VendorBAA]:
        """List all vendor BAA statuses for a practice."""
        return self._call("hipaa_vendor_baa_list", {"npi": npi})

    def check_vendor(
        self,
        vendor_name: str,
        *,
        vendor_domain: Optional[str] = None,
    ) -> VendorCheck:
        """Check a vendor's HIPAA compliance posture."""
        args: Dict[str, Any] = {"vendor_name": vendor_name}
        if vendor_domain is not None:
            args["vendor_domain"] = vendor_domain
        return self._call("hipaa_check_vendor", args)

    # ------------------------------------------------------------------
    # Policies & training
    # ------------------------------------------------------------------

    def get_policies(
        self,
        *,
        category: Optional[str] = None,
    ) -> List[Policy]:
        """Retrieve HIPAA policy documents, optionally filtered by category."""
        args: Dict[str, Any] = {}
        if category is not None:
            args["category"] = category
        return self._call("hipaa_policies", args)

    def get_training_status(self, npi: str) -> TrainingStatus:
        """Get staff training module completion status."""
        return self._call("hipaa_training_status", {"npi": npi})

    # ------------------------------------------------------------------
    # Security Risk Assessment
    # ------------------------------------------------------------------

    def generate_sra(self, npi: str) -> SRAResult:
        """Generate a Security Risk Assessment for a practice."""
        return self._call("hipaa_sra", {"npi": npi})

    # ------------------------------------------------------------------
    # Breach intelligence
    # ------------------------------------------------------------------

    def get_breach(
        self,
        *,
        entity_name: Optional[str] = None,
        breach_id: Optional[str] = None,
        state: Optional[str] = None,
        limit: int = 10,
    ) -> Union[BreachRecord, List[BreachRecord]]:
        """Look up HHS breach records by entity, ID, or state."""
        args: Dict[str, Any] = {"limit": limit}
        if entity_name is not None:
            args["entity_name"] = entity_name
        if breach_id is not None:
            args["breach_id"] = breach_id
        if state is not None:
            args["state"] = state
        return self._call("hipaa_breach_check", args)

    def get_breach_score(self, npi: str) -> BreachScore:
        """Get a breach-risk score for a practice based on historical data."""
        return self._call("hipaa_breach_score", {"npi": npi})

    def get_breach_probability(
        self,
        npi: str,
        *,
        time_horizon_days: int = 365,
    ) -> BreachProbability:
        """Calculate breach probability for a practice.

        Parameters
        ----------
        npi:
            10-digit NPI.
        time_horizon_days:
            Prediction window in days. Defaults to 365.
        """
        return self._call(
            "hipaa_breach_probability",
            {"npi": npi, "time_horizon_days": time_horizon_days},
        )

    # ------------------------------------------------------------------
    # Internal (network agent) scanning
    # ------------------------------------------------------------------

    def trigger_internal_scan(
        self,
        npi: str,
        agent_id: str,
    ) -> InternalScanResult:
        """Trigger an internal network scan via a deployed agent."""
        return self._call(
            "hipaa_internal_scan",
            {"npi": npi, "agent_id": agent_id},
        )

    def get_internal_scan_status(self, scan_id: str) -> InternalScanStatus:
        """Poll status of an internal network scan."""
        return self._call("hipaa_internal_scan_status", {"scan_id": scan_id})

    def get_internal_findings(
        self,
        scan_id: str,
        *,
        severity: Optional[str] = None,
    ) -> List[InternalFinding]:
        """Retrieve findings from a completed internal scan."""
        args: Dict[str, Any] = {"scan_id": scan_id}
        if severity is not None:
            args["severity"] = severity
        return self._call("hipaa_internal_findings", args)

    # ------------------------------------------------------------------
    # Practice lookup & outreach
    # ------------------------------------------------------------------

    def lookup_practice(self, npi: str) -> PracticeLookup:
        """Look up a practice by NPI and return scan summary."""
        return self._call("hipaa_lookup_practice", {"npi": npi})

    def get_outreach_status(self, npi: str) -> OutreachStatus:
        """Get the outreach / drip-campaign status for a practice."""
        return self._call("hipaa_outreach_status", {"npi": npi})

    def get_practice_summary(self, npi: str) -> PracticeSummary:
        """Get a comprehensive summary including findings for a practice."""
        return self._call("hipaa_practice_summary", {"npi": npi})

    # ------------------------------------------------------------------
    # Incidents & controls
    # ------------------------------------------------------------------

    def log_incident(
        self,
        npi: str,
        title: str,
        *,
        severity: str = "medium",
        description: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
    ) -> Incident:
        """Log a security or privacy incident."""
        args: Dict[str, Any] = {
            "npi": npi,
            "title": title,
            "severity": severity,
        }
        if description is not None:
            args["description"] = description
        if details is not None:
            args["details"] = details
        return self._call("hipaa_log_incident", args)

    def get_incidents(
        self,
        npi: str,
        *,
        status: Optional[str] = None,
        limit: int = 50,
    ) -> List[Incident]:
        """List incidents for a practice, optionally filtered by status."""
        args: Dict[str, Any] = {"npi": npi, "limit": limit}
        if status is not None:
            args["status"] = status
        return self._call("hipaa_incidents", args)

    def get_controls(
        self,
        npi: str,
        *,
        category: Optional[str] = None,
    ) -> List[Control]:
        """List implemented HIPAA controls for a practice."""
        args: Dict[str, Any] = {"npi": npi}
        if category is not None:
            args["category"] = category
        return self._call("hipaa_controls", args)

    # ------------------------------------------------------------------
    # Workflow & validation
    # ------------------------------------------------------------------

    def validate_workflow(
        self,
        workflow_type: str,
        *,
        npi: Optional[str] = None,
        steps: Optional[List[Dict[str, Any]]] = None,
    ) -> WorkflowValidation:
        """Validate a HIPAA compliance workflow for correctness.

        Parameters
        ----------
        workflow_type:
            E.g. ``"breach_response"``, ``"access_request"``, ``"sra"``.
        npi:
            Optional practice context for the validation.
        steps:
            Optional list of workflow steps to validate.
        """
        args: Dict[str, Any] = {"workflow_type": workflow_type}
        if npi is not None:
            args["npi"] = npi
        if steps is not None:
            args["steps"] = steps
        return self._call("hipaa_validate_workflow", args)

    # ------------------------------------------------------------------
    # Webhooks
    # ------------------------------------------------------------------

    def subscribe_webhook(
        self,
        url: str,
        events: Sequence[str],
        *,
        secret: Optional[str] = None,
    ) -> WebhookSubscription:
        """Subscribe a URL to receive webhook events.

        Parameters
        ----------
        url:
            HTTPS endpoint that will receive POST notifications.
        events:
            List of event types, e.g. ``["scan.completed", "breach.detected"]``.
        secret:
            Optional shared secret for HMAC signature verification.
        """
        args: Dict[str, Any] = {"url": url, "events": list(events)}
        if secret is not None:
            args["secret"] = secret
        return self._call("hipaa_subscribe_webhook", args)

    def list_webhooks(self) -> List[WebhookSubscription]:
        """List all active webhook subscriptions."""
        return self._call("hipaa_list_webhooks", {})

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Close the underlying HTTP client and release resources."""
        self._client.close()

    def __enter__(self) -> "HIPAAAgent":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def __repr__(self) -> str:
        return f"HIPAAAgent(base_url={self._base_url!r})"
