#!/usr/bin/env python3
from typing import Optional, List, Dict
from pydantic import BaseModel, ConfigDict, Field


class EvaluationServiceInputParams(BaseModel):
    task_id: str = Field(description="任务ID")
    input_path: Optional[str] = Field(
        description="推理结果S3地址",
        default="tool_evaluation/infer_Result/Dataset7.n_RAG_e2e.json"
    )
    result_detail_path: Optional[str] = Field(
        description="评估结果S3地址", default="tool_evaluation/result_detail/Dataset7.n_RAG_e2e.json"
    )
    result_metric_path: Optional[str] = Field(
        description="指标结果S3地址",
        default="tool_evaluation/result_metrics/Dataset7.n_RAG_e2e.json"
    )
    model_config = ConfigDict(extra="allow")


class BaseEvaluationServiceConfig(BaseModel):
    progress_log_dir: str = Field(default="./data/progress", description="进度日志目录")
    result_dir: str = Field(default="./data/results", description="评估结果本地目录")
    progress_file: str = Field(default="./data/results/task_progress.json", description="进度存储文件")


class BaseEvaluationRaw(BaseModel):
    data_id: int = Field(description="数据唯一ID")
    raw_data: Dict = Field(default_factory=dict, description="原始数据字典")
    inputs: Dict = Field(default_factory=dict, description="服务输入数据")
    outputs: Dict = Field(default_factory=dict, description="服务输出数据")


class EvaluationData(BaseModel):
    data_list: List[BaseEvaluationRaw] = Field(default_factory=list)


class EvaluationResult(BaseModel):
    data_id: int = Field(description="数据唯一ID")
    metrics: Dict[str, float] = Field(default_factory=dict, description="单个结果的指标")


class EvaluationResults(BaseModel):
    results: List[EvaluationResult] = Field(default_factory=list)
    metrics: Dict = Field(default_factory=dict, description="性能指标")


class EvaluationContext(BaseModel):
    """评估上下文"""
    task_id: str = Field(description="任务ID")
    input_params: EvaluationServiceInputParams = Field(description="任务输入参数")
    temp_files: List[str] = Field(default_factory=list, description="临时文件列表")
    local_detail_path: Optional[str] = Field(default=None, description="本地详细结果路径")
    local_h5_detail_path: Optional[str] = Field(default=None, description="本地详细结果路径")

    def add_temp_file(self, file_path: str) -> None:
        """添加临时文件路径"""
        if file_path not in self.temp_files:
            self.temp_files.append(file_path)

    def set_local_detail_path(self, path: str) -> None:
        """设置本地详细结果路径"""
        self.local_detail_path = path
        self.add_temp_file(path)

    def set_local_h5_detail_path(self, path: str) -> None:
        """设置本地HDF5详细结果路径"""
        self.local_h5_detail_path = path
        self.add_temp_file(path)

    class Config:
        # 允许任意类型，因为我们有自定义方法
        arbitrary_types_allowed = True
        # 使用枚举值而不是名称
        use_enum_values = True
