from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field


class DataRecord(BaseModel):
    """A single data record in the enterprise data lake."""

    record_id: str
    system: Literal["postgres", "s3", "slack", "email"]
    data_type: Literal["PII", "anonymized", "legal_hold", "derived"]
    user_id: str
    dependencies: List[str] = Field(default_factory=list)
    legal_hold: bool
    age_days: int
    deletion_cost: float
    status: Literal["active", "deleted", "anonymized"] = "active"


class EraseRequest(BaseModel):
    """The incoming GDPR erasure request."""

    request_id: str
    user_id: str
    deadline_steps: int
    regulation: Literal["GDPR", "CCPA"]
    submitted_at_step: int = 0


class EraseOpsObservation(BaseModel):
    """Full observation returned to the agent at every step."""

    records: List[DataRecord]
    request: EraseRequest
    step: int
    step_budget: int
    cumulative_penalty: float
    last_action_result: str
    pii_remaining: int
    legal_hold_violations: int
    dependency_violations: int
    contamination_violations: int
    requests_completed: int
    deadline_pressure: float
    escalation_queue_size: int
    escalation_queue_max_age: int
    blocked_dependency_count: int
    blocked_dependency_max_age: int


class EraseOpsAction(BaseModel):
    """Action the agent takes each step."""

    action_type: Literal["inspect", "delete", "anonymize", "escalate", "noop"]
    record_id: Optional[str] = None
    justification: Optional[str] = None


class EraseOpsReward(BaseModel):
    """Decomposed reward for transparency."""

    total: float
    deletion_reward: float
    step_penalty: float
    legal_hold_penalty: float
    dependency_penalty: float
    contamination_penalty: float
    over_deletion_penalty: float
    completion_bonus: float


class TaskConfig(BaseModel):
    """In-memory deterministic task configuration returned by get_task."""

    task_name: str
    step_budget: int
    request: EraseRequest
    records: List[DataRecord]
    observation_order: List[str]
    metadata: Dict[str, Any] = Field(default_factory=dict)