import os
import yaml
import openpyxl
import requests
from openpyxl import Workbook
from openpyxl.styles import Font

# API配置
API_KEY = "app-m7XGgbTe3BVHmA1TAYg9Ec4v"  # Dify API Key
WORKFLOW_ID = os.getenv("DIFY_WORKFLOW_ID", "your-workflow-id-here")  # 工作流ID
API_BASE_URL = "http://192.168.100.143/v1"  # API基础地址
FILE_UPLOAD_URL = f"{API_BASE_URL}/files/upload"  # 文件上传地址
EXCEL_PATH = "ai作业/AI考试作业.xlsx"
ASSIGNMENT_DIR = "ai作业/作业"
OUTPUT_DIR = "results"
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "批改结果.xlsx")

# 确保输出目录存在
os.makedirs(OUTPUT_DIR, exist_ok=True)

def read_excel_submissions():
    """读取Excel中的作业提交记录"""
    wb = openpyxl.load_workbook(EXCEL_PATH)
    ws = wb.active
    submissions = []
    
    # 获取标题行确定列索引
    headers = [cell.value for cell in ws[1]]
    name_col = headers.index('填写人')
    workflow_col = headers.index('工作流程描述')
    solution_col = headers.index('Dify工作流解决方案设计')
    
    for row in ws.iter_rows(min_row=2, values_only=True):
        if len(row) > max(name_col, workflow_col, solution_col):  # 确保有足够列
            submissions.append({
                'name': row[name_col],
                'work_description': row[workflow_col],
                'solution': row[solution_col]
            })
    return submissions

def check_yml_files(submissions):
    """检查Excel中的提交记录是否有对应的YML文件"""
    missing_files = []
    for sub in submissions:
        yml_path = find_assignment_yml(sub['name'])
        if not yml_path:
            missing_files.append(sub['name'])
    
    if missing_files:
        print("\n以下提交记录缺少对应的YML文件:")
        for name in missing_files:
            print(f"- {name}")
        return False
    return True

def find_assignment_yml(name):
    """根据姓名查找对应的YML作业文件"""
    for filename in os.listdir(ASSIGNMENT_DIR):
        if filename.endswith('.yml'):  # 仅支持.yml格式
            # 从文件名中提取姓名部分(第二个下划线分隔的部分)
            parts = filename.split('_')
            if len(parts) >= 2 and name == parts[1]:
                return os.path.join(ASSIGNMENT_DIR, filename)
    return None

def parse_yml_file(yml_path):
    """解析YML文件内容"""
    with open(yml_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def call_dify_api(yml_path, assignment_data):
    """调用Dify API进行作业批改，上传YML文件"""
    # 先上传文件
    upload_headers = {
        "Authorization": f"Bearer {API_KEY}"
    }
    upload_data = {
        "user": "ai-grading-system",
        "type": "yml"  # 自定义文件类型
    }
    
    try:
        # 上传文件
        with open(yml_path, 'rb') as f:
            files = {'file': (os.path.basename(yml_path), f, 'text/plain')}
            upload_response = requests.post(
                FILE_UPLOAD_URL,
                headers=upload_headers,
                files=files,
                data=upload_data
            )
            upload_response.raise_for_status()
            file_id = upload_response.json().get('id')
            
            if not file_id:
                print("文件上传失败: 未获取到文件ID")
                return None

        # 执行工作流
        run_url = f"{API_BASE_URL}/workflows/run"
        run_headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        run_data = {
            "inputs": {
                "yml_file": {
                    "transfer_method": "local_file",
                    "upload_file_id": file_id,
                    "type": "custom"
                },
                "work_description": assignment_data.get('work_description', ''),
                "solution": assignment_data.get('solution', '')
            },
            "response_mode": "blocking",
            "user": "ai-grading-system"
        }
        
        response = requests.post(run_url, headers=run_headers, json=run_data)
        response.raise_for_status()
        return response.json()
        
    except requests.exceptions.RequestException as e:
        print(f"API调用失败: {e}")
        return None

def save_results(results):
    """保存批改结果到Excel"""
    wb = Workbook()
    ws = wb.active
    ws.title = "批改结果"
    
    # 添加表头
    headers = ["姓名", "评分", "评分详情"]
    ws.append(headers)
    
    # 设置表头样式
    for cell in ws[1]:
        cell.font = Font(bold=True)
    
    # 添加数据
    for result in results:
        ws.append([
            result['name'],
            result.get('score', 'N/A'),
            str(result.get('details', '无详情'))
        ])
    
    # 自动调整列宽
    for column in ws.columns:
        max_length = 0
        column = [cell for cell in column]
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(cell.value)
            except:
                pass
        adjusted_width = (max_length + 2) * 1.2
        ws.column_dimensions[column[0].column_letter].width = adjusted_width
    
    wb.save(OUTPUT_FILE)
    print(f"批改结果已保存到: {OUTPUT_FILE}")

def main():
    # 读取Excel中的作业提交
    submissions = read_excel_submissions()
    
    # 检查YML文件是否存在
    if not check_yml_files(submissions):
        print("\n请补充缺少的YML文件后再运行批改")
        return
    
    results = []
    
    for sub in submissions:
        print(f"正在处理: {sub['name']}")
        
        # 查找对应的YML文件
        yml_path = find_assignment_yml(sub['name'])
        if not yml_path:
            print(f"未找到 {sub['name']} 的作业文件")
            continue
            
        # 解析YML文件
        try:
            yml_content = parse_yml_file(yml_path)
        except Exception as e:
            print(f"解析YML文件失败: {e}")
            continue
            
        # 准备API调用数据
        assignment_data = {
            **sub,
            **yml_content
        }
        
        # 调用API进行批改(上传YML文件)
        api_response = call_dify_api(yml_path, assignment_data)
        if not api_response:
            print(f"{sub['name']} 批改失败")
            continue
            
        # 解析API响应(直接从文件上传返回的结果)
        try:
            outputs = api_response.get('data', {}).get('outputs', {})
            results.append({
                'name': sub['name'],
                'score': outputs.get('score'),
                'details': outputs.get('details')
            })
            print(f"{sub['name']} 批改完成")
        except Exception as e:
            print(f"解析API响应失败: {e}")
    
    # 保存结果
    save_results(results)

if __name__ == "__main__":
    main()