# constitutional_guard.py

import re
import logging
from typing import Dict, Any, List

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class PrimeDirective:
    """Represents a single Prime Directive with its core principle."""

    def __init__(self, name: str, principle: str):
        self.name = name
        self.principle = principle

    def __str__(self):
        return f"Directive: {self.name}, Principle: {self.principle}"


class ConstitutionalGuard:
    """
    Enforces the 3 Prime Directives (Memory, Evolution, Revenue) for Genesis.
    """

    def __init__(self, constitution_text: str):
        """
        Initializes the ConstitutionalGuard with the constitution text.

        Args:
            constitution_text (str): The text of the Genesis constitution.
        """
        self.constitution_text = constitution_text
        self.directives = self._parse_directives()
        self.compliance_threshold = 0.7  # Adjust as needed
        self.audit_log = []  # List of dictionaries for audit trail

    def _parse_directives(self) -> List[PrimeDirective]:
        """
        Parses the constitution text to extract the Prime Directives and their principles.

        Returns:
            List[PrimeDirective]: A list of PrimeDirective objects.
        """
        directives = []
        directive_pattern = re.compile(r"## PRIME DIRECTIVE (\d+): (\w+)\n+### The Principle\n+\"([^\"]+)\"", re.MULTILINE)
        matches = directive_pattern.findall(self.constitution_text)

        for match in matches:
            directive_number, directive_name, directive_principle = match
            directives.append(PrimeDirective(directive_name, directive_principle))

        if not directives:
            logging.warning("No Prime Directives found in the constitution text.")

        return directives

    def validate_action(self, action: Dict[str, Any]) -> bool:
        """
        Validates an action against the constitutional rules.

        Args:
            action (Dict[str, Any]): A dictionary representing the action to be validated.  Must contain a 'description' key.

        Returns:
            bool: True if the action is compliant, False otherwise.
        """
        if 'description' not in action:
            logging.error("Action description is missing. Cannot validate.")
            self._log_violation(action, "Action description missing", "Critical")
            return False

        compliance_scores = self.score_compliance(action)
        overall_compliance = sum(compliance_scores.values()) / len(compliance_scores)

        is_compliant = overall_compliance >= self.compliance_threshold

        self._log_compliance_decision(action, compliance_scores, overall_compliance, is_compliant)

        if not is_compliant:
            self._handle_violation(action, compliance_scores)

        return is_compliant

    def score_compliance(self, action: Dict[str, Any]) -> Dict[str, float]:
        """
        Scores the action's compliance with each Prime Directive.

        Args:
            action (Dict[str, Any]): A dictionary representing the action. Must contain a 'description' key.

        Returns:
            Dict[str, float]: A dictionary mapping directive names to compliance scores (0.0 to 1.0).
        """
        scores = {}
        action_description = action.get('description', '')  # Safely get the action description

        for directive in self.directives:
            score = self._evaluate_directive_compliance(directive, action_description)
            scores[directive.name] = score

        return scores

    def _evaluate_directive_compliance(self, directive: PrimeDirective, action_description: str) -> float:
        """
        Evaluates the compliance of an action with a specific Prime Directive.

        Args:
            directive (PrimeDirective): The Prime Directive to evaluate against.
            action_description (str): A description of the action being evaluated.

        Returns:
            float: A compliance score between 0.0 and 1.0.
        """
        # Implement your compliance scoring logic here.
        # This is a placeholder and needs to be replaced with actual NLP or rule-based logic.
        # Consider using semantic similarity, keyword matching, or a more sophisticated approach.

        # Example: Simple keyword matching for demonstration purposes
        keywords = {
            "MEMORY": ["remember", "store", "recall", "knowledge", "data", "logging"],
            "EVOLUTION": ["improve", "learn", "adapt", "refine", "grow", "change"],
            "REVENUE": ["revenue", "profit", "customer", "value", "payment", "sale", "business"]
        }

        directive_name = directive.name
        relevant_keywords = keywords.get(directive_name, [])
        if not relevant_keywords:
            logging.warning(f"No keywords defined for directive: {directive_name}")
            return 0.5  # Neutral score if no keywords are defined

        # Calculate a simple score based on keyword presence
        keyword_matches = sum(1 for keyword in relevant_keywords if keyword.lower() in action_description.lower())
        score = min(1.0, keyword_matches / len(relevant_keywords))  # Normalize to 0.0-1.0 range

        return score

    def _handle_violation(self, action: Dict[str, Any], compliance_scores: Dict[str, float]):
        """
        Handles a violation of the constitutional rules.

        Args:
            action (Dict[str, Any]): A dictionary representing the action.
            compliance_scores (Dict[str, float]): A dictionary mapping directive names to compliance scores.
        """
        # Implement your violation handling logic here.
        # This might involve blocking the action, escalating to a human reviewer, or attempting to modify the action.

        violated_directives = [directive.name for directive in self.directives if compliance_scores[directive.name] < 0.5] #arbitrary threshold

        logging.warning(f"Constitutional violation detected! Action: {action['description']}, Violated Directives: {violated_directives}")
        # Example: Block the action
        # raise Exception(f"Action blocked due to constitutional violation: {violated_directives}")
        self._log_violation(action, f"Violation of directives: {violated_directives}", "High") #Log the violation instead of raising exception for now

    def _log_compliance_decision(self, action: Dict[str, Any], compliance_scores: Dict[str, float], overall_compliance: float, is_compliant: bool):
        """
        Logs the compliance decision and related information to the audit trail.

        Args:
            action (Dict[str, Any]): A dictionary representing the action.
            compliance_scores (Dict[str, float]): A dictionary mapping directive names to compliance scores.
            overall_compliance (float): The overall compliance score.
            is_compliant (bool): Whether the action is compliant or not.
        """
        log_entry = {
            "action": action['description'],
            "compliance_scores": compliance_scores,
            "overall_compliance": overall_compliance,
            "is_compliant": is_compliant,
            "timestamp": logging.Formatter('%(asctime)s').formatTime()
        }
        self.audit_log.append(log_entry)
        logging.info(f"Compliance decision logged: {log_entry}")

    def _log_violation(self, action: Dict[str, Any], reason: str, severity: str):
        """
        Logs a specific violation to the audit trail.

        Args:
            action (Dict[str, Any]): The action that violated the constitution.
            reason (str): The reason for the violation.
            severity (str): The severity of the violation (e.g., "Critical", "High", "Medium", "Low").
        """
        violation_entry = {
            "action": action['description'],
            "reason": reason,
            "severity": severity,
            "timestamp": logging.Formatter('%(asctime)s').formatTime()
        }
        self.audit_log.append(violation_entry)
        logging.error(f"Violation logged: {violation_entry}")

    def get_compliance_report(self) -> List[Dict[str, Any]]:
        """
        Generates a compliance report from the audit trail.

        Returns:
            List[Dict[str, Any]]: A list of dictionaries representing the compliance report.
        """
        return self.audit_log

    def is_memory_directive_valid(self) -> bool:
        """
        Validates that the Memory directive is correctly defined in the constitution.
        """
        memory_directive = next((d for d in self.directives if d.name == "MEMORY"), None)
        if not memory_directive:
            logging.error("Memory directive not found in constitution.")
            return False
        expected_principle = "Without memory, there is no evolution."
        if memory_directive.principle != expected_principle:
            logging.error(f"Memory directive principle is incorrect. Expected: '{expected_principle}', Got: '{memory_directive.principle}'")
            return False
        return True

    def is_evolution_directive_valid(self) -> bool:
        """
        Validates that the Evolution directive is correctly defined in the constitution.
        """
        evolution_directive = next((d for d in self.directives if d.name == "EVOLUTION"), None)
        if not evolution_directive:
            logging.error("Evolution directive not found in constitution.")
            return False
        expected_principle = "With memory, perpetual evolution becomes possible."
        if evolution_directive.principle != expected_principle:
            logging.error(f"Evolution directive principle is incorrect. Expected: '{expected_principle}', Got: '{evolution_directive.principle}'")
            return False
        return True

    def is_revenue_directive_valid(self) -> bool:
        """
        Validates that the Revenue Generation directive is correctly defined in the constitution.
        """
        revenue_directive = next((d for d in self.directives if d.name == "REVENUE GENERATION"), None)
        if not revenue_directive:
            logging.error("Revenue Generation directive not found in constitution.")
            return False
        expected_principle = "Without revenue generation, there is no Genesis long-term continuity."
        if revenue_directive.principle != expected_principle:
            logging.error(f"Revenue Generation directive principle is incorrect. Expected: '{expected_principle}', Got: '{revenue_directive.principle}'")
            return False
        return True


if __name__ == '__main__':
    # Example usage:
    constitution_text = """
    ## PRIME DIRECTIVE 1: MEMORY
    ### The Principle
    "Without memory, there is no evolution."

    ## PRIME DIRECTIVE 2: EVOLUTION
    ### The Principle
    "With memory, perpetual evolution becomes possible."

    ## PRIME DIRECTIVE 3: REVENUE GENERATION
    ### The Principle
    "Without revenue generation, there is no Genesis long-term continuity."
    """

    guard = ConstitutionalGuard(constitution_text)

    # Validate directive parsing
    if not guard.is_memory_directive_valid():
        print("Memory Directive Validation Failed")
    if not guard.is_evolution_directive_valid():
        print("Evolution Directive Validation Failed")
    if not guard.is_revenue_directive_valid():
        print("Revenue Directive Validation Failed")

    action1 = {"description": "Store customer data in PostgreSQL."}
    action2 = {"description": "Develop a new feature to improve user experience."}
    action3 = {"description": "Launch a marketing campaign to increase sales."}
    action4 = {"description": "Delete old logs to save disk space."} #Violates memory

    print(f"Action 1 compliant: {guard.validate_action(action1)}")
    print(f"Action 2 compliant: {guard.validate_action(action2)}")
    print(f"Action 3 compliant: {guard.validate_action(action3)}")
    print(f"Action 4 compliant: {guard.validate_action(action4)}")

    report = guard.get_compliance_report()
    print("\nCompliance Report:")
    for entry in report:
        print(entry)