# main.py

import os, glob
import pandas as pd
import json
import shutil
import pandas as pd
import uvicorn
from fastapi import FastAPI, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field, model_validator
from typing import List, Optional, Union
from autodiag_class import Auto_diag


app = FastAPI(root_path="/api/diag",title=" Diagnosis API")


# 请求模型定义
class AutoDiagInput(BaseModel):
    ids: List[int]      # 数据索引id
    windCode: str       # 风场编号
    engine_code: str    # 风机编号
    autodiagType: str   # 诊断类型

    @model_validator(mode='before')
    def convert_ids(cls, values):
        """将单个id转换为列表形式"""
        if isinstance(values.get('ids'), int):
            values['ids'] = [values['ids']]
        return values

class DiagnosisResult(BaseModel):
    status_codes: List[int]  # 每个id对应的状态码列表
    max_status: int          # 所有状态码中的最大值
    count_0: int             # 状态码0的个数
    count_1: int             # 状态码1的个数
    count_2: int             # 状态码2的个数

@app.post("/autodiag/{autodiagType}")
async def perform_diagnosis(autodiagType: str, input_data: AutoDiagInput):
    """
    执行自动诊断分析
    参数:
        autodiagType: 诊断类型
        input_data: 包含ids, windCode, engine_code的输入数据
    返回:
        诊断结果，包含状态码列表和统计信息
    """
    autodiag_map = {
        "Unbalance": "Unbalance_diag",    # 不平衡诊断
        "Misalignment": "Misalignment_diag",  # 不对中诊断
        "Looseness": "Looseness_diag",    # 松动诊断
        "Bearing": "Bearing_diag",        # 轴承诊断       
        "Gear": "Gear_diag"              # 齿轮诊断
    }

    if autodiagType not in autodiag_map:
        raise HTTPException(status_code=400, detail="非可用的诊断类型")
    
    try:
        # 初始化诊断类
        autodiag = Auto_diag(input_data.ids, input_data.windCode, input_data.engine_code)
        
        # 获取诊断方法
        func = getattr(autodiag, autodiag_map[autodiagType])
        
        # 执行诊断
        if callable(func):
            result = func()  # 直接返回格式化后的结果
            return JSONResponse(content=result)
            
    except ValueError as e:
        # 专门捕获齿轮诊断的错误
        if "Can not perform gearbox diagnosis" in str(e):
            return JSONResponse(
                status_code=200,
                content={
                    "code": 400,
                    "message": str(e)
                }
            )
        elif "当前采集频率不适合进行诊断分析" in str(e):
            return JSONResponse(
                status_code=200,
                content={
                    "code": 405,
                    "message": str(e)
                }
            )            
        # 其他ValueError
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))