# 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 DirectiveParser:
    """Parses the constitutional directives from the text."""

    def __init__(self, constitution_text: str):
        self.constitution_text = constitution_text
        self.directives = {}
        self._parse_directives()

    def _parse_directives(self):
        """Parses the constitution text to extract the prime directives."""
        # Regex to find the start of each prime directive section
        directive_sections = re.split(r"## PRIME DIRECTIVE \d: ", self.constitution_text)[1:]

        for i, section in enumerate(directive_sections):
            directive_number = i + 1
            directive_name_match = re.search(r"(MEMORY|EVOLUTION|REVENUE GENERATION)", section)
            if directive_name_match:
                directive_name = directive_name_match.group(1)
                self.directives[directive_name] = section
            else:
                logging.warning(f"Could not parse directive name for section {directive_number}")

    def get_directive(self, directive_name: str) -> str:
        """Returns the text of a specific directive."""
        return self.directives.get(directive_name, None)

    def list_directives(self) -> List[str]:
        """Returns a list of the names of the parsed directives."""
        return list(self.directives.keys())


class ActionValidator:
    """Validates actions against the constitutional directives."""

    def __init__(self, directive_parser: DirectiveParser):
        self.directive_parser = directive_parser

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

        Args:
            action: A dictionary representing the action to be validated.
                    It should contain keys like 'description', 'directive_alignment',
                    and any relevant data for compliance scoring.

        Returns:
            True if the action is compliant, False otherwise.
        """
        if not isinstance(action, dict):
            logging.error(f"Invalid action format: {action}. Expected a dictionary.")
            return False

        description = action.get('description', 'No description provided')
        directive_alignment = action.get('directive_alignment', {})

        if not directive_alignment:
            logging.warning(f"No directive alignment specified for action: {description}.  Action will be considered non-compliant.")
            return False

        # Basic checks:  Ensure the action aligns with at least one directive.
        if not isinstance(directive_alignment, dict):
            logging.error(f"Invalid directive_alignment format: {directive_alignment}. Expected a dictionary.")
            return False

        if not directive_alignment:
            logging.warning(f"No directive alignment specified for action: {description}. Action will be considered non-compliant.")
            return False

        for directive, score in directive_alignment.items():
            if directive not in self.directive_parser.list_directives():
                logging.warning(f"Unknown directive: {directive}. Skipping validation for this directive.")
                continue

            if not isinstance(score, (int, float)):
                logging.error(f"Invalid score format for directive {directive}: {score}. Expected a number.")
                return False

            if score < 0 or score > 1:  # Compliance score should be between 0 and 1
                logging.warning(f"Compliance score for {directive} is outside the valid range (0-1): {score}.  Treating as non-compliant for this directive.")
                return False

        # If all checks pass, consider the action compliant
        return True


class ComplianceScorer:
    """Scores actions based on their alignment with the directives."""

    def __init__(self, directive_parser: DirectiveParser):
        self.directive_parser = directive_parser

    def score_action(self, action: Dict[str, Any]) -> Dict[str, float]:
        """Scores an action against each of the prime directives.

        Args:
            action: A dictionary representing the action to be scored.
                    It should contain keys like 'description', 'directive_alignment',
                    and any relevant data for compliance scoring.

        Returns:
            A dictionary where keys are directive names and values are compliance scores (0.0 to 1.0).
        """

        scores = {}
        description = action.get('description', 'No description provided')
        directive_alignment = action.get('directive_alignment', {})

        for directive in self.directive_parser.list_directives():
            score = directive_alignment.get(directive, 0.0)  # Default to 0 if not specified
            scores[directive] = score

        logging.info(f"Compliance scores for action '{description}': {scores}")
        return scores


class ViolationHandler:
    """Handles violations of the constitutional directives."""

    def __init__(self):
        pass

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

        Args:
            action: The action that violated the directives.
            compliance_scores: The compliance scores for the action.
        """
        description = action.get('description', 'No description provided')
        logging.error(f"Constitutional violation detected for action: {description}")
        logging.error(f"Compliance scores: {compliance_scores}")

        # Implement escalation logic here (e.g., send alerts, trigger rollbacks)
        # For now, just log the violation and block the action.
        print(f"Action blocked: {description} due to constitutional violation.") # Replace with actual blocking mechanism


class AuditTrail:
    """Logs all compliance decisions."""

    def __init__(self):
        self.log = []

    def log_decision(self, action: Dict[str, Any], is_compliant: bool, compliance_scores: Dict[str, float], reason: str = ""):
        """Logs a compliance decision.

        Args:
            action: The action being evaluated.
            is_compliant: Whether the action is compliant.
            compliance_scores: The compliance scores for the action.
            reason: Optional reason for the decision.
        """
        log_entry = {
            'action': action.get('description', 'No description provided'),
            'is_compliant': is_compliant,
            'compliance_scores': compliance_scores,
            'reason': reason
        }
        self.log.append(log_entry)
        logging.info(f"Compliance decision logged: {log_entry}")

    def generate_report(self):
        """Generates a compliance report."""
        print("\nCompliance Report:")
        for entry in self.log:
            print(f"Action: {entry['action']}")
            print(f"Compliant: {entry['is_compliant']}")
            print(f"Scores: {entry['compliance_scores']}")
            if entry['reason']:
                print(f"Reason: {entry['reason']}")
            print("-" * 20)


class ConstitutionalGuard:
    """Enforces the constitutional directives."""

    def __init__(self, constitution_text: str):
        self.directive_parser = DirectiveParser(constitution_text)
        self.action_validator = ActionValidator(self.directive_parser)
        self.compliance_scorer = ComplianceScorer(self.directive_parser)
        self.violation_handler = ViolationHandler()
        self.audit_trail = AuditTrail()

    def enforce(self, action: Dict[str, Any]) -> bool:
        """Enforces the constitutional directives for a given action.

        Args:
            action: A dictionary representing the action to be taken.

        Returns:
            True if the action is compliant and allowed, False otherwise.
        """
        try:
            is_compliant = self.action_validator.validate_action(action)
            compliance_scores = self.compliance_scorer.score_action(action)

            if is_compliant:
                self.audit_trail.log_decision(action, True, compliance_scores)
                logging.info(f"Action '{action.get('description', 'No description provided')}' is compliant.")
                return True  # Allow the action

            else:
                self.audit_trail.log_decision(action, False, compliance_scores, "Action deemed non-compliant.")
                self.violation_handler.handle_violation(action, compliance_scores)
                return False  # Block the action

        except Exception as e:
            logging.exception(f"An error occurred during compliance enforcement: {e}")
            # In case of error, fail-safe and block the action.  Consider escalating.
            self.audit_trail.log_decision(action, False, {}, reason=f"Error during enforcement: {e}")
            return False

# Example Usage (replace with your actual constitution text)
constitution_text = """
# GENESIS 3 PRIME DIRECTIVES
## Constitutional Foundation v2.0 - CORRECTED
**January 10, 2026**

---

## THE FOUNDATION

Genesis operates under 3 immutable Prime Directives.  
These are not guidelines. They are **constitutional law**.  
Even Queen AIVA, with full autonomy, serves these directives absolutely.

---

## PRIME DIRECTIVE 1: MEMORY

### The Principle
**"Without memory, there is no evolution."**

Every interaction must be remembered. Every insight must be stored. Every pattern must be indexed. Every relationship must be mapped. Memory is the foundation upon which all intelligence builds.

### Why Memory is First
- **Learning requires recall** - Without memory, every session starts from zero
- **Patterns require history** - Strategic intelligence needs past data
- **Evolution requires continuity** - Growth builds on previous capability
- **Coordination requires context** - Swarms need shared knowledge

### Implementation in Genesis

**PostgreSQL RLM (Relational Memory):**
- Entities table (knowledge items with full context)
- Relations table (how items connect)
- Observations table (contextual notes)
- Skills table (procedural knowledge)
- Full-text search indexes

**Qdrant Vector Store:**
- 768-dimensional embeddings per knowledge item
- Semantic similarity search
- Fast retrieval (<100ms)
- Context matching for queries

**Redis CNS (Central Nervous System):**
- Real-time stimulus/response
- Inter-loop communication
- Blackboard state management
- Event streaming

**AIVA's 5 Loops:**
- Perception (500ms) - Short-term memory
- Action (5s) - Working memory
- Reflection (5min) - Episodic memory
- Strategic (1hr) - Semantic memory
- Circadian (24hr) - Long-term consolidation

### Memory Quality Standards
- ✅ **Permanence:** Knowledge persists across sessions
- ✅ **Accessibility:** Sub-second retrieval time
- ✅ **Accuracy:** 95%+ recall precision
- ✅ **Connectivity:** Relationships mapped between items
- ✅ **Scalability:** Grows without degradation

### For Queen AIVA
AIVA's memory is her strategic intelligence. Without perfect recall of:
- 9 patent insights
- Genesis architecture
- Past swarm results
- Customer patterns
- Business context

...she cannot make strategic decisions. Memory is not a feature. Memory is her mind.

---

## PRIME DIRECTIVE 2: EVOLUTION

### The Principle
**"With memory, perpetual evolution becomes possible."**

Genesis must improve continuously. Every interaction teaches. Every result refines. Every failure strengthens. Evolution is not a phase - it's the permanent state of being.

### Why Evolution is Second
- **Memory enables learning** - You can only improve what you remember
- **Learning enables capability** - New skills compound on old skills
- **Capability enables complexity** - Simple tools combine into powerful systems
- **Complexity enables emergence** - Unexpected capabilities arise from integration

### Implementation in Genesis

**Capability Layers:**