data_manager.py

"""
数据管理模块 - 负责加载数据集和管理审核记录
"""
import os
import json
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict
from pathlib import Path


# ============== 数据类定义 ==============

@dataclass
class LabelInfo:
    """图表标签信息"""
    Number: str
    Type: str
    Source: str
    Weblink: str
    Topic: str
    Describe: str
    Other: str = ""


@dataclass
class QAItem:
    """问答对"""
    id: str
    chart: str
    question: str
    answer: str


@dataclass
class ReviewRecord:
    """审核记录"""
    review_id: str
    chart_id: str
    qa_id: str
    source: str
    chart_type: str
    model: str
    original_question: str
    original_answer: str
    status: str  # correct, incorrect, needs_modification, pending
    modified_question: str = ""
    modified_answer: str = ""
    issue_type: str = ""  # question_ambiguous, answer_wrong, chart_unclear, other
    comment: str = ""
    reviewer: str = "default"
    review_time: str = ""

    def to_dict(self) -> Dict:
        return asdict(self)


# ============== 数据管理器 ==============

class DataManager:
    """数据管理器"""
    
    def __init__(self, dataset_path: str = "./dataset", reviews_path: str = "./reviews"):
        """
        初始化数据管理器
        
        Args:
            dataset_path: 数据集根目录路径
            reviews_path: 审核记录保存目录路径
        """
        self.dataset_path = Path(dataset_path)
        self.reviews_path = Path(reviews_path)
        self.web_path = self.dataset_path / "web"
        self.label_path = self.dataset_path / "label"
        self.qa_path = self.dataset_path / "question_answer"
        self.reviews_file = self.reviews_path / "reviews.json"
        
        # 确保目录存在
        self.reviews_path.mkdir(parents=True, exist_ok=True)
        if not self.reviews_file.exists():
            self._save_reviews([])
    
    # ============== 数据集结构获取 ==============
    
    def get_dataset_structure(self) -> Dict[str, Any]:
        """
        获取数据集的目录结构
        
        Returns:
            包含 source -> chart_type -> models 的树形结构
        """
        structure = {"sources": {}}
        
        if not self.web_path.exists():
            return structure
        
        # 遍历所有 source
        for source_name in self._list_dirs(self.web_path):
            source_web_path = self.web_path / source_name
            source_qa_path = self.qa_path / source_name
            
            structure["sources"][source_name] = {"chart_types": {}}
            
            # 遍历所有 chart_type
            for chart_type_name in self._list_dirs(source_web_path):
                chart_type_web_path = source_web_path / chart_type_name
                chart_type_qa_path = source_qa_path / chart_type_name
                
                # 获取图表数量
                chart_files = [f for f in self._list_files(chart_type_web_path) if f.endswith('.html')]
                chart_count = len(chart_files)
                
                # 获取模型列表
                models = self._list_dirs(chart_type_qa_path) if chart_type_qa_path.exists() else []
                
                # 获取已审核数量
                reviews = self.get_all_reviews()
                reviewed_charts = set()
                for r in reviews:
                    if r.get('source') == source_name and r.get('chart_type') == chart_type_name:
                        reviewed_charts.add(r.get('chart_id'))
                reviewed_count = len(reviewed_charts)
                
                structure["sources"][source_name]["chart_types"][chart_type_name] = {
                    "chart_count": chart_count,
                    "reviewed_count": reviewed_count,
                    "models": models
                }
        
        return structure
    
    def get_chart_list(self, source: str, chart_type: str) -> List[str]:
        """获取指定 source 和 chart_type 下的所有图表 ID"""
        chart_type_path = self.web_path / source / chart_type
        if not chart_type_path.exists():
            return []
        
        return [f.replace('.html', '') for f in self._list_files(chart_type_path) if f.endswith('.html')]
    
    def get_all_chart_paths(self) -> List[Dict[str, str]]:
        """
        获取所有图表的完整路径（用于导航）
        
        Returns:
            包含 {source, chart_type, chart_id, model} 的列表
        """
        paths = []
        structure = self.get_dataset_structure()
        
        for source_name, source_data in structure["sources"].items():
            for chart_type_name, chart_type_data in source_data["chart_types"].items():
                chart_ids = self.get_chart_list(source_name, chart_type_name)
                for chart_id in chart_ids:
                    for model in chart_type_data.get("models", []):
                        paths.append({
                            "source": source_name,
                            "chart_type": chart_type_name,
                            "chart_id": chart_id,
                            "model": model
                        })
        
        return paths
    
    # ============== 图表数据获取 ==============
    
    def get_chart_data(self, source: str, chart_type: str, chart_id: str) -> Dict[str, Any]:
        """
        获取图表数据（HTML内容和标签信息）
        
        Args:
            source: 数据来源
            chart_type: 图表类型
            chart_id: 图表ID
            
        Returns:
            包含 html_content, html_path, label_info 的字典
        """
        html_path = self.web_path / source / chart_type / f"{chart_id}.html"
        label_path = self.label_path / source / chart_type / f"{chart_id}.json"
        
        result = {
            "html_content": "",
            "html_path": str(html_path) if html_path.exists() else "",
            "label_info": None
        }
        
        # 读取 HTML 内容
        if html_path.exists():
            with open(html_path, 'r', encoding='utf-8') as f:
                result["html_content"] = f.read()
        
        # 读取标签信息
        if label_path.exists():
            try:
                with open(label_path, 'r', encoding='utf-8') as f:
                    label_data = json.load(f)
                    result["label_info"] = label_data
            except Exception as e:
                print(f"Error reading label file: {e}")
        
        return result
    
    # ============== QA 数据获取 ==============
    
    def get_qa_list(self, source: str, chart_type: str, model: str, chart_id: str) -> List[QAItem]:
        """
        获取指定图表的 QA 列表
        
        Args:
            source: 数据来源
            chart_type: 图表类型
            model: 模型名称
            chart_id: 图表ID
            
        Returns:
            QAItem 列表
        """
        qa_model_path = self.qa_path / source / chart_type / model
        
        if not qa_model_path.exists():
            return []
        
        qa_list = []
        for qa_file in self._list_files(qa_model_path):
            if not qa_file.endswith('.json'):
                continue
            
            try:
                with open(qa_model_path / qa_file, 'r', encoding='utf-8') as f:
                    qa_data = json.load(f)
                    # 筛选属于当前图表的 QA
                    if qa_data.get('chart') == chart_id or chart_id in qa_data.get('id', ''):
                        qa_list.append(QAItem(
                            id=qa_data.get('id', ''),
                            chart=qa_data.get('chart', ''),
                            question=qa_data.get('question', ''),
                            answer=qa_data.get('answer', '')
                        ))
            except Exception as e:
                print(f"Error reading QA file {qa_file}: {e}")
        
        return qa_list
    
    def get_all_qa_for_chart(self, source: str, chart_type: str, chart_id: str) -> Dict[str, List[QAItem]]:
        """
        获取指定图表所有模型的 QA 数据
        
        Returns:
            {model_name: [QAItem, ...], ...} 的字典
        """
        chart_qa_path = self.qa_path / source / chart_type
        
        if not chart_qa_path.exists():
            return {}
        
        result = {}
        for model in self._list_dirs(chart_qa_path):
            qa_list = self.get_qa_list(source, chart_type, model, chart_id)
            if qa_list:
                result[model] = qa_list
        
        return result
    
    # ============== 审核记录管理 ==============
    
    def _save_reviews(self, reviews: List[Dict]):
        """保存审核记录到文件"""
        with open(self.reviews_file, 'w', encoding='utf-8') as f:
            json.dump(reviews, f, ensure_ascii=False, indent=2)
    
    def get_all_reviews(self) -> List[Dict]:
        """获取所有审核记录"""
        if not self.reviews_file.exists():
            return []
        
        try:
            with open(self.reviews_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except:
            return []
    
    def get_review_by_qa_id(self, qa_id: str) -> Optional[Dict]:
        """根据 QA ID 获取审核记录"""
        reviews = self.get_all_reviews()
        for r in reviews:
            if r.get('qa_id') == qa_id:
                return r
        return None
    
    def get_reviews_by_chart(self, chart_id: str, model: Optional[str] = None) -> List[Dict]:
        """获取指定图表的审核记录"""
        reviews = self.get_all_reviews()
        result = []
        for r in reviews:
            if r.get('chart_id') == chart_id:
                if model is None or r.get('model') == model:
                    result.append(r)
        return result
    
    def save_review(self, review_data: Dict) -> Dict:
        """
        保存审核记录
        
        Args:
            review_data: 审核数据字典
            
        Returns:
            保存后的审核记录
        """
        reviews = self.get_all_reviews()
        
        # 生成审核记录
        review = {
            "review_id": review_data.get('review_id') or str(uuid.uuid4()),
            "chart_id": review_data.get('chart_id', ''),
            "qa_id": review_data.get('qa_id', ''),
            "source": review_data.get('source', ''),
            "chart_type": review_data.get('chart_type', ''),
            "model": review_data.get('model', ''),
            "original_question": review_data.get('original_question', ''),
            "original_answer": review_data.get('original_answer', ''),
            "status": review_data.get('status', 'pending'),
            "modified_question": review_data.get('modified_question', ''),
            "modified_answer": review_data.get('modified_answer', ''),
            "issue_type": review_data.get('issue_type', ''),
            "comment": review_data.get('comment', ''),
            "reviewer": review_data.get('reviewer', 'default'),
            "review_time": datetime.now().isoformat()
        }
        
        # 更新或添加
        existing_index = None
        for i, r in enumerate(reviews):
            if r.get('qa_id') == review['qa_id']:
                existing_index = i
                break
        
        if existing_index is not None:
            reviews[existing_index] = review
        else:
            reviews.append(review)
        
        self._save_reviews(reviews)
        return review
    
    def get_review_stats(self) -> Dict[str, int]:
        """获取审核统计信息"""
        reviews = self.get_all_reviews()
        return {
            "total": len(reviews),
            "correct": len([r for r in reviews if r.get('status') == 'correct']),
            "incorrect": len([r for r in reviews if r.get('status') == 'incorrect']),
            "needs_modification": len([r for r in reviews if r.get('status') == 'needs_modification']),
            "pending": len([r for r in reviews if r.get('status') == 'pending'])
        }
    
    def export_reviews(self, output_path: str = None) -> str:
        """
        导出审核记录
        
        Args:
            output_path: 输出文件路径，如果为 None 则自动生成
            
        Returns:
            导出文件的路径
        """
        reviews = self.get_all_reviews()
        
        if output_path is None:
            output_path = f"./reviews_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(reviews, f, ensure_ascii=False, indent=2)
        
        return output_path
    
    # ============== 辅助方法 ==============
    
    def _list_dirs(self, path: Path) -> List[str]:
        """列出目录下的所有子目录"""
        if not path.exists():
            return []
        return [d.name for d in path.iterdir() if d.is_dir() and not d.name.startswith('.')]
    
    def _list_files(self, path: Path) -> List[str]:
        """列出目录下的所有文件"""
        if not path.exists():
            return []
        return [f.name for f in path.iterdir() if f.is_file() and not f.name.startswith('.')]


# ============== 全局实例 ==============

# 默认数据管理器实例
data_manager = DataManager()