import gradio as gr
import pandas as pd
from datetime import datetime, timedelta
import io
import os
from huggingface_hub import HfApi, hf_hub_download, login
from typing import List, Tuple, Optional
import tempfile
import zipfile

class DataManager:
    def __init__(self, token=None):
        self.repo_id = "CIV3283/Data_Storage"
        self.branch = "data_branch"
        
        # 优先使用传入的token，否则从环境变量获取
        self.token = token or os.getenv('CIV3283_access')
        
        # 如果有token，进行认证
        if self.token:
            try:
                login(token=self.token)
                print("✅ 认证成功")
            except Exception as e:
                print(f"❌ 认证失败: {e}")
        
        self.api = HfApi(token=self.token)
        
    def get_available_files(self) -> Tuple[List[str], List[str]]:
        """获取所有可用的feedback和query文件"""
        try:
            # 尝试获取仓库中的所有文件
            repo_files = self.api.list_repo_files(
                repo_id=self.repo_id,
                revision=self.branch,
                token=self.token,
                repo_type="space"  # 明确指定这是一个space
            )
            
            feedback_files = []
            query_files = []
            
            for file in repo_files:
                if file.endswith('_feedback_log.csv') and 'CIV3283_Student_' in file:
                    feedback_files.append(file)
                elif file.endswith('_query_log.csv') and 'CIV3283_Student_' in file:
                    query_files.append(file)
            
            return sorted(feedback_files), sorted(query_files)
            
        except Exception as e:
            print(f"❌ 获取文件列表错误: {e}")
            # 如果API调用失败，尝试手动生成可能的文件名列表
            return self._generate_possible_filenames()
    
    def _generate_possible_filenames(self) -> Tuple[List[str], List[str]]:
        """生成可能的文件名列表作为后备方案"""
        feedback_files = []
        query_files = []
        
        # 生成01到95的所有可能文件名
        for i in range(1, 96):
            student_id = f"{i:02d}"  # 格式化为两位数
            feedback_file = f"CIV3283_Student_{student_id}_feedback_log.csv"
            query_file = f"CIV3283_Student_{student_id}_query_log.csv"
            feedback_files.append(feedback_file)
            query_files.append(query_file)
        
        return feedback_files, query_files
    
    def download_and_read_csv(self, filename: str) -> pd.DataFrame:
        """下载并读取CSV文件"""
        try:
            # 下载文件到临时位置
            local_file = hf_hub_download(
                repo_id=self.repo_id,
                filename=filename,
                revision=self.branch,
                repo_type="space",
                token=self.token,
                cache_dir=tempfile.gettempdir()
            )
            
            # 读取CSV
            df = pd.read_csv(local_file)
            
            # 确保timestamp列存在并转换为datetime
            if 'timestamp' in df.columns:
                df['timestamp'] = pd.to_datetime(df['timestamp'])
            
            return df
            
        except Exception as e:
            print(f"❌ 读取文件 {filename} 失败: {e}")
            # 返回空DataFrame而不是完全失败
            return pd.DataFrame()
    
    def filter_data_by_date(self, df: pd.DataFrame, start_date: str, end_date: str) -> pd.DataFrame:
        """根据日期范围筛选数据"""
        if df.empty or 'timestamp' not in df.columns:
            return df
        
        try:
            # 转换输入的日期字符串
            start_dt = pd.to_datetime(start_date)
            
            if end_date == "现在":
                end_dt = pd.Timestamp.now()
            else:
                end_dt = pd.to_datetime(end_date) + timedelta(days=1) - timedelta(seconds=1)  # 包含整天
            
            # 筛选数据
            mask = (df['timestamp'] >= start_dt) & (df['timestamp'] <= end_dt)
            return df[mask]
            
        except Exception as e:
            print(f"Error filtering data: {e}")
            return df
    
    def process_files(self, file_type: str, start_date: str, end_date: str) -> Tuple[str, str, str]:
        """处理文件并生成下载内容"""
        feedback_files, query_files = self.get_available_files()
        
        result_data = []
        file_info = []
        total_records = 0
        
        # 根据选择的文件类型确定要处理的文件
        files_to_process = []
        if file_type in ["反馈文件", "两者都要"]:
            files_to_process.extend([(f, "feedback") for f in feedback_files])
        if file_type in ["查询文件", "两者都要"]:
            files_to_process.extend([(f, "query") for f in query_files])
        
        # 处理每个文件
        for filename, ftype in files_to_process:
            df = self.download_and_read_csv(filename)
            if not df.empty:
                original_count = len(df)
                filtered_df = self.filter_data_by_date(df, start_date, end_date)
                filtered_count = len(filtered_df)
                
                if filtered_count > 0:
                    result_data.append(filtered_df)
                    total_records += filtered_count
                
                file_info.append(f"{filename}: {original_count}条记录 → {filtered_count}条符合条件")
        
        # 合并所有数据
        if result_data:
            combined_df = pd.concat(result_data, ignore_index=True)
            
            # 生成CSV内容
            csv_buffer = io.StringIO()
            combined_df.to_csv(csv_buffer, index=False, encoding='utf-8')
            csv_content = csv_buffer.getvalue()
            
            # 生成文件名
            end_date_str = "现在" if end_date == "现在" else end_date
            filename = f"CIV3283_数据导出_{file_type}_{start_date}_to_{end_date_str}.csv"
            
            # 生成状态信息
            status = f"✅ 成功处理 {len(files_to_process)} 个文件\n"
            status += f"📊 总共找到 {total_records} 条符合条件的记录\n"
            status += f"📅 时间范围: {start_date} 到 {end_date_str}\n\n"
            status += "文件详情:\n" + "\n".join(file_info)
            
            return csv_content, filename, status
        else:
            return "", "", "❌ 没有找到符合条件的数据"

# 全局变量存储数据管理器
data_manager = None

def initialize_data_manager(hf_token=None):
    """初始化数据管理器"""
    global data_manager
    data_manager = DataManager(token=hf_token)
    return data_manager

def process_and_download(file_type, start_date, end_date, hf_token):
    """处理请求并返回下载文件"""
    try:
        # 确保数据管理器已初始化
        global data_manager
        env_token = os.getenv('CIV3283_access')
        current_token = hf_token or env_token
        
        if data_manager is None or data_manager.token != current_token:
            data_manager = initialize_data_manager(current_token)
        
        csv_content, filename, status = data_manager.process_files(file_type, start_date, end_date)
        
        if csv_content:
            # 创建临时文件
            temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, encoding='utf-8')
            temp_file.write(csv_content)
            temp_file.close()
            
            return status, temp_file.name, filename
        else:
            return status, None, None
            
    except Exception as e:
        return f"❌ 处理过程中出现错误: {str(e)}", None, None

def get_file_summary(hf_token=None):
    """获取文件概览"""
    try:
        # 如果还没有初始化或token发生变化，重新初始化
        global data_manager
        env_token = os.getenv('CIV3283_access')
        current_token = hf_token or env_token
        
        if data_manager is None or data_manager.token != current_token:
            data_manager = initialize_data_manager(current_token)
        
        feedback_files, query_files = data_manager.get_available_files()
        summary = f"📁 当前可用文件:\n"
        summary += f"   反馈文件: {len(feedback_files)} 个\n"
        summary += f"   查询文件: {len(query_files)} 个\n"
        summary += f"   总计: {len(feedback_files) + len(query_files)} 个文件"
        
        if env_token:
            summary += f"\n✅ 使用环境变量 CIV3283_access 认证"
        elif hf_token:
            summary += f"\n✅ 使用手动输入的token认证"
        else:
            summary += f"\n⚠️  未找到认证token（可能无法访问私有仓库）"
            
        return summary
    except Exception as e:
        return f"❌ 获取文件列表失败: {str(e)}"

# 创建Gradio界面
def create_interface():
    with gr.Blocks(title="CIV3283 数据管理器", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🗃️ CIV3283 数据存储管理器")
        gr.Markdown("从 Hugging Face Space 的 data_branch 分支整理和下载学生数据")
        
        # 认证部分
        with gr.Row():
            with gr.Column(scale=3):
                hf_token = gr.Textbox(
                    label="🔐 Hugging Face Token (可选)",
                    placeholder="hf_xxxxxxxxxxxx",
                    type="password",
                    info="可选：手动输入token，留空将使用环境变量 CIV3283_access"
                )
            with gr.Column(scale=1):
                gr.Markdown(f"""
                **环境变量状态:**
                {'✅ CIV3283_access 已设置' if os.getenv('CIV3283_access') else '❌ CIV3283_access 未找到'}
                """)

        
        # 文件概览
        with gr.Row():
            file_summary = gr.Textbox(
                label="📊 文件概览",
                value="点击刷新按钮获取文件列表",
                interactive=False,
                max_lines=6
            )
            refresh_btn = gr.Button("🔄 刷新文件列表", variant="secondary")
        
        # 自动刷新一次（使用环境变量）
        demo.load(
            fn=lambda: get_file_summary(None),
            outputs=file_summary
        )
        
        refresh_btn.click(
            fn=get_file_summary,
            inputs=[hf_token],
            outputs=file_summary
        )
        
        # 主要控制面板
        with gr.Row():
            with gr.Column(scale=1):
                file_type = gr.Radio(
                    choices=["查询文件", "反馈文件", "两者都要"],
                    value="查询文件",
                    label="📂 选择文件类型"
                )
                
                start_date = gr.Textbox(
                    label="📅 开始日期",
                    placeholder="2025-06-16",
                    value="2025-06-16",
                    info="格式: YYYY-MM-DD"
                )
                
                end_date = gr.Textbox(
                    label="📅 结束日期",
                    placeholder="2025-06-20 或 现在",
                    value="现在",
                    info="格式: YYYY-MM-DD 或输入'现在'"
                )
                
                process_btn = gr.Button("🚀 处理并生成下载文件", variant="primary", size="lg")
            
            with gr.Column(scale=2):
                status_box = gr.Textbox(
                    label="📋 处理状态",
                    max_lines=15,
                    interactive=False
                )
                
                download_file = gr.File(
                    label="📥 下载文件",
                    interactive=False
                )
        
        # 使用说明
        with gr.Accordion("📖 使用说明", open=False):
            gr.Markdown("""
            ### 认证说明
            - **环境变量优先**: 系统会自动使用环境变量 `CIV3283_access` 中的token
            - **手动覆盖**: 如果需要使用不同的token，可以在上方手动输入
            - **获取token**: 如果需要新token，访问 https://huggingface.co/settings/tokens
            
            ### 功能说明
            1. **文件类型选择**: 选择要导出的数据类型
               - 查询文件: 包含搜索信息和查询响应
               - 反馈文件: 包含学生反馈评论
               - 两者都要: 合并所有数据
            
            2. **日期范围**: 
               - 开始日期: 必须填写，格式 YYYY-MM-DD
               - 结束日期: 可以是具体日期或"现在"
            
            3. **数据格式**:
               - Query文件列: student_space, student_id, timestamp, search_info, query_and_response, thumb_feedback
               - Feedback文件列: student_space, student_id, timestamp, comment
            
            ### 注意事项
            - 系统会自动发现所有可用的学生文件（01-95，跳过缺失的编号）
            - 时间戳格式: 2025-06-16 08:55:33
            - 导出的数据保持原始顺序，不进行时间排序
            - 每次处理会显示详细的文件统计信息
            - 如果无法访问仓库API，系统会尝试预设的文件名列表
            """)
        
        # 绑定处理函数
        process_btn.click(
            fn=process_and_download,
            inputs=[file_type, start_date, end_date, hf_token],
            outputs=[status_box, download_file, gr.State()]
        )
    
    return demo

# 启动应用
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(

    )