"""
Attention Allocation Receipts (X-BAM Task Prioritization)

Tracks X-BAM (eXplanation-Based Attention Management) task prioritization
and resource allocation decisions.

Key Features:
- χ-based priority computation (MemoryGraph integration)
- Resource allocation tracking (CPU, memory, tokens, duration)
- Evidence-based justification (receipt chains)
- Fairness metrics (group parity, individual fairness)
- Guardian integration (policy compliance)
- Deadline and dependency tracking

Reference: WORK_BREAKDOWN_PARALLEL_TEAMS_2025-11-06.md (Workstream D.2.3)
"""

from datetime import datetime, timedelta, timezone
from typing import Dict, List, Literal, Optional
from pydantic import BaseModel, Field, field_validator, model_validator


class ResourceAllocation(BaseModel):
    """Resource limits for a task."""

    cpu_cores: Optional[float] = Field(None, ge=0, description="CPU cores allocated")
    memory_mb: Optional[float] = Field(None, ge=0, description="Memory in MB")
    tokens_allocated: Optional[int] = Field(None, ge=0, description="Token budget")
    max_duration_seconds: Optional[int] = Field(None, ge=0, description="Max execution time")


class PolicyConstraints(BaseModel):
    """Policy constraints affecting allocation."""

    max_priority: Optional[float] = Field(None, description="Maximum allowed priority")
    min_resources: Optional[Dict] = Field(None, description="Minimum resource requirements")
    fairness_group: Optional[str] = Field(None, description="Fairness group identifier")


class FairnessMetrics(BaseModel):
    """Fairness metrics for allocation."""

    group_parity: Optional[float] = Field(
        None,
        ge=0.0,
        le=1.0,
        description="Group parity score (1.0 = perfect parity)"
    )
    individual_fairness: Optional[float] = Field(
        None,
        ge=0.0,
        le=1.0,
        description="Individual fairness score (1.0 = perfectly fair)"
    )


class AttentionAllocationReceipt(BaseModel):
    """
    Receipt documenting X-BAM task prioritization and resource allocation.

    Tracks how tasks are prioritized, resources allocated, and provides
    evidence-based justification for allocation decisions.
    """

    schema: Literal["receipt.attention_allocation.v1"] = "receipt.attention_allocation.v1"

    receipt_id: str = Field(
        ...,
        pattern=r"^attn-[0-9A-Za-z_-]+$",
        description="Unique attention allocation receipt ID"
    )

    task_id: str = Field(..., description="Unique task identifier")

    agent_id: str = Field(..., description="Agent performing the prioritization")

    session_id: Optional[str] = Field(None, description="Session or context ID")

    task_type: Literal["execution", "review", "coordination", "planning", "analysis", "synthesis"] = Field(
        ...,
        description="Type of task being prioritized"
    )

    priority_score: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Computed priority score (0.0 = lowest, 1.0 = highest)"
    )

    allocation_method: Literal["chi-based", "deadline-driven", "policy-mandated", "user-specified", "learned"] = Field(
        ...,
        description="Method used to determine allocation"
    )

    chi_score: Optional[float] = Field(
        None,
        ge=0.0,
        le=1.0,
        description="χ (chi) heuristic score from MemoryGraph"
    )

    delta_chi: Optional[float] = Field(
        None,
        description="Change in χ score (can be negative)"
    )

    resource_allocation: Optional[ResourceAllocation] = Field(
        None,
        description="Resource limits for this task"
    )

    evidence_refs: Optional[List[str]] = Field(
        default_factory=list,
        description="Receipt IDs or memory node IDs that informed prioritization"
    )

    justification: Optional[str] = Field(
        None,
        description="Human-readable explanation for priority assignment"
    )

    owner: Optional[str] = Field(None, description="Agent or entity owning the task")

    status: Literal["allocated", "executing", "completed", "cancelled", "expired"] = Field(
        ...,
        description="Task allocation status"
    )

    created_at: Optional[str] = Field(None, description="Task creation timestamp")

    allocated_at: Optional[str] = Field(None, description="Allocation decision timestamp")

    expires_at: Optional[str] = Field(None, description="Expiration timestamp for allocation")

    deadline: Optional[str] = Field(None, description="Task deadline (if applicable)")

    dependencies: Optional[List[str]] = Field(
        default_factory=list,
        description="Task IDs this task depends on"
    )

    blocking: Optional[List[str]] = Field(
        default_factory=list,
        description="Task IDs that depend on this task"
    )

    guardian_verdict: Optional[Literal["pass", "refusal", "error"]] = Field(
        None,
        description="Guardian's verdict on allocation decision"
    )

    guardian_receipt_id: Optional[str] = Field(
        None,
        description="Reference to Guardian's receipt"
    )

    dignity_score: Optional[float] = Field(
        None,
        ge=0.0,
        le=1.0,
        description="EVHD dignity score for allocation decision"
    )

    policy_constraints: Optional[PolicyConstraints] = Field(
        None,
        description="Policy constraints that affected allocation"
    )

    fairness_metrics: Optional[FairnessMetrics] = Field(
        None,
        description="Fairness metrics for allocation"
    )

    timestamp: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat(),
        description="ISO 8601 timestamp"
    )

    signature: str = Field(
        ...,
        pattern=r"^ed25519:.+$",
        description="Ed25519 signature"
    )

    metadata: Optional[Dict] = Field(
        None,
        description="Optional metadata (scheduler, queue_depth, wait_time_ms, allocation_latency_ms)"
    )

    @model_validator(mode='after')
    def validate_chi_consistency(self) -> 'AttentionAllocationReceipt':
        """Ensure chi_score and delta_chi are consistent with allocation_method."""
        if self.allocation_method == "chi-based" and self.chi_score is None:
            raise ValueError("chi-based allocation must provide chi_score")
        return self

    @model_validator(mode='after')
    def validate_deadline_consistency(self) -> 'AttentionAllocationReceipt':
        """Ensure deadline-driven allocation has a deadline."""
        if self.allocation_method == "deadline-driven" and not self.deadline:
            raise ValueError("deadline-driven allocation must provide deadline")
        return self

    @model_validator(mode='after')
    def validate_expires_after_created(self) -> 'AttentionAllocationReceipt':
        """Ensure expires_at is after created_at."""
        if self.created_at and self.expires_at:
            created = datetime.fromisoformat(self.created_at.replace('Z', '+00:00'))
            expires = datetime.fromisoformat(self.expires_at.replace('Z', '+00:00'))

            if expires < created:
                raise ValueError(
                    f"expires_at ({self.expires_at}) must be after created_at ({self.created_at})"
                )
        return self


def create_chi_based_allocation(
    task_id: str,
    agent_id: str,
    task_type: str,
    chi_score: float,
    delta_chi: float,
    justification: str,
    resource_allocation: Optional[ResourceAllocation] = None,
    evidence_refs: Optional[List[str]] = None,
    ttl_seconds: int = 3600,
    signature: str = "ed25519:stub"
) -> AttentionAllocationReceipt:
    """
    Create an attention allocation receipt using χ-based prioritization.

    Args:
        task_id: Unique task identifier
        agent_id: Agent performing prioritization
        task_type: Type of task
        chi_score: χ (chi) heuristic score
        delta_chi: Change in χ score
        justification: Explanation for priority
        resource_allocation: Resource limits
        evidence_refs: Evidence that informed decision
        ttl_seconds: Time-to-live for allocation
        signature: Ed25519 signature

    Returns:
        AttentionAllocationReceipt with chi-based allocation
    """
    timestamp = datetime.now(timezone.utc)
    expires_at = timestamp + timedelta(seconds=ttl_seconds)

    # Compute priority from chi_score
    priority_score = min(1.0, max(0.0, chi_score))

    # Sanitize task_id for receipt_id (replace invalid characters)
    sanitized_task_id = task_id.replace("/", "-").replace(" ", "-").replace(".", "-")
    receipt_id = f"attn-chi-{sanitized_task_id}-{int(timestamp.timestamp())}"

    return AttentionAllocationReceipt(
        receipt_id=receipt_id,
        task_id=task_id,
        agent_id=agent_id,
        task_type=task_type,
        priority_score=priority_score,
        allocation_method="chi-based",
        chi_score=chi_score,
        delta_chi=delta_chi,
        resource_allocation=resource_allocation,
        evidence_refs=evidence_refs or [],
        justification=justification,
        status="allocated",
        created_at=timestamp.isoformat(),
        allocated_at=timestamp.isoformat(),
        expires_at=expires_at.isoformat(),
        timestamp=timestamp.isoformat(),
        signature=signature
    )


def create_deadline_driven_allocation(
    task_id: str,
    agent_id: str,
    task_type: str,
    deadline: datetime,
    resource_allocation: Optional[ResourceAllocation] = None,
    justification: Optional[str] = None,
    signature: str = "ed25519:stub"
) -> AttentionAllocationReceipt:
    """
    Create an attention allocation receipt using deadline-driven prioritization.

    Args:
        task_id: Unique task identifier
        agent_id: Agent performing prioritization
        task_type: Type of task
        deadline: Task deadline
        resource_allocation: Resource limits
        justification: Explanation for priority
        signature: Ed25519 signature

    Returns:
        AttentionAllocationReceipt with deadline-driven allocation
    """
    timestamp = datetime.now(timezone.utc)

    # Compute priority based on time until deadline
    time_until_deadline = (deadline - timestamp).total_seconds()
    if time_until_deadline <= 0:
        priority_score = 1.0  # Overdue
    elif time_until_deadline < 3600:  # < 1 hour
        priority_score = 0.9
    elif time_until_deadline < 86400:  # < 1 day
        priority_score = 0.7
    else:
        priority_score = 0.5

    # Sanitize task_id for receipt_id (replace invalid characters)
    sanitized_task_id = task_id.replace("/", "-").replace(" ", "-").replace(".", "-")
    receipt_id = f"attn-deadline-{sanitized_task_id}-{int(timestamp.timestamp())}"

    return AttentionAllocationReceipt(
        receipt_id=receipt_id,
        task_id=task_id,
        agent_id=agent_id,
        task_type=task_type,
        priority_score=priority_score,
        allocation_method="deadline-driven",
        resource_allocation=resource_allocation,
        justification=justification or f"Deadline in {time_until_deadline/3600:.1f} hours",
        status="allocated",
        created_at=timestamp.isoformat(),
        allocated_at=timestamp.isoformat(),
        deadline=deadline.isoformat(),
        timestamp=timestamp.isoformat(),
        signature=signature
    )


# IdempotentWriter integration
def write_attention_allocation_receipt(
    receipt: AttentionAllocationReceipt,
    writer: Optional['IdempotentWriter'] = None,
    force_write: bool = False
) -> Dict:
    """
    Write attention allocation receipt with deduplication.

    Args:
        receipt: AttentionAllocationReceipt to write
        writer: IdempotentWriter instance (creates default if None)
        force_write: If True, bypass deduplication

    Returns:
        Dict with receipt_id, is_duplicate, content_hash, path
    """
    from receipts.idempotent_writer import IdempotentWriter
    from pathlib import Path

    if writer is None:
        writer = IdempotentWriter(base_dir=Path("receipts"))

    # Convert to dict
    receipt_dict = receipt.model_dump(mode='json', by_alias=True)

    # Write with deduplication
    return writer.write_receipt(receipt_dict, force_write=force_write)


# Example usage
if __name__ == "__main__":
    print("AttentionAllocationReceipt Demo")
    print("=" * 70)

    # Example 1: χ-based allocation
    print("\n1. χ-based allocation (high χ score):")
    chi_allocation = create_chi_based_allocation(
        task_id="task-123",
        agent_id="agent-001",
        task_type="execution",
        chi_score=0.85,
        delta_chi=0.15,
        justification="High relevance + novelty; aligns with dignity polyhedron",
        resource_allocation=ResourceAllocation(
            cpu_cores=2.0,
            memory_mb=4096,
            tokens_allocated=8000,
            max_duration_seconds=300
        ),
        evidence_refs=["memory-node-456", "receipt-prophecy-789"]
    )
    print(f"   Receipt ID: {chi_allocation.receipt_id}")
    print(f"   Priority: {chi_allocation.priority_score:.2f}")
    print(f"   χ score: {chi_allocation.chi_score:.2f}")
    print(f"   Δχ: {chi_allocation.delta_chi:+.2f}")

    # Example 2: Deadline-driven allocation
    print("\n2. Deadline-driven allocation (urgent):")
    deadline = datetime.now(timezone.utc) + timedelta(minutes=30)
    deadline_allocation = create_deadline_driven_allocation(
        task_id="task-456",
        agent_id="agent-002",
        task_type="review",
        deadline=deadline,
        resource_allocation=ResourceAllocation(
            cpu_cores=1.0,
            memory_mb=2048,
            tokens_allocated=4000
        )
    )
    print(f"   Receipt ID: {deadline_allocation.receipt_id}")
    print(f"   Priority: {deadline_allocation.priority_score:.2f}")
    print(f"   Deadline: {deadline_allocation.deadline}")
    print(f"   Justification: {deadline_allocation.justification}")

    # Example 3: With Guardian verdict
    print("\n3. Allocation with Guardian approval:")
    guarded_allocation = create_chi_based_allocation(
        task_id="task-789",
        agent_id="agent-003",
        task_type="coordination",
        chi_score=0.72,
        delta_chi=0.05,
        justification="Moderate priority; dignity score acceptable"
    )
    guarded_allocation.guardian_verdict = "pass"
    guarded_allocation.guardian_receipt_id = "guardian-receipt-123"
    guarded_allocation.dignity_score = 0.78
    print(f"   Guardian: {guarded_allocation.guardian_verdict}")
    print(f"   Dignity: {guarded_allocation.dignity_score:.2f}")

    # Test serialization
    print("\n4. Serialization test:")
    json_data = chi_allocation.model_dump(mode='json', by_alias=True)
    print(f"   JSON keys: {list(json_data.keys())[:5]}...")
    print(f"   Schema: {json_data['schema']}")

    print("\n✅ All examples working!")
