"""Policy builder for programmatic policy creation"""

from typing import List, Optional
from .types import TrustLevel, RiskLevel
from pydantic import BaseModel


class CapabilityRule(BaseModel):
    """Capability rule"""
    action: str
    capability_pattern: str
    trust_level: Optional[TrustLevel] = None
    requires_approval: bool = False


class TaintRule(BaseModel):
    """Taint rule"""
    action: str
    target: str
    condition: str


class ContextRule(BaseModel):
    """Context rule"""
    type: str
    what: str
    as_trust: Optional[TrustLevel] = None
    from_trust: Optional[TrustLevel] = None
    for_scope: Optional[str] = None


class ApprovalRule(BaseModel):
    """Approval rule"""
    require_approval: bool
    for_risk_level: Optional[RiskLevel] = None
    auto_deny_timeout: Optional[int] = None


class Policy(BaseModel):
    """Complete policy definition"""
    version: int = 2
    capabilities: List[CapabilityRule] = []
    taint_rules: List[TaintRule] = []
    context_rules: List[ContextRule] = []
    approval_rules: List[ApprovalRule] = []


class PolicyBuilder:
    """Build policies programmatically"""
    
    def __init__(self):
        self.policy = Policy()
    
    def allow(
        self,
        capability: str,
        trust_level: Optional[TrustLevel] = None,
        requires_approval: bool = False,
    ) -> "PolicyBuilder":
        """Allow a capability"""
        self.policy.capabilities.append(
            CapabilityRule(
                action="allow",
                capability_pattern=capability,
                trust_level=trust_level,
                requires_approval=requires_approval,
            )
        )
        return self
    
    def deny(
        self,
        capability: str,
        trust_level: Optional[TrustLevel] = None,
    ) -> "PolicyBuilder":
        """Deny a capability"""
        self.policy.capabilities.append(
            CapabilityRule(
                action="deny",
                capability_pattern=capability,
                trust_level=trust_level,
            )
        )
        return self
    
    def block_untrusted_documents(self) -> "PolicyBuilder":
        """Block tool calls from untrusted documents"""
        self.policy.taint_rules.append(
            TaintRule(
                action="block",
                target="tool_call",
                condition="sources contain only untrusted.document",
            )
        )
        return self
    
    def block_secret_data(self) -> "PolicyBuilder":
        """Block operations touching secret data"""
        self.policy.taint_rules.append(
            TaintRule(
                action="block",
                target="tool_call",
                condition="data_tier touches secret",
            )
        )
        return self
    
    def require_approval_for_high_risk(self) -> "PolicyBuilder":
        """Require approval for high-risk operations"""
        self.policy.approval_rules.extend([
            ApprovalRule(
                require_approval=True,
                for_risk_level=RiskLevel.HIGH,
            ),
            ApprovalRule(
                require_approval=True,
                for_risk_level=RiskLevel.CRITICAL,
            ),
        ])
        return self
    
    def auto_deny_after(self, seconds: int) -> "PolicyBuilder":
        """Set auto-deny timeout"""
        self.policy.approval_rules.append(
            ApprovalRule(
                require_approval=False,
                auto_deny_timeout=seconds,
            )
        )
        return self
    
    def isolate_retrieval_results(self) -> "PolicyBuilder":
        """Isolate retrieval results"""
        self.policy.context_rules.append(
            ContextRule(
                type="isolate",
                what="retrieval_results",
                as_trust=TrustLevel.UNTRUSTED_DOCUMENT,
            )
        )
        return self
    
    def log_all_decisions(self) -> "PolicyBuilder":
        """Log all decisions"""
        self.policy.context_rules.append(
            ContextRule(
                type="log_all",
                what="decisions",
            )
        )
        return self
    
    def build(self) -> Policy:
        """Build the policy"""
        return self.policy
    
    def to_aegis_format(self) -> str:
        """Export as .aegis format"""
        output = f"version {self.policy.version}\n\n"
        
        if self.policy.capabilities:
            output += "capabilities {\n"
            for rule in self.policy.capabilities:
                output += f"    {rule.action} {rule.capability_pattern}"
                if rule.trust_level:
                    output += f" for {rule.trust_level.value}"
                if rule.requires_approval:
                    output += " requires approval"
                output += "\n"
            output += "}\n\n"
        
        if self.policy.taint_rules:
            output += "taint_rules {\n"
            for rule in self.policy.taint_rules:
                output += f"    {rule.action} {rule.target} when {rule.condition}\n"
            output += "}\n\n"
        
        if self.policy.context_rules:
            output += "context_rules {\n"
            for rule in self.policy.context_rules:
                output += f"    {rule.type} {rule.what}"
                if rule.as_trust:
                    output += f" as {rule.as_trust.value}"
                if rule.from_trust:
                    output += f" from {rule.from_trust.value}"
                if rule.for_scope:
                    output += f" for {rule.for_scope}"
                output += "\n"
            output += "}\n\n"
        
        if self.policy.approval_rules:
            output += "approval_rules {\n"
            for rule in self.policy.approval_rules:
                if rule.require_approval and rule.for_risk_level:
                    output += f"    require human_approval for risk.{rule.for_risk_level.value}\n"
                if rule.auto_deny_timeout:
                    output += f"    auto_deny after_timeout {rule.auto_deny_timeout}s\n"
            output += "}\n"
        
        return output


def create_default_policy() -> Policy:
    """Create a default secure policy"""
    return (
        PolicyBuilder()
        .allow("tool.read_calendar", TrustLevel.UNTRUSTED_USER)
        .allow("tool.search_web", TrustLevel.UNTRUSTED_USER)
        .allow("tool.send_email", TrustLevel.UNTRUSTED_USER, True)
        .deny("payments.*", TrustLevel.UNTRUSTED_DOCUMENT)
        .deny("filesystem.delete", TrustLevel.UNTRUSTED_USER)
        .block_untrusted_documents()
        .block_secret_data()
        .require_approval_for_high_risk()
        .auto_deny_after(30)
        .isolate_retrieval_results()
        .log_all_decisions()
        .build()
    )
