import yaml
import os

def load_persona(path):
    """
    Load a LitDigitalTwin persona from YAML file.
    Ensures required fields and initializes missing defaults.
    """
    if not os.path.exists(path):
        raise FileNotFoundError(f"Persona file not found: {path}")
    
    with open(path, "r", encoding="utf-8") as f:
        persona = yaml.safe_load(f)

    # Required top-level keys
    required_keys = [
        "persona_name",
        "age",
        "role",
        "system_prompt",
        "facts"
    ]
    for key in required_keys:
        if key not in persona:
            raise ValueError(f"Missing required key in persona: {key}")

    # Ensure facts is a list
    if not isinstance(persona.get("facts"), list):
        print("Warning: 'facts' should be a list. Converting.")
        facts = persona.get("facts", [])
        if isinstance(facts, dict):
            persona["facts"] = list(facts.values())
        else:
            persona["facts"] = []

    # Initialize default_state if missing
    if "default_state" not in persona:
        persona["default_state"] = {}

    state = persona["default_state"]
    default_state_keys = {
        "anxiety": 0.5,
        "trust": 0.5,
        "openness": 0.5,
        "mode": "baseline",
        "emotional_memory": []
    }

    for key, default in default_state_keys.items():
        if key not in state:
            print(f"Warning: Missing state key '{key}' in persona. Using default value.")
            state[key] = default

    # Ensure tone_guidance exists and has at least 'baseline'
    if "tone_guidance" not in persona:
        persona["tone_guidance"] = {
            "baseline": {
                "voice": "Natural and authentic",
                "example": "I'm doing okay today, thanks for asking."
            }
        }

    return persona


def validate_persona(persona):
    """
    Validate that a persona has all necessary components for simulation.
    Returns (is_valid, error_message)
    """
    if not persona.get("persona_name"):
        return False, "Persona must have a name"

    if not persona.get("system_prompt"):
        return False, "Persona must have a system_prompt"

    if not isinstance(persona.get("facts"), list):
        return False, "Persona 'facts' must be a list"

    state = persona.get("default_state", {})
    for key in ["anxiety", "trust", "openness"]:
        value = state.get(key)
        if value is None:
            return False, f"default_state missing required key: {key}"
        if not isinstance(value, (int, float)):
            return False, f"default_state.{key} must be numeric"
        if not 0 <= value <= 1:
            return False, f"default_state.{key} must be between 0 and 1"

    return True, "Persona is valid"


def save_persona(persona, path):
    """
    Save a persona to YAML file.
    """
    with open(path, "w", encoding="utf-8") as f:
        yaml.dump(persona, f, sort_keys=False, default_flow_style=False)
    return path


def create_default_persona(name, age, role):
    """
    Create a basic LitDigitalTwin persona template.
    """
    persona = {
        "persona_name": name,
        "age": age,
        "role": role,
        "system_prompt": f"You are {name}, a {age}-year-old {role}. Respond naturally and stay in character.",
        "facts": [
            f"{name} is {age} years old.",
            f"{name} works as a {role}."
        ],
        "tone_guidance": {
            "baseline": {
                "voice": "Natural and authentic",
                "example": "I'm doing okay today, thanks for asking."
            }
        },
        "default_state": {
            "anxiety": 0.5,
            "trust": 0.5,
            "openness": 0.5,
            "mode": "baseline",
            "emotional_memory": []
        }
    }
    return persona


def list_available_personas(persona_dir="./personas"):
    """
    List all available persona files.
    """
    if not os.path.exists(persona_dir):
        return []

    personas = []
    for filename in os.listdir(persona_dir):
        if filename.endswith(".yml") or filename.endswith(".yaml"):
            path = os.path.join(persona_dir, filename)
            try:
                persona = load_persona(path)
                personas.append({
                    "filename": filename,
                    "name": persona.get("persona_name", "Unknown"),
                    "age": persona.get("age", ""),
                    "role": persona.get("role", "")
                })
            except Exception as e:
                print(f"Error loading {filename}: {e}")
    return personas