# data_models/metrics.py
from dataclasses import dataclass, field
from typing import List, Dict, Any, Literal

# Define literal types for more specific type hinting
MetricType = Literal['time_series', 'aggregate', 'categorical']
TimeGranularity = Literal['daily', 'weekly', 'monthly', 'yearly', 'other'] # Added 'yearly' and 'other'

@dataclass
class TimeSeriesMetric:
    """Structure for time-series based metrics"""
    metric_name: str
    values: List[float] = field(default_factory=list)
    timestamps: List[str] = field(default_factory=list) # Consider using datetime objects or ISO format strings
    metric_type: MetricType = 'time_series'
    time_granularity: TimeGranularity = 'monthly'
    unit: Optional[str] = None # e.g., 'count', '%', 'USD'
    description: Optional[str] = None # Optional description of the metric

    def __post_init__(self):
        if len(self.values) != len(self.timestamps):
            # Or log a warning, or handle as appropriate for your application
            raise ValueError(f"Length of values ({len(self.values)}) and timestamps ({len(self.timestamps)}) must match for metric '{self.metric_name}'.")

@dataclass
class AgentMetrics:
    """
    Enhanced structure for agent metrics with time-awareness and more details.
    """
    agent_name: str
    analysis_summary: str # Summary text from the agent's analysis
    
    # Specific metric categories
    time_series_metrics: List[TimeSeriesMetric] = field(default_factory=list)
    aggregate_metrics: Dict[str, float] = field(default_factory=dict) # Key-value pairs for single value metrics
    categorical_metrics: Dict[str, Any] = field(default_factory=dict) # For distributions, counts by category, etc.
                                                                    # Example: {'industry_distribution': {'Tech': 100, 'Finance': 50}}
    
    # Contextual information
    time_periods_covered: List[str] = field(default_factory=list) # e.g., ["2023-01", "2023-02"] or ["Q1 2023", "Q2 2023"]
    data_sources_used: List[str] = field(default_factory=list) # Information about the input data
    generation_timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat()) # When these metrics were generated

    # Optional fields for richer reporting
    key_insights: List[str] = field(default_factory=list) # Bullet points of key findings
    potential_errors_or_warnings: List[str] = field(default_factory=list) # Any issues encountered during analysis