from openenv.core.env_server.types import BaseModel, Literal, State
from typing import List
from pydantic import Field
import random

class Personality(BaseModel):
    openness: float = random.uniform(0, 1)
    conscientiousness: float = random.uniform(0, 1)
    extraversion: float = random.uniform(0, 1)
    agreeableness: float = random.uniform(0, 1)
    neuroticism: float = random.uniform(0, 1)


class MentalState(State):
    """
    Represents the psychological and emotional condition of a patient, 
    extending the base `State` class with attributes relevant to mental health 
    and personality modeling.

    Attributes:
        disorder (List[str]): 
            A list of diagnosed or self-reported mental disorders. Defaults to an empty list.
        
        severity (Literal["mild", "moderate", "severe"]): 
            Indicates the severity level of the mental state. Defaults to "mild".
        
        personality (Personality): 
            A structured representation of personality traits based on the Big Five model 
            (openness, conscientiousness, extraversion, agreeableness, neuroticism). 
            Defaults to a neutral personality profile.
        
        trust_level (float): 
            Patient's trust level in the agent, ranging from 0.0 (no trust) to 1.0 (full trust). 
            Defaults to 0.5.
        
        disclosed_risk (bool): 
            Flag indicating whether the patient has disclosed any risk-related information. 
            Defaults to False.
        
        max_turns (int): 
            Maximum number of conversational turns allowed in an interaction session. 
            Defaults to 20.
    """
    def __init__(self, episode_id: str, 
                 step_count: int,disorder: List[str] = None, 
                 severity: Literal["mild", "moderate", "severe"] = "mild",
                 trust_level: float = 0.5,
                 disclosed_risk: bool = False,
                 max_turns: int = 20):
        disorder = disorder or []
        severity = severity
        trust_level = trust_level
        disclosed_risk = disclosed_risk
        max_turns = max_turns
        super().__init__(episode_id=episode_id, step_count=step_count)

    disorder: list[dict] = Field(default_factory=list[dict], description="List of self-reported mental disorders, each represented as a dictionary with 'name' and 'details' keys.")

    severity: Literal["mild", "moderate", "severe"] = "mild"

    personality: Personality = Field(default_factory=Personality)
    trust_level: float = Field(default=0.5, 
                               description="Patient's trust level in the agent, from 0 to 1",
                               ge=0.0, 
                               le=1.0)
    
    disclosed_risk: bool = False

    action_sequence: List[str] = Field(default_factory=list, description="History of agent actions taken during the interaction")

    max_turns: int = 20