from enum import Enum
from typing import Optional
from pydantic import BaseModel


class RagStatus(str, Enum):
    """RAG status for agent reviews."""
    RED = "Red"
    AMBER = "Amber"
    GREEN = "Green"
    ERROR = "Error"


class OverallStatus(str, Enum):
    """Overall status for the proof review."""
    PASSED = "Passed"
    FAILED = "Failed"
    ANALYSIS_ERROR = "Analysis Error"
    REQUIRES_MANUAL_LEGAL_REVIEW = "Requires Manual Legal Review"


class SubReview(BaseModel):
    """Individual agent review result."""
    ragStatus: RagStatus
    feedback: str
    issues: list[str]
    isFinancialPromotion: Optional[bool] = None
    financialPromotionReason: Optional[str] = None

    class Config:
        use_enum_values = True


class AgentReview(BaseModel):
    """Complete review from all agents."""
    legalAgentReview: SubReview
    brandAgentReview: SubReview
    toneAgentReview: SubReview
    channelAgentReview: SubReview
    leadAgentSummary: str
    overallStatus: OverallStatus
    financialPromotionReason: Optional[str] = None

    class Config:
        use_enum_values = True


# WebSocket message types
class AnalyzeRequest(BaseModel):
    """Request to analyze a proof via WebSocket."""
    type: str  # Should be "analyze"
    file_data: str  # Base64 encoded file
    file_type: str  # MIME type
    is_wip: bool = False


class AgentStartedMessage(BaseModel):
    """Message sent when an agent starts processing."""
    type: str = "agent_started"
    agent_name: str


class AgentCompletedMessage(BaseModel):
    """Message sent when an agent completes."""
    type: str = "agent_completed"
    agent_name: str
    review: SubReview

    class Config:
        use_enum_values = True


class SummaryMessage(BaseModel):
    """Message sent when summary is ready."""
    type: str = "summary"
    lead_agent_summary: str
    overall_status: OverallStatus
    financial_promotion_reason: Optional[str] = None

    class Config:
        use_enum_values = True


class CompleteMessage(BaseModel):
    """Message sent when analysis is complete."""
    type: str = "complete"
    result: AgentReview

    class Config:
        use_enum_values = True


class ErrorMessage(BaseModel):
    """Message sent when an error occurs."""
    type: str = "error"
    message: str
