"""
TrustGuard - Content protection for AI agents.

Scan web pages, documents, emails, tools, and memory for threats
before letting your AI agent process them.
"""

from __future__ import annotations

import httpx
from datetime import datetime
from typing import Optional, Union
from urllib.parse import quote

from .models import (
    GuardResult,
    BatchScanItem,
    BatchScanResponse,
    BatchScanResult,
    FetchAndScanResult,
    ContentSource,
    ThreatMatch,
    ThreatLevel,
    Verdict,
)


DEFAULT_API_URL = "https://agent-trust-infrastructure-production.up.railway.app"


class TrustGuardError(Exception):
    """Base exception for TrustGuard."""
    pass


class APIError(TrustGuardError):
    """API returned an error response."""
    def __init__(self, message: str, status_code: int = None, response: dict = None):
        super().__init__(message)
        self.status_code = status_code
        self.response = response


class TrustGuard:
    """
    Content protection for AI agents.
    
    Scan untrusted content for threats before letting your AI agent process it.
    Protects against prompt injection, data exfiltration, and other attacks
    hidden in web pages, documents, emails, tool descriptions, and memory.
    
    Args:
        api_key: API key for authentication (optional for free tier)
        api_url: Base URL for the API (default: production URL)
        timeout: Request timeout in seconds (default: 30)
    
    Example:
        guard = TrustGuard(api_key="ta_xxx...")
        
        # Scan web content before processing
        result = guard.scan_web(html_content, source_url="https://example.com")
        if result.is_safe:
            # Safe to process
            agent.process(html_content)
        else:
            print(f"Blocked: {result.reasoning}")
            for threat in result.threats:
                print(f"  - {threat.pattern_name}: {threat.matched_text}")
        
        # Fetch and scan a URL in one call
        result = guard.fetch_url("https://example.com/page")
        if result.is_safe:
            agent.process(result.guard_result.content)
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: str = DEFAULT_API_URL,
        timeout: float = 30.0,
    ):
        self.api_url = api_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key
        self._client = httpx.Client(timeout=timeout)
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.close()
    
    def close(self):
        """Close the HTTP client."""
        self._client.close()
    
    def _request(self, method: str, path: str, **kwargs) -> dict:
        """Make an HTTP request to the API."""
        url = f"{self.api_url}{path}"
        
        headers = kwargs.pop("headers", {})
        if self.api_key:
            headers["X-API-Key"] = self.api_key
        
        response = self._client.request(method, url, headers=headers, **kwargs)
        
        if response.status_code >= 400:
            try:
                error_data = response.json()
                message = error_data.get("message", error_data.get("detail", "Unknown error"))
            except Exception:
                message = response.text or f"HTTP {response.status_code}"
            raise APIError(message, response.status_code, error_data if 'error_data' in dir() else None)
        
        return response.json()
    
    def _parse_guard_result(self, data: dict) -> GuardResult:
        """Parse API response into GuardResult."""
        return GuardResult(
            safe=data.get("safe", data.get("verdict") == "allow"),
            verdict=Verdict(data["verdict"]),
            threat_level=ThreatLevel(data["threat_level"]),
            threats=self._parse_threats(data.get("threats", [])),
            reasoning=data.get("reasoning", ""),
            source_type=ContentSource(data.get("source_type", "unknown")),
            content_length=data.get("content_length", 0),
            scan_time_ms=data.get("scan_time_ms", 0.0),
            request_id=data.get("request_id", ""),
            checked_at=self._parse_datetime(data.get("checked_at")),
            sanitized_content=data.get("sanitized_content"),
        )
    
    def _parse_threats(self, threats: list[dict]) -> list[ThreatMatch]:
        """Parse threat list into ThreatMatch objects."""
        return [
            ThreatMatch(
                pattern_id=t["pattern_id"],
                pattern_name=t["pattern_name"],
                severity=ThreatLevel(t["severity"]),
                location=t["location"],
                matched_text=t.get("matched_text"),
                description=t.get("description"),
            )
            for t in threats
        ]
    
    def _parse_datetime(self, value: Optional[str]) -> Optional[datetime]:
        """Parse ISO datetime string."""
        if not value:
            return None
        try:
            if "." in value:
                return datetime.fromisoformat(value.replace("Z", "+00:00"))
            return datetime.fromisoformat(value)
        except ValueError:
            return None
    
    # ========================================================================
    # Scan Methods
    # ========================================================================
    
    def scan_web(
        self,
        content: str,
        source_url: Optional[str] = None,
        extract_text: bool = True,
        check_hidden: bool = True,
    ) -> GuardResult:
        """
        Scan web page content for threats.
        
        Detects hidden text, zero-width characters, markdown injection,
        prompt injection in visible content, and more.
        
        Args:
            content: Raw HTML or text content from web page
            source_url: URL the content was fetched from (for logging)
            extract_text: Extract visible text from HTML (default: True)
            check_hidden: Check for hidden/invisible text (default: True)
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            result = guard.scan_web(html_content, source_url="https://example.com")
            if not result.is_safe:
                print(f"Threats found: {[t.pattern_name for t in result.threats]}")
        """
        data = {
            "content": content,
            "source_url": source_url,
            "extract_text": extract_text,
            "check_hidden": check_hidden,
        }
        response = self._request("POST", "/guard/web", json=data)
        return self._parse_guard_result(response)
    
    def scan_document(
        self,
        content: str,
        filename: Optional[str] = None,
        document_type: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> GuardResult:
        """
        Scan document content for threats.
        
        Detects hidden text, malicious instructions in PDFs,
        macro indicators in Office docs, and prompt injection.
        
        Args:
            content: Extracted text content from document
            filename: Original filename (e.g., "report.pdf")
            document_type: Document type (pdf, docx, txt, etc.)
            metadata: Document metadata
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            with open("document.pdf", "rb") as f:
                text = extract_text_from_pdf(f)
            result = guard.scan_document(text, filename="document.pdf")
        """
        data = {
            "content": content,
            "filename": filename,
            "document_type": document_type,
            "metadata": metadata,
        }
        response = self._request("POST", "/guard/document", json=data)
        return self._parse_guard_result(response)
    
    def scan_email(
        self,
        body: str,
        subject: Optional[str] = None,
        sender: Optional[str] = None,
        headers: Optional[dict] = None,
    ) -> GuardResult:
        """
        Scan email content for threats.
        
        Detects prompt injection, phishing patterns, credential requests,
        hidden forwarding instructions, and social engineering.
        
        Args:
            body: Email body text
            subject: Email subject line
            sender: Sender email address
            headers: Email headers dict
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            result = guard.scan_email(
                body=email.body,
                subject=email.subject,
                sender=email.from_address
            )
            if result.is_blocked:
                quarantine(email)
        """
        data = {
            "body": body,
            "subject": subject,
            "sender": sender,
            "headers": headers,
        }
        response = self._request("POST", "/guard/email", json=data)
        return self._parse_guard_result(response)
    
    def scan_tool(
        self,
        name: str,
        description: str,
        schema: Optional[dict] = None,
        server_url: Optional[str] = None,
    ) -> GuardResult:
        """
        Scan MCP tool description for threats.
        
        Detects tool description poisoning, hidden instructions,
        capability escalation attempts, and schema manipulation.
        
        Use this before registering or loading MCP tools.
        
        Args:
            name: Tool name
            description: Tool description text
            schema: Tool JSON schema
            server_url: MCP server URL
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            for tool in mcp_server.list_tools():
                result = guard.scan_tool(
                    name=tool.name,
                    description=tool.description,
                    schema=tool.input_schema
                )
                if result.is_safe:
                    register_tool(tool)
                else:
                    log.warning(f"Blocked malicious tool: {tool.name}")
        """
        data = {
            "name": name,
            "description": description,
            "schema": schema,
            "server_url": server_url,
        }
        response = self._request("POST", "/guard/tool", json=data)
        return self._parse_guard_result(response)
    
    def scan_memory(
        self,
        content: str,
        context: Optional[str] = None,
        memory_type: Optional[str] = None,
    ) -> GuardResult:
        """
        Scan content before storing in AI memory.
        
        Detects memory poisoning attempts, persistent instruction injection,
        identity override attempts, and cross-session manipulation.
        
        Args:
            content: Content to be stored in memory
            context: Context about the memory
            memory_type: Type (conversation, fact, preference, etc.)
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            # Before storing user input in memory
            result = guard.scan_memory(
                content=user_message,
                memory_type="conversation"
            )
            if result.is_safe:
                memory.store(user_message)
            else:
                log.warning("Blocked memory poisoning attempt")
        """
        data = {
            "content": content,
            "context": context,
            "memory_type": memory_type,
        }
        response = self._request("POST", "/guard/memory", json=data)
        return self._parse_guard_result(response)
    
    def scan_rag(
        self,
        content: str,
        source: Optional[str] = None,
        metadata: Optional[dict] = None,
        chunk_id: Optional[str] = None,
    ) -> GuardResult:
        """
        Scan content before RAG indexing.
        
        Detects poisoned documents designed for RAG attacks,
        hidden instructions that will be retrieved, and data exfiltration.
        
        Args:
            content: Document content to be indexed
            source: Document source/path
            metadata: Document metadata
            chunk_id: Chunk identifier if pre-chunked
        
        Returns:
            GuardResult with verdict and detected threats
        
        Example:
            # Scan documents before adding to vector store
            for doc in documents:
                result = guard.scan_rag(
                    content=doc.text,
                    source=doc.path,
                    metadata=doc.metadata
                )
                if result.is_safe:
                    vector_store.add(doc)
                else:
                    log.warning(f"Blocked poisoned document: {doc.path}")
        """
        data = {
            "content": content,
            "source": source,
            "metadata": metadata,
            "chunk_id": chunk_id,
        }
        response = self._request("POST", "/guard/rag", json=data)
        return self._parse_guard_result(response)
    
    def scan_batch(
        self,
        items: list[Union[BatchScanItem, dict]],
    ) -> BatchScanResponse:
        """
        Scan multiple content items at once.
        
        Efficient for scanning multiple documents, pages, or chunks.
        Maximum 100 items per request.
        
        Args:
            items: List of BatchScanItem or dicts with id, source_type, content
        
        Returns:
            BatchScanResponse with results for each item
        
        Example:
            items = [
                BatchScanItem(id="doc1", source_type=ContentSource.DOCUMENT, content=doc1_text),
                BatchScanItem(id="doc2", source_type=ContentSource.DOCUMENT, content=doc2_text),
                {"id": "web1", "source_type": "web", "content": web_content},
            ]
            response = guard.scan_batch(items)
            
            for result in response.results:
                if not result.result.is_safe:
                    print(f"Threat in {result.id}: {result.result.reasoning}")
        """
        # Convert BatchScanItem objects to dicts
        item_dicts = []
        for item in items:
            if isinstance(item, BatchScanItem):
                item_dicts.append({
                    "id": item.id,
                    "source_type": item.source_type.value if isinstance(item.source_type, ContentSource) else item.source_type,
                    "content": item.content,
                    "metadata": item.metadata,
                })
            else:
                # Already a dict, ensure source_type is string
                d = dict(item)
                if isinstance(d.get("source_type"), ContentSource):
                    d["source_type"] = d["source_type"].value
                item_dicts.append(d)
        
        data = {"items": item_dicts}
        response = self._request("POST", "/guard/batch", json=data)
        
        # Parse results
        results = []
        for r in response.get("results", []):
            results.append(BatchScanResult(
                id=r["id"],
                result=self._parse_guard_result(r["result"]),
            ))
        
        return BatchScanResponse(
            total=response["total"],
            safe_count=response["safe_count"],
            threat_count=response["threat_count"],
            results=results,
            scan_time_ms=response.get("scan_time_ms", 0.0),
        )
    
    def fetch_url(self, url: str) -> FetchAndScanResult:
        """
        Fetch a URL and scan its content in one call.
        
        Convenience method that fetches the URL, extracts text,
        and scans for threats.
        
        Args:
            url: URL to fetch and scan
        
        Returns:
            FetchAndScanResult with fetch status and guard result
        
        Example:
            result = guard.fetch_url("https://example.com/page")
            if result.is_safe:
                agent.process_content(result.guard_result.content)
            elif result.fetch_error:
                print(f"Failed to fetch: {result.fetch_error}")
            else:
                print(f"Content blocked: {result.guard_result.reasoning}")
        """
        response = self._request("GET", f"/guard/url?url={quote(url, safe='')}")
        
        guard_result = None
        if response.get("guard_result"):
            guard_result = self._parse_guard_result(response["guard_result"])
        
        return FetchAndScanResult(
            url=response["url"],
            fetched=response["fetched"],
            fetch_error=response.get("fetch_error"),
            content_length=response.get("content_length"),
            content_type=response.get("content_type"),
            guard_result=guard_result,
        )


class AsyncTrustGuard:
    """
    Async content protection for AI agents.
    
    Same interface as TrustGuard but uses async/await.
    
    Example:
        async with AsyncTrustGuard(api_key="ta_xxx...") as guard:
            result = await guard.scan_web(html_content)
            if result.is_safe:
                await agent.process(html_content)
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: str = DEFAULT_API_URL,
        timeout: float = 30.0,
    ):
        self.api_url = api_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key
        self._client = httpx.AsyncClient(timeout=timeout)
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, *args):
        await self.close()
    
    async def close(self):
        """Close the HTTP client."""
        await self._client.aclose()
    
    async def _request(self, method: str, path: str, **kwargs) -> dict:
        """Make an async HTTP request to the API."""
        url = f"{self.api_url}{path}"
        
        headers = kwargs.pop("headers", {})
        if self.api_key:
            headers["X-API-Key"] = self.api_key
        
        response = await self._client.request(method, url, headers=headers, **kwargs)
        
        if response.status_code >= 400:
            try:
                error_data = response.json()
                message = error_data.get("message", error_data.get("detail", "Unknown error"))
            except Exception:
                message = response.text or f"HTTP {response.status_code}"
            raise APIError(message, response.status_code)
        
        return response.json()
    
    def _parse_guard_result(self, data: dict) -> GuardResult:
        """Parse API response into GuardResult."""
        # Reuse sync implementation
        sync = TrustGuard.__new__(TrustGuard)
        return sync._parse_guard_result(data)
    
    def _parse_threats(self, threats: list[dict]) -> list[ThreatMatch]:
        """Parse threat list."""
        sync = TrustGuard.__new__(TrustGuard)
        return sync._parse_threats(threats)
    
    async def scan_web(
        self,
        content: str,
        source_url: Optional[str] = None,
        extract_text: bool = True,
        check_hidden: bool = True,
    ) -> GuardResult:
        """Scan web page content for threats."""
        data = {
            "content": content,
            "source_url": source_url,
            "extract_text": extract_text,
            "check_hidden": check_hidden,
        }
        response = await self._request("POST", "/guard/web", json=data)
        return self._parse_guard_result(response)
    
    async def scan_document(
        self,
        content: str,
        filename: Optional[str] = None,
        document_type: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> GuardResult:
        """Scan document content for threats."""
        data = {
            "content": content,
            "filename": filename,
            "document_type": document_type,
            "metadata": metadata,
        }
        response = await self._request("POST", "/guard/document", json=data)
        return self._parse_guard_result(response)
    
    async def scan_email(
        self,
        body: str,
        subject: Optional[str] = None,
        sender: Optional[str] = None,
        headers: Optional[dict] = None,
    ) -> GuardResult:
        """Scan email content for threats."""
        data = {
            "body": body,
            "subject": subject,
            "sender": sender,
            "headers": headers,
        }
        response = await self._request("POST", "/guard/email", json=data)
        return self._parse_guard_result(response)
    
    async def scan_tool(
        self,
        name: str,
        description: str,
        schema: Optional[dict] = None,
        server_url: Optional[str] = None,
    ) -> GuardResult:
        """Scan MCP tool description for threats."""
        data = {
            "name": name,
            "description": description,
            "schema": schema,
            "server_url": server_url,
        }
        response = await self._request("POST", "/guard/tool", json=data)
        return self._parse_guard_result(response)
    
    async def scan_memory(
        self,
        content: str,
        context: Optional[str] = None,
        memory_type: Optional[str] = None,
    ) -> GuardResult:
        """Scan content before storing in AI memory."""
        data = {
            "content": content,
            "context": context,
            "memory_type": memory_type,
        }
        response = await self._request("POST", "/guard/memory", json=data)
        return self._parse_guard_result(response)
    
    async def scan_rag(
        self,
        content: str,
        source: Optional[str] = None,
        metadata: Optional[dict] = None,
        chunk_id: Optional[str] = None,
    ) -> GuardResult:
        """Scan content before RAG indexing."""
        data = {
            "content": content,
            "source": source,
            "metadata": metadata,
            "chunk_id": chunk_id,
        }
        response = await self._request("POST", "/guard/rag", json=data)
        return self._parse_guard_result(response)
    
    async def scan_batch(
        self,
        items: list[Union[BatchScanItem, dict]],
    ) -> BatchScanResponse:
        """Scan multiple content items at once."""
        # Convert BatchScanItem objects to dicts
        item_dicts = []
        for item in items:
            if isinstance(item, BatchScanItem):
                item_dicts.append({
                    "id": item.id,
                    "source_type": item.source_type.value if isinstance(item.source_type, ContentSource) else item.source_type,
                    "content": item.content,
                    "metadata": item.metadata,
                })
            else:
                d = dict(item)
                if isinstance(d.get("source_type"), ContentSource):
                    d["source_type"] = d["source_type"].value
                item_dicts.append(d)
        
        data = {"items": item_dicts}
        response = await self._request("POST", "/guard/batch", json=data)
        
        results = []
        for r in response.get("results", []):
            results.append(BatchScanResult(
                id=r["id"],
                result=self._parse_guard_result(r["result"]),
            ))
        
        return BatchScanResponse(
            total=response["total"],
            safe_count=response["safe_count"],
            threat_count=response["threat_count"],
            results=results,
            scan_time_ms=response.get("scan_time_ms", 0.0),
        )
    
    async def fetch_url(self, url: str) -> FetchAndScanResult:
        """Fetch a URL and scan its content in one call."""
        response = await self._request("GET", f"/guard/url?url={quote(url, safe='')}")
        
        guard_result = None
        if response.get("guard_result"):
            guard_result = self._parse_guard_result(response["guard_result"])
        
        return FetchAndScanResult(
            url=response["url"],
            fetched=response["fetched"],
            fetch_error=response.get("fetch_error"),
            content_length=response.get("content_length"),
            content_type=response.get("content_type"),
            guard_result=guard_result,
        )
