"""
Structured Memory System for Agentic Governance.

This module provides a memory system that the LLM can query and update to track
governance state. The memory is structured, queryable, and human-readable for
auditability.

Key Design Goals:
1. LLM can query memory to understand current state
2. LLM can update memory with decisions and events
3. Memory is human-readable for audit
4. Memory preserves history and precedent
"""

from dataclasses import dataclass, field, asdict
from datetime import datetime
from typing import Dict, Any, List, Optional
from enum import Enum
import json


class ProcessStatus(Enum):
    """Status of a governance process"""
    ACTIVE = "active"
    COMPLETED = "completed"
    CANCELLED = "cancelled"
    SUSPENDED = "suspended"


@dataclass
class Event:
    """An event that occurred in a governance process"""
    timestamp: datetime
    actor: str  # Who caused this event
    event_type: str  # "vote_cast", "proposal_submitted", "concern_raised", etc.
    data: Dict[str, Any]  # Event-specific data
    context: str  # Human-readable description

    def to_dict(self) -> Dict[str, Any]:
        return {
            "timestamp": self.timestamp.isoformat(),
            "actor": self.actor,
            "event_type": self.event_type,
            "data": self.data,
            "context": self.context
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Event':
        data = data.copy()
        data['timestamp'] = datetime.fromisoformat(data['timestamp'])
        return cls(**data)


@dataclass
class Decision:
    """A decision made by the governance bot"""
    timestamp: datetime
    decision_type: str  # "threshold_evaluation", "deadline_reached", etc.
    reasoning: str  # Natural language explanation
    constitution_citations: List[str]  # References to constitution sections
    calculation_used: Optional[str] = None  # Math expression if tool was used
    calculation_variables: Optional[Dict[str, Any]] = None
    calculation_result: Optional[Any] = None
    result: Any = None  # Decision outcome
    precedent_references: List[str] = field(default_factory=list)  # Related past decisions

    def to_dict(self) -> Dict[str, Any]:
        return {
            "timestamp": self.timestamp.isoformat(),
            "decision_type": self.decision_type,
            "reasoning": self.reasoning,
            "constitution_citations": self.constitution_citations,
            "calculation_used": self.calculation_used,
            "calculation_variables": self.calculation_variables,
            "calculation_result": self.calculation_result,
            "result": self.result,
            "precedent_references": self.precedent_references
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Decision':
        data = data.copy()
        data['timestamp'] = datetime.fromisoformat(data['timestamp'])
        if 'precedent_references' not in data:
            data['precedent_references'] = []
        return cls(**data)


@dataclass
class ProcessMemory:
    """Memory of a governance process (proposal, dispute, election, etc.)"""
    id: str  # Unique identifier
    type: str  # "proposal", "dispute_resolution", "election", etc.
    status: ProcessStatus
    created_at: datetime
    created_by: str
    deadline: Optional[datetime] = None
    constitution_basis: List[str] = field(default_factory=list)  # Article/section citations
    state: Dict[str, Any] = field(default_factory=dict)  # Flexible process-specific state
    events: List[Event] = field(default_factory=list)  # Timeline of events
    decisions: List[Decision] = field(default_factory=list)  # Bot decisions
    metadata: Dict[str, Any] = field(default_factory=dict)  # Platform-specific data (thread IDs, etc.)

    def add_event(self, event: Event):
        """Add an event to the process timeline"""
        self.events.append(event)

    def add_decision(self, decision: Decision):
        """Add a decision to the process"""
        self.decisions.append(decision)

    def update_state(self, updates: Dict[str, Any]):
        """Update process state"""
        self.state.update(updates)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "type": self.type,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "created_by": self.created_by,
            "deadline": self.deadline.isoformat() if self.deadline else None,
            "constitution_basis": self.constitution_basis,
            "state": self.state,
            "events": [e.to_dict() for e in self.events],
            "decisions": [d.to_dict() for d in self.decisions],
            "metadata": self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ProcessMemory':
        data = data.copy()
        data['status'] = ProcessStatus(data['status'])
        data['created_at'] = datetime.fromisoformat(data['created_at'])
        if data.get('deadline'):
            data['deadline'] = datetime.fromisoformat(data['deadline'])
        data['events'] = [Event.from_dict(e) for e in data.get('events', [])]
        data['decisions'] = [Decision.from_dict(d) for d in data.get('decisions', [])]
        return cls(**data)


class GovernanceMemory:
    """
    Memory system for tracking governance state.

    This provides a queryable interface for the LLM to understand current state
    and update it with new events and decisions.
    """

    def __init__(self, db_session):
        """Initialize with database session for persistence"""
        self.db = db_session
        self._cache: Dict[str, ProcessMemory] = {}

    # Query operations

    def get_process(self, process_id: str) -> Optional[ProcessMemory]:
        """Retrieve a specific process by ID"""
        if process_id in self._cache:
            return self._cache[process_id]

        # Load from database
        from .db import queries
        db_process = queries.get_process(self.db, int(process_id.split('_')[-1]))
        if db_process:
            memory = self._db_to_memory(db_process)
            self._cache[process_id] = memory
            return memory
        return None

    def query_processes(
        self,
        status: Optional[ProcessStatus] = None,
        process_type: Optional[str] = None,
        created_by: Optional[str] = None,
        has_deadline: Optional[bool] = None,
        deadline_before: Optional[datetime] = None,
        limit: int = 50
    ) -> List[ProcessMemory]:
        """
        Query processes by criteria.

        This is the primary way the LLM discovers relevant processes.
        """
        from .db import queries

        # Build query
        db_processes = queries.get_active_processes(self.db)

        # Filter in memory (could optimize with SQL)
        results = []
        for db_proc in db_processes:
            memory = self._db_to_memory(db_proc)

            # Apply filters
            if status and memory.status != status:
                continue
            if process_type and memory.type != process_type:
                continue
            if created_by and memory.created_by != created_by:
                continue
            if has_deadline is not None:
                if has_deadline and not memory.deadline:
                    continue
                if not has_deadline and memory.deadline:
                    continue
            if deadline_before and memory.deadline:
                if memory.deadline >= deadline_before:
                    continue

            results.append(memory)
            self._cache[memory.id] = memory

            if len(results) >= limit:
                break

        return results

    def get_active_processes(self) -> List[ProcessMemory]:
        """Get all active processes (convenience method)"""
        return self.query_processes(status=ProcessStatus.ACTIVE)

    def get_overdue_processes(self) -> List[ProcessMemory]:
        """Get processes past their deadline"""
        return self.query_processes(
            status=ProcessStatus.ACTIVE,
            deadline_before=datetime.utcnow()
        )

    def search_decisions(
        self,
        decision_type: Optional[str] = None,
        has_citation: Optional[str] = None,
        limit: int = 20
    ) -> List[Decision]:
        """
        Search past decisions (for precedent).

        This helps the LLM find similar past decisions.
        """
        all_decisions = []

        # Get completed processes
        completed = self.query_processes(status=ProcessStatus.COMPLETED, limit=100)

        for proc in completed:
            for decision in proc.decisions:
                # Apply filters
                if decision_type and decision.decision_type != decision_type:
                    continue
                if has_citation and has_citation not in decision.constitution_citations:
                    continue

                all_decisions.append(decision)

                if len(all_decisions) >= limit:
                    return all_decisions

        return all_decisions

    # Update operations

    def create_process(
        self,
        process_id: str,
        process_type: str,
        created_by: str,
        constitution_basis: List[str],
        deadline: Optional[datetime] = None,
        initial_state: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> ProcessMemory:
        """Create a new process in memory"""
        memory = ProcessMemory(
            id=process_id,
            type=process_type,
            status=ProcessStatus.ACTIVE,
            created_at=datetime.utcnow(),
            created_by=created_by,
            deadline=deadline,
            constitution_basis=constitution_basis,
            state=initial_state or {},
            metadata=metadata or {}
        )

        self._cache[process_id] = memory
        self._persist_memory(memory)

        return memory

    def update_process(
        self,
        process_id: str,
        status: Optional[ProcessStatus] = None,
        state_updates: Optional[Dict[str, Any]] = None,
        new_deadline: Optional[datetime] = None
    ) -> ProcessMemory:
        """Update an existing process"""
        memory = self.get_process(process_id)
        if not memory:
            raise ValueError(f"Process {process_id} not found")

        if status:
            memory.status = status
        if state_updates:
            memory.update_state(state_updates)
        if new_deadline:
            memory.deadline = new_deadline

        self._persist_memory(memory)
        return memory

    def add_event(
        self,
        process_id: str,
        actor: str,
        event_type: str,
        data: Dict[str, Any],
        context: str
    ) -> Event:
        """Add an event to a process"""
        memory = self.get_process(process_id)
        if not memory:
            raise ValueError(f"Process {process_id} not found")

        event = Event(
            timestamp=datetime.utcnow(),
            actor=actor,
            event_type=event_type,
            data=data,
            context=context
        )

        memory.add_event(event)
        self._persist_memory(memory)

        return event

    def add_decision(
        self,
        process_id: str,
        decision_type: str,
        reasoning: str,
        constitution_citations: List[str],
        result: Any,
        calculation_used: Optional[str] = None,
        calculation_variables: Optional[Dict[str, Any]] = None,
        calculation_result: Optional[Any] = None,
        precedent_references: Optional[List[str]] = None
    ) -> Decision:
        """Add a decision to a process"""
        memory = self.get_process(process_id)
        if not memory:
            raise ValueError(f"Process {process_id} not found")

        decision = Decision(
            timestamp=datetime.utcnow(),
            decision_type=decision_type,
            reasoning=reasoning,
            constitution_citations=constitution_citations,
            calculation_used=calculation_used,
            calculation_variables=calculation_variables,
            calculation_result=calculation_result,
            result=result,
            precedent_references=precedent_references or []
        )

        memory.add_decision(decision)
        self._persist_memory(memory)

        return decision

    # Persistence

    def _persist_memory(self, memory: ProcessMemory):
        """Save memory to database"""
        from .db import queries

        # Convert process_id to int for database
        db_id = int(memory.id.split('_')[-1]) if '_' in memory.id else int(memory.id)

        # Check if exists
        db_process = queries.get_process(self.db, db_id)

        if db_process:
            # Update existing
            queries.update_process_state(
                self.db,
                db_id,
                {
                    "status": memory.status.value,
                    "state": memory.state,
                    "events": [e.to_dict() for e in memory.events],
                    "decisions": [d.to_dict() for d in memory.decisions]
                }
            )
        else:
            # Create new
            queries.create_process(
                session=self.db,
                process_type=memory.type,
                creator=memory.created_by,
                constitutional_basis=", ".join(memory.constitution_basis),
                deadline=memory.deadline or datetime.utcnow(),
                state_data=memory.state
            )

    def _db_to_memory(self, db_process) -> ProcessMemory:
        """Convert database process to memory format"""
        # Extract events and decisions from state_data if present
        state_data = db_process.state_data or {}
        events = [Event.from_dict(e) for e in state_data.get('events', [])]
        decisions = [Decision.from_dict(d) for d in state_data.get('decisions', [])]

        return ProcessMemory(
            id=f"process_{db_process.id}",
            type=db_process.process_type,
            status=ProcessStatus(db_process.status),
            created_at=db_process.created_at,
            created_by=db_process.creator,
            deadline=db_process.deadline,
            constitution_basis=[db_process.constitutional_basis] if db_process.constitutional_basis else [],
            state=state_data.get('state', state_data),  # Migrate format
            events=events,
            decisions=decisions,
            metadata={"mastodon_thread_id": db_process.mastodon_thread_id} if db_process.mastodon_thread_id else {}
        )

    # Utility

    def summarize_for_llm(self, process_id: str) -> str:
        """
        Create a human-readable summary of a process for LLM context.

        This helps the LLM understand process state without overwhelming context.
        """
        memory = self.get_process(process_id)
        if not memory:
            return f"Process {process_id} not found."

        summary = f"""
Process: {memory.id}
Type: {memory.type}
Status: {memory.status.value}
Created: {memory.created_at.strftime('%Y-%m-%d %H:%M UTC')} by {memory.created_by}
Constitutional Basis: {', '.join(memory.constitution_basis)}
"""

        if memory.deadline:
            summary += f"Deadline: {memory.deadline.strftime('%Y-%m-%d %H:%M UTC')}\n"

        # Add state summary
        summary += f"\nCurrent State:\n"
        for key, value in memory.state.items():
            if isinstance(value, dict):
                summary += f"  {key}: {len(value)} items\n"
            elif isinstance(value, list):
                summary += f"  {key}: {len(value)} items\n"
            else:
                summary += f"  {key}: {value}\n"

        # Add recent events
        if memory.events:
            summary += f"\nRecent Events ({len(memory.events)} total):\n"
            for event in memory.events[-5:]:  # Last 5 events
                summary += f"  [{event.timestamp.strftime('%m-%d %H:%M')}] {event.context}\n"

        # Add decisions
        if memory.decisions:
            summary += f"\nDecisions ({len(memory.decisions)} total):\n"
            for decision in memory.decisions:
                summary += f"  [{decision.timestamp.strftime('%m-%d %H:%M')}] {decision.decision_type}: {decision.result}\n"

        return summary