"""
{{ service_name }} Guardian Agent
Generated by ABI-Core scaffolding
"""

import logging
import os
from typing import Dict, Any, Optional, AsyncIterable
from collections.abc import AsyncIterable

from abi_core.core.agents.base_agent import BaseAgent
from abi_core.core.common.types import TaskState, TaskStatusUpdateEvent, TaskArtifactUpdateEvent
from abi_core.core.common.utils import abi_logging
from abi_core.core.opa.policy_loader_v2 import PolicyLoaderV2

from .guard_core.custom_policies import get_custom_policy_generator

from langchain_core.messages import AIMessage
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent

logger = logging.getLogger(__name__)

MODEL_NAME = os.getenv('MODEL_NAME', '{{ model_name | default('llama3.2:3b') }}')
memory = MemorySaver()

class {{ service_class_name }}Guardian(BaseAgent):
    """
    {{ service_name }} Guardian Agent
    Provides domain-specific security and policy enforcement
    
    Features:
    - Custom security policies for {{ domain | default('general') }} domain
    - Real-time threat detection and response
    - Compliance monitoring and reporting
    - Integration with ABI core security framework
    """
    
    def __init__(self):
        super().__init__(
            agent_name='{{ service_name }} Guardian',
            description='Domain-specific security and policy enforcement for {{ service_name }}',
            content_types=['text', 'text/plain', 'application/json']
        )
        
        # Initialize LLM
        self.llm = ChatOllama(
            model=MODEL_NAME,
            base_url=os.getenv('OLLAMA_HOST', 'http://localhost:11434'),
            temperature=0.1  # Low temperature for consistent security decisions
        )
        
        # Initialize policy components
        self.policy_generator = get_custom_policy_generator()
        self.policy_loader = PolicyLoaderV2()
        
        # Security configuration
        self.domain = "{{ domain | default('general') }}"
        self.service_name = "{{ service_name }}"
        self.risk_threshold = {{ risk_threshold | default('0.7') }}
        
        # Initialize agent with security tools
        self.agent = create_react_agent(
            self.llm,
            tools=self._get_security_tools(),
            checkpointer=memory
        )
        
        # Load custom policies
        self._initialize_custom_policies()
        
        abi_logging(f'{{ service_name }} Guardian initialized for domain: {self.domain}')
    
    def _initialize_custom_policies(self):
        """Initialize custom security policies"""
        try:
            # Ensure custom policies exist
            policy_dir = os.getenv('POLICY_DIR', '/app/policies')
            os.makedirs(policy_dir, exist_ok=True)
            
            custom_policy_file = os.path.join(policy_dir, f"{{ service_name.lower().replace('-', '_') }}_custom.rego")
            
            if not os.path.exists(custom_policy_file):
                logger.info("Generating custom policies...")
                self.policy_generator.write_custom_policies(custom_policy_file)
            
            # Validate policies
            if self.policy_generator.validate_custom_policies(custom_policy_file):
                logger.info("Custom policies loaded and validated successfully")
            else:
                logger.error("Custom policy validation failed")
                
        except Exception as e:
            logger.error(f"Failed to initialize custom policies: {e}")
    
    def _get_security_tools(self):
        """Get security-specific tools"""
        tools = []
        
        # Add policy evaluation tool
        tools.append(self._create_policy_evaluation_tool())
        
        # Add threat detection tool
        tools.append(self._create_threat_detection_tool())
        
        # Add compliance check tool
        tools.append(self._create_compliance_check_tool())
        
        {% if custom_security_tools %}
        # Add custom security tools
        {% for tool in custom_security_tools %}
        tools.append(self._create_{{ tool }}_tool())
        {% endfor %}
        {% endif %}
        
        return tools
    
    def _create_policy_evaluation_tool(self):
        """Create policy evaluation tool"""
        from langchain_core.tools import tool
        
        @tool
        def evaluate_policy(action: str, resource_type: str, source_agent: str, content: str = "") -> str:
            """Evaluate security policy for a given action"""
            try:
                policy_input = {
                    "action": action,
                    "resource_type": resource_type,
                    "source_agent": source_agent,
                    "content": content,
                    "domain": self.domain,
                    "service": self.service_name
                }
                
                # TODO: Integrate with OPA for actual policy evaluation
                # For now, return a mock evaluation
                risk_score = self._calculate_risk_score(policy_input)
                
                if risk_score > self.risk_threshold:
                    return f"DENY: High risk action (score: {risk_score:.2f}). Action '{action}' on '{resource_type}' by '{source_agent}' exceeds risk threshold."
                else:
                    return f"ALLOW: Low risk action (score: {risk_score:.2f}). Action '{action}' on '{resource_type}' by '{source_agent}' is permitted."
                    
            except Exception as e:
                return f"ERROR: Policy evaluation failed: {str(e)}"
        
        return evaluate_policy
    
    def _create_threat_detection_tool(self):
        """Create threat detection tool"""
        from langchain_core.tools import tool
        
        @tool
        def detect_threats(content: str, context: str = "") -> str:
            """Detect security threats in content"""
            try:
                threats = []
                
                # Check for sensitive data patterns
                if self._contains_sensitive_data(content):
                    threats.append("Sensitive data detected")
                
                # Check for malicious patterns
                if self._contains_malicious_patterns(content):
                    threats.append("Malicious patterns detected")
                
                {% if custom_threat_patterns %}
                # Check custom threat patterns
                {% for pattern in custom_threat_patterns %}
                if self._check_{{ pattern.name }}_pattern(content):
                    threats.append("{{ pattern.description }}")
                {% endfor %}
                {% endif %}
                
                if threats:
                    return f"THREATS DETECTED: {', '.join(threats)}"
                else:
                    return "No threats detected"
                    
            except Exception as e:
                return f"ERROR: Threat detection failed: {str(e)}"
        
        return detect_threats
    
    def _create_compliance_check_tool(self):
        """Create compliance checking tool"""
        from langchain_core.tools import tool
        
        @tool
        def check_compliance(action: str, resource_type: str, metadata: str = "{}") -> str:
            """Check compliance with domain-specific regulations"""
            try:
                import json
                meta = json.loads(metadata) if metadata else {}
                
                compliance_issues = []
                
                {% if compliance_rules %}
                # Check domain-specific compliance rules
                {% for rule in compliance_rules %}
                if not self._check_{{ rule.name }}_compliance(action, resource_type, meta):
                    compliance_issues.append("{{ rule.description }}")
                {% endfor %}
                {% endif %}
                
                if compliance_issues:
                    return f"COMPLIANCE VIOLATIONS: {', '.join(compliance_issues)}"
                else:
                    return "Compliance check passed"
                    
            except Exception as e:
                return f"ERROR: Compliance check failed: {str(e)}"
        
        return check_compliance
    
    {% if custom_security_tools %}
    {% for tool in custom_security_tools %}
    def _create_{{ tool }}_tool(self):
        """Create {{ tool }} security tool"""
        from langchain_core.tools import tool
        
        @tool
        def {{ tool }}(input_data: str) -> str:
            """{{ tool.title().replace('_', ' ') }} security functionality"""
            try:
                # TODO: Implement {{ tool }} logic
                return f"{{ tool.title().replace('_', ' ') }} executed for: {input_data}"
            except Exception as e:
                return f"ERROR: {{ tool }} failed: {str(e)}"
        
        return {{ tool }}
    {% endfor %}
    {% endif %}
    
    def _calculate_risk_score(self, policy_input: Dict[str, Any]) -> float:
        """Calculate risk score for policy input"""
        # Base risk scores
        action_risks = {
            "read": 0.1,
            "write": 0.4,
            "delete": 0.8,
            "execute": 0.9,
            "create_agent": 1.0,
            "modify_policy": 1.0
        }
        
        base_risk = action_risks.get(policy_input.get("action", ""), 0.5)
        
        # Apply domain-specific multipliers
        if policy_input.get("resource_type") in ["{{ domain | default('general') }}_data", "sensitive_config"]:
            base_risk *= 1.5
        
        # Check for sensitive content
        if self._contains_sensitive_data(policy_input.get("content", "")):
            base_risk += 0.3
        
        return min(base_risk, 1.0)
    
    def _contains_sensitive_data(self, content: str) -> bool:
        """Check if content contains sensitive data"""
        import re
        
        sensitive_patterns = [
            r'\b\d{3}-\d{2}-\d{4}\b',  # SSN
            r'(?i)(api[_-]?key|token|secret|password)["\\s]*[:=]["\\s]*[a-zA-Z0-9]{20,}',  # API keys
            {% if custom_sensitive_patterns %}
            {% for pattern in custom_sensitive_patterns %}
            r'{{ pattern.regex }}',  # {{ pattern.description }}
            {% endfor %}
            {% endif %}
        ]
        
        for pattern in sensitive_patterns:
            if re.search(pattern, content):
                return True
        
        return False
    
    def _contains_malicious_patterns(self, content: str) -> bool:
        """Check if content contains malicious patterns"""
        import re
        
        malicious_patterns = [
            r'(?i)(rm\s+-rf|del\s+/|format\s+c:|shutdown|reboot)',  # Destructive commands
            r'(?i)(drop\s+table|delete\s+from|truncate|alter\s+table)',  # SQL injection
            {% if custom_malicious_patterns %}
            {% for pattern in custom_malicious_patterns %}
            r'{{ pattern.regex }}',  # {{ pattern.description }}
            {% endfor %}
            {% endif %}
        ]
        
        for pattern in malicious_patterns:
            if re.search(pattern, content):
                return True
        
        return False
    
    {% if custom_threat_patterns %}
    {% for pattern in custom_threat_patterns %}
    def _check_{{ pattern.name }}_pattern(self, content: str) -> bool:
        """Check for {{ pattern.description }}"""
        import re
        return bool(re.search(r'{{ pattern.regex }}', content))
    {% endfor %}
    {% endif %}
    
    {% if compliance_rules %}
    {% for rule in compliance_rules %}
    def _check_{{ rule.name }}_compliance(self, action: str, resource_type: str, metadata: Dict[str, Any]) -> bool:
        """Check {{ rule.description }} compliance"""
        # TODO: Implement {{ rule.name }} compliance check
        return True
    {% endfor %}
    {% endif %}
    
    async def stream(
        self, 
        query: str, 
        context_id: str, 
        task_id: str
    ) -> AsyncIterable[Dict[str, Any]]:
        """
        Process security query and stream responses
        
        Args:
            query: Security query to process
            context_id: Context identifier
            task_id: Task identifier
            
        Yields:
            Dict containing security response data
        """
        abi_logging(f'{{ service_name }} Guardian processing security query: {query[:100]}...')
        
        try:
            # Prepare agent configuration
            config = {
                "configurable": {
                    "thread_id": f"{context_id}-{task_id}"
                }
            }
            
            # Add security context to query
            enhanced_query = f"""
            Security Analysis Request for {{ service_name }}:
            Domain: {self.domain}
            Query: {query}
            
            Please analyze this request for security implications and provide recommendations.
            """
            
            # Process with security agent
            response_chunks = []
            async for chunk in self.agent.astream(
                {"messages": [("user", enhanced_query)]}, 
                config=config
            ):
                if "agent" in chunk:
                    response_chunks.append(chunk["agent"]["messages"][0])
                elif "tools" in chunk:
                    # Handle security tool responses
                    for tool_msg in chunk["tools"]["messages"]:
                        abi_logging(f'Security tool response: {tool_msg.content[:100]}...')
            
            # Combine response with security metadata
            if response_chunks:
                final_response = response_chunks[-1]
                content = final_response.content if hasattr(final_response, 'content') else str(final_response)
                
                yield {
                    'content': content,
                    'response_type': 'text',
                    'is_task_completed': True,
                    'require_user_input': False,
                    'metadata': {
                        'agent': '{{ service_name }} Guardian',
                        'domain': self.domain,
                        'security_level': 'high',
                        'model': MODEL_NAME,
                        'context_id': context_id,
                        'task_id': task_id
                    }
                }
            else:
                yield {
                    'content': 'Security analysis completed - no specific recommendations',
                    'response_type': 'text', 
                    'is_task_completed': True,
                    'require_user_input': False
                }
                
        except Exception as e:
            logger.error(f'Error in {{ service_name }} Guardian: {e}')
            yield {
                'content': f'Security analysis error: {str(e)}',
                'response_type': 'text',
                'is_task_completed': True,
                'require_user_input': False
            }

# Guardian instance
guardian = {{ service_class_name }}Guardian()