"""
ActionUnitSet: Container for loading and saving sets of Action Units.
"""
import json
import os
from datetime import datetime
from typing import List, Optional

from core.context import ActionUnit


class ActionUnitSet:
    """Container for a set of Action Units with metadata."""
    
    def __init__(self,
                 agent_id: str,
                 action_units: List[ActionUnit],
                 version: str = "1.0",
                 is_generated: bool = False,
                 generated_at: Optional[str] = None):
        self.agent_id = agent_id
        self.action_units = action_units
        self.version = version
        self.is_generated = is_generated
        self.generated_at = generated_at
    
    def get_action_units(self) -> List[ActionUnit]:
        """Get the list of ActionUnit objects."""
        return self.action_units
    
    def to_dict(self) -> dict:
        """Serialize to dictionary format."""
        return {
            'agent_id': self.agent_id,
            'version': self.version,
            'is_generated': self.is_generated,
            'generated_at': self.generated_at,
            'action_units': [au.to_dict() for au in self.action_units]
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> 'ActionUnitSet':
        """Create ActionUnitSet from dictionary with validation."""
        # Validate required fields
        if 'agent_id' not in data:
            raise ValueError("Missing required field 'agent_id' in ActionUnitSet data")
        if 'action_units' not in data:
            raise ValueError("Missing required field 'action_units' in ActionUnitSet data")
        
        if not isinstance(data['action_units'], list):
            raise ValueError(f"Field 'action_units' must be a list in {data.get('agent_id', 'unknown')}")
        
        if len(data['action_units']) == 0:
            raise ValueError(f"ActionUnitSet for {data['agent_id']} has no action units")
        
        # Convert to ActionUnit objects
        try:
            action_units = [ActionUnit.from_dict(au_data) for au_data in data['action_units']]
        except (KeyError, TypeError) as e:
            raise ValueError(f"Invalid action unit data in {data['agent_id']}: {e}")
        
        # Validate AU names are unique
        names = [au.name for au in action_units]
        if len(names) != len(set(names)):
            duplicates = [name for name in names if names.count(name) > 1]
            raise ValueError(f"Duplicate Action Unit names in {data['agent_id']}: {set(duplicates)}")
        
        return cls(
            agent_id=data['agent_id'],
            action_units=action_units,
            version=data.get('version', '1.0'),
            is_generated=data.get('is_generated', data.get('is_autogenerated', False)),
            generated_at=data.get('generated_at')
        )
    
    def to_file(self, filepath: str) -> None:
        """Save ActionUnitSet to JSON file."""
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
    
    @classmethod
    def from_file(cls, filepath: str) -> 'ActionUnitSet':
        """Load ActionUnitSet from JSON file with error handling."""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
        except FileNotFoundError:
            raise FileNotFoundError(f"Action unit file not found: {filepath}")
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in action unit file {filepath}: {e}")
        
        return cls.from_dict(data)
    
    @classmethod
    def from_agent_id(cls, agent_id: str, use_generated: bool = False, 
                     base_dir: str = 'prompts/action_units') -> 'ActionUnitSet':
        """
        Load ActionUnitSet for a given agent.
        
        Args:
            agent_id: Agent identifier (e.g., 'geo', 'tcga', 'statistician')
            use_generated: If True, try to load generated version first
            base_dir: Base directory for action unit files
            
        Returns:
            ActionUnitSet instance
        """
        # Normalize agent_id to lowercase for file naming
        agent_id_lower = agent_id.lower()
        filename = f"{agent_id_lower}_action_units.json"
        
        # Try generated first if requested
        if use_generated:
            generated_path = os.path.join(base_dir, 'generated', filename)
            if os.path.exists(generated_path):
                return cls.from_file(generated_path)
        
        # Fall back to base
        base_path = os.path.join(base_dir, 'base', filename)
        if os.path.exists(base_path):
            return cls.from_file(base_path)
        
        raise FileNotFoundError(
            f"Action unit file not found for agent '{agent_id}' in {base_dir}"
        )
