"""agent-shield — Compliance middleware for AI agents.

Add GDPR/NDPA compliance to any AI agent in 3 lines of code.
PII detection, audit logging, consent management, and DPIA generation.

    from agent_shield import Shield

    shield = Shield()
    result = shield.scan("Contact me at john@example.com or 08012345678")
    print(result.pii_found)     # {'EMAIL': 1, 'PHONE_NG': 1}
    print(result.redacted)      # "Contact me at [EMAIL_REDACTED] or [PHONE_REDACTED]"

Built by Janus Compliance — https://januscompliance.co.uk
"""

__version__ = "0.1.0"

from dataclasses import dataclass
from typing import Any, Callable, Optional

from . import pii
from . import dpia
from . import dataflow
from .audit import AuditLog
from .consent import ConsentManager
from .wrapper import wrap_openai_call, wrap_anthropic_call, wrap_generic_call, ShieldedCall


@dataclass
class ScanResult:
    """Result of scanning text for PII."""
    original: str
    redacted: str
    pii_found: dict[str, int]
    has_pii: bool
    matches: list


class Shield:
    """Main interface for agent-shield.

    Usage:
        shield = Shield(audit_path="./audit.jsonl")

        # Scan text for PII
        result = shield.scan("Email me at test@example.com")

        # Wrap an OpenAI call
        response = shield.call_openai(
            client=openai_client,
            messages=[{"role": "user", "content": "Hello"}],
            purpose="customer_support",
            user_id="user_123",
        )

        # Generate a DPIA
        dpia_text = shield.generate_dpia(system_name="My Agent")

        # Generate a data flow map
        flow_map = shield.generate_dataflow()
    """

    def __init__(
        self,
        audit_path: str = "./agent_shield_audit.jsonl",
        consent_path: str = "./agent_shield_consent.jsonl",
        redact_by_default: bool = False,
        min_confidence: float = 0.5,
    ):
        self.audit = AuditLog(audit_path)
        self.consent = ConsentManager(consent_path)
        self.redact_by_default = redact_by_default
        self.min_confidence = min_confidence

    def scan(self, text: str) -> ScanResult:
        """Scan text for PII."""
        matches = pii.detect(text, min_confidence=self.min_confidence)
        redacted = pii.redact(text, min_confidence=self.min_confidence)
        found = pii.summary(text, min_confidence=self.min_confidence)

        return ScanResult(
            original=text,
            redacted=redacted,
            pii_found=found,
            has_pii=bool(found),
            matches=matches,
        )

    def call_openai(
        self,
        client: Any,
        messages: list[dict],
        model: str = "gpt-4",
        purpose: str = "",
        user_id: Optional[str] = None,
        redact_input: Optional[bool] = None,
        **kwargs,
    ) -> ShieldedCall:
        """Make a shielded OpenAI API call.

        Requires: pip install agent-shield[openai]
        """
        return wrap_openai_call(
            client=client,
            messages=messages,
            model=model,
            audit_log=self.audit,
            purpose=purpose,
            user_id=user_id,
            redact_input=redact_input if redact_input is not None else self.redact_by_default,
            **kwargs,
        )

    def call_anthropic(
        self,
        client: Any,
        messages: list[dict],
        model: str = "claude-sonnet-4-20250514",
        purpose: str = "",
        user_id: Optional[str] = None,
        redact_input: Optional[bool] = None,
        max_tokens: int = 1024,
        **kwargs,
    ) -> ShieldedCall:
        """Make a shielded Anthropic API call.

        Requires: pip install agent-shield[anthropic]
        """
        return wrap_anthropic_call(
            client=client,
            messages=messages,
            model=model,
            audit_log=self.audit,
            purpose=purpose,
            user_id=user_id,
            redact_input=redact_input if redact_input is not None else self.redact_by_default,
            max_tokens=max_tokens,
            **kwargs,
        )

    def call(
        self,
        fn: Callable,
        input_text: str,
        provider: str = "custom",
        model: str = "custom",
        purpose: str = "",
        user_id: Optional[str] = None,
        redact_input: Optional[bool] = None,
        **kwargs,
    ) -> ShieldedCall:
        """Wrap any LLM call function with compliance features."""
        return wrap_generic_call(
            fn=fn,
            input_text=input_text,
            provider=provider,
            model=model,
            audit_log=self.audit,
            purpose=purpose,
            user_id=user_id,
            redact_input=redact_input if redact_input is not None else self.redact_by_default,
            **kwargs,
        )

    def generate_dpia(
        self,
        system_name: str = "AI Agent System",
        controller_name: str = "",
        dpo_contact: str = "",
    ) -> str:
        """Generate a DPIA skeleton from audit log data. Returns Markdown."""
        return dpia.generate(
            audit_log=self.audit,
            system_name=system_name,
            controller_name=controller_name,
            dpo_contact=dpo_contact,
        )

    def generate_dataflow(self) -> str:
        """Generate a data flow map from audit log data. Returns Markdown."""
        return dataflow.generate_map(self.audit)

    def generate_dataflow_diagram(self, app_name: str = "Your App") -> str:
        """Generate a Mermaid diagram of data flows."""
        return dataflow.generate_mermaid(self.audit, app_name=app_name)

    def verify_audit(self) -> tuple[bool, int]:
        """Verify the audit log chain integrity. Returns (is_valid, entry_count)."""
        return self.audit.verify_chain()

    def user_report(self, user_id: str) -> dict:
        """Generate a complete data report for a user (data subject access request)."""
        return {
            "user_id": user_id,
            "audit_entries": self.audit.export_for_user(user_id),
            "consent_history": self.consent.export(user_id),
            "active_consents": self.consent.active_consents(user_id),
        }
