from __future__ import annotations

import ast
import operator
from dataclasses import dataclass
from math import isfinite
from typing import Any

from app.simulation.core.peng_robinson import HELIUM_PR, PengRobinsonFluid
from app.simulation.examples.test_mql.system import COMPONENT_SPECS, GLOBAL_PARAMETERS


_BINARY_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Pow: operator.pow,
}
_UNARY_OPERATORS = {
    ast.UAdd: operator.pos,
    ast.USub: operator.neg,
}


class TestMqlExpressionError(ValueError):
    """Raised when an AMESim parameter expression cannot be resolved safely."""


@dataclass(frozen=True)
class TestMqlResolvedParameter:
    name: str
    title: str
    raw_value: str
    units: str
    value: float | None

    @property
    def is_numeric(self) -> bool:
        return self.value is not None


@dataclass(frozen=True)
class TestMqlResolvedComponent:
    alias: str
    component_name: str
    submodel: str
    label: str
    parameters: dict[str, TestMqlResolvedParameter]

    def parameter_value(self, name: str) -> float:
        parameter = self.parameters[name]
        if parameter.value is None:
            raise KeyError(f"Parameter {name!r} on {self.alias!r} is not numeric")
        return parameter.value


@dataclass(frozen=True)
class TestMqlConfig:
    raw_global_parameters: dict[str, str]
    global_parameters: dict[str, float]
    fluid: PengRobinsonFluid
    components: tuple[TestMqlResolvedComponent, ...]

    @classmethod
    def from_amesim_specs(cls) -> "TestMqlConfig":
        raw_globals = dict(GLOBAL_PARAMETERS)
        numeric_globals = {
            name: value
            for name, raw in raw_globals.items()
            if (value := resolve_numeric_expression(raw, {})) is not None
        }
        components = tuple(
            _resolve_component(spec, numeric_globals)
            for spec in COMPONENT_SPECS
        )
        return cls(
            raw_global_parameters=raw_globals,
            global_parameters=numeric_globals,
            fluid=HELIUM_PR,
            components=components,
        )

    def component(self, alias: str) -> TestMqlResolvedComponent:
        for component in self.components:
            if component.alias == alias:
                return component
        raise KeyError(alias)

    def components_by_submodel(self, submodel: str) -> tuple[TestMqlResolvedComponent, ...]:
        return tuple(component for component in self.components if component.submodel == submodel)


def _resolve_component(
    spec: dict[str, Any],
    variables: dict[str, float],
) -> TestMqlResolvedComponent:
    parameters = {}
    for parameter in spec.get("parameters", []):
        name = str(parameter["name"])
        raw_value = str(parameter["value"])
        parameters[name] = TestMqlResolvedParameter(
            name=name,
            title=str(parameter["title"]),
            raw_value=raw_value,
            units=str(parameter["units"]),
            value=resolve_numeric_expression(raw_value, variables),
        )
    return TestMqlResolvedComponent(
        alias=str(spec["alias"]),
        component_name=str(spec["component_name"]),
        submodel=str(spec["submodel"]),
        label=str(spec["label"]),
        parameters=parameters,
    )


def resolve_numeric_expression(
    expression: str,
    variables: dict[str, float],
) -> float | None:
    expression = expression.strip()
    if not expression:
        return None
    normalized = expression.replace("^", "**")
    try:
        parsed = ast.parse(normalized, mode="eval")
        value = float(_eval_node(parsed.body, variables))
    except (SyntaxError, TestMqlExpressionError, ValueError, TypeError, ZeroDivisionError):
        return None
    return value if isfinite(value) else None


def _eval_node(node: ast.AST, variables: dict[str, float]) -> float:
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return float(node.value)
    if isinstance(node, ast.Name):
        if node.id not in variables:
            raise TestMqlExpressionError(f"Unknown variable: {node.id}")
        return float(variables[node.id])
    if isinstance(node, ast.BinOp):
        operator_type = type(node.op)
        if operator_type not in _BINARY_OPERATORS:
            raise TestMqlExpressionError(f"Unsupported binary operator: {operator_type}")
        return float(_BINARY_OPERATORS[operator_type](_eval_node(node.left, variables), _eval_node(node.right, variables)))
    if isinstance(node, ast.UnaryOp):
        operator_type = type(node.op)
        if operator_type not in _UNARY_OPERATORS:
            raise TestMqlExpressionError(f"Unsupported unary operator: {operator_type}")
        return float(_UNARY_OPERATORS[operator_type](_eval_node(node.operand, variables)))
    raise TestMqlExpressionError(f"Unsupported expression node: {type(node)}")