import ast
import json
import os
from typing import Dict, Union
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from starlette import status

ast_router = APIRouter()


class CalculationInput(BaseModel):
    params: Dict[str, Union[int, float]]
    formula_name: str


def calculate_expression(expression: ast.Expression, params: Dict) -> float:
    """
    使用AST模块计算安全的表达式
    """
    try:
        code = compile(expression, "<string>", "eval")
        result = eval(code, {}, params)
        return result
    except SyntaxError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="公式语法错误：" + str(e))
    except NameError as e:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="公式中存在未定义的变量：" + str(e))
    except TypeError as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST,
                            detail="公式中的某些操作不支持所提供的参数类型：" + str(e))
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="执行公式时出现错误：" + str(e))


@ast_router.post("/calculate_ast", summary="传入指标名称与参数进行计算", tags=['指标计算'])
def calculate(request: CalculationInput):
    name = request.formula_name
    data = request.params

    # 查询公式
    with open(os.getcwd() + '\\App\\FormulaData.json', encoding='utf-8') as f:
        calculation = json.load(f)
    try:
        item = calculation.get(name)
        formula = item.get('formula')
        expression_ast = ast.parse(formula, mode="eval")
        result = calculate_expression(expression_ast, data)
        return {"result": round(result, request.precision) if hasattr(request, 'precision') else round(result, 2)}
    except AttributeError:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="公式名称错误")