from __future__ import annotations

from datetime import datetime, timezone
import json
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
from xml.etree import ElementTree as ET

from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import FileResponse, HTMLResponse
from pydantic import BaseModel, Field

from app.system_xml import (
    SystemXmlDocument,
    SystemXmlValidationReport,
    validate_system_xml_document,
)

if TYPE_CHECKING:
    from PythonModels.core.network import SimulationNetwork
    from PythonModels.core.ports import PortDefinition


app = FastAPI(title="System Simulation ReactFlow App")
FRONTEND_DIST_DIR = Path(__file__).resolve().parent.parent / "frontend" / "dist"
PROJECT_STORAGE_DIR = Path(__file__).parent / "data" / "reactflow-projects"
SYSTEM_XML_SCHEMA_VERSION = "2"
SYSTEM_XML_UNIT_SYSTEM = "SI"


class ReactFlowPosition(BaseModel):
    x: float = 0.0
    y: float = 0.0


class ReactFlowPortDefinition(BaseModel):
    name: str
    kind: Literal["physical", "signal"] = "physical"
    domain: str = "pneumatic"
    nominalRole: Literal[
        "inlet",
        "outlet",
        "bidirectional",
        "input",
        "output",
    ] = "bidirectional"
    positiveFlowDirection: Literal["intoComponent"] | None = None
    side: Literal["left", "right"] = "left"


class ReactFlowNodeData(BaseModel):
    label: str = ""
    componentType: str = "component"
    modelType: str = "component"
    ports: list[ReactFlowPortDefinition | str] = Field(default_factory=list)
    parameters: dict[str, Any] = Field(default_factory=dict)
    rotation: Literal[0, 90, 180, 270] = 0
    mirrored: bool = False


class ReactFlowNodePayload(BaseModel):
    id: str
    type: str | None = None
    position: ReactFlowPosition = Field(default_factory=ReactFlowPosition)
    data: ReactFlowNodeData


class ReactFlowEdgePayload(BaseModel):
    id: str
    source: str
    target: str
    sourceHandle: str | None = None
    targetHandle: str | None = None


class ReactFlowSimulationConfig(BaseModel):
    t_start: float = 0.0
    t_stop: float = 2.0
    step: float = 0.1
    max_step: float = 0.005
    method: str = "BDF"


class ReactFlowProjectPayload(BaseModel):
    name: str = "untitled"
    nodes: list[ReactFlowNodePayload] = Field(default_factory=list)
    edges: list[ReactFlowEdgePayload] = Field(default_factory=list)
    simulation: ReactFlowSimulationConfig = Field(default_factory=ReactFlowSimulationConfig)


def pydantic_to_jsonable(model: BaseModel) -> dict[str, Any]:
    if hasattr(model, "model_dump"):
        return model.model_dump(mode="json")
    return model.dict()


@app.get("/")
def index() -> HTMLResponse:
    frontend_index = FRONTEND_DIST_DIR / "index.html"
    if frontend_index.exists():
        return HTMLResponse(
            frontend_index.read_text(encoding="utf-8"),
            headers={"Cache-Control": "no-store"},
        )

    return HTMLResponse(
        """
        <!doctype html>
        <html lang="zh-CN">
          <head>
            <meta charset="utf-8" />
            <title>System Simulation Flow</title>
            <style>
              body {
                margin: 0;
                font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
                color: #1f2933;
                background: #f8fafc;
              }
              main {
                max-width: 720px;
                margin: 96px auto;
                padding: 0 24px;
              }
              code {
                padding: 2px 6px;
                border-radius: 4px;
                background: #e5e7eb;
              }
            </style>
          </head>
          <body>
            <main>
              <h1>System Simulation Flow</h1>
              <p>ReactFlow 前端未构建。开发时请在 <code>frontend</code> 目录运行 <code>npm run dev</code>，然后打开 <code>http://127.0.0.1:5173</code>。</p>
              <p>如果要让 FastAPI 直接服务页面，请先运行 <code>npm run build</code> 生成 <code>frontend/dist</code>。</p>
            </main>
          </body>
        </html>
        """,
        headers={"Cache-Control": "no-store"},
    )


@app.get("/assets/{path:path}")
def frontend_asset(path: str) -> FileResponse:
    asset_path = (FRONTEND_DIST_DIR / "assets" / path).resolve()
    assets_root = (FRONTEND_DIST_DIR / "assets").resolve()
    try:
        asset_path.relative_to(assets_root)
    except ValueError as exc:
        raise HTTPException(status_code=404, detail="Frontend asset was not found.") from exc
    if not asset_path.exists():
        raise HTTPException(status_code=404, detail="Frontend asset was not found.")
    return FileResponse(asset_path)


@app.post("/api/reactflow/system-xml")
def export_reactflow_system_xml(payload: ReactFlowProjectPayload) -> Response:
    try:
        xml = build_reactflow_system_xml(payload)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return Response(content=xml, media_type="application/xml")


@app.get("/api/reactflow/projects")
def list_reactflow_projects() -> dict[str, object]:
    PROJECT_STORAGE_DIR.mkdir(parents=True, exist_ok=True)
    projects = []
    for path in sorted(PROJECT_STORAGE_DIR.glob("*.json")):
        projects.append(
            {
                "id": path.stem,
                "updatedAt": datetime.fromtimestamp(
                    path.stat().st_mtime,
                    timezone.utc,
                ).isoformat(),
            }
        )
    return {"projects": projects}


@app.get("/api/reactflow/projects/{project_id}")
def load_reactflow_project(project_id: str) -> dict[str, object]:
    path = reactflow_project_path(project_id)
    if not path.exists():
        raise HTTPException(status_code=404, detail="Project was not found.")
    return json.loads(path.read_text(encoding="utf-8"))


@app.post("/api/reactflow/projects/{project_id}")
def save_reactflow_project(
    project_id: str,
    payload: ReactFlowProjectPayload,
) -> dict[str, object]:
    path = reactflow_project_path(project_id)
    PROJECT_STORAGE_DIR.mkdir(parents=True, exist_ok=True)
    data = pydantic_to_jsonable(payload)
    data["name"] = payload.name or project_id
    path.write_text(
        json.dumps(data, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return {
        "id": path.stem,
        "path": str(path),
        "componentCount": len(payload.nodes),
        "connectionCount": len(payload.edges),
    }


@app.post("/api/reactflow/simulate-testmodel")
def simulate_reactflow_testmodel(payload: ReactFlowProjectPayload) -> dict[str, object]:
    try:
        result = run_reactflow_testmodel(payload)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return result


@app.post("/api/reactflow/compile-model")
def compile_reactflow_model(payload: ReactFlowProjectPayload) -> dict[str, object]:
    try:
        network = compile_reactflow_network(payload)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"success": True, **network.as_interface_dict()}


@app.post("/api/system-xml/validate")
async def validate_system_xml(request: Request) -> dict[str, object]:
    report = validate_system_xml_document(await request.body())
    return report.as_dict()


@app.post("/api/system-xml/parse")
async def parse_system_xml(request: Request) -> dict[str, object]:
    report = validate_system_xml_document(await request.body())
    document = _validated_xml_document_or_422(report)
    return {
        "success": True,
        "validation": report.as_dict(),
        "project": document.as_project_data(),
    }


@app.post("/api/system-xml/compile-model")
async def compile_system_xml_model(request: Request) -> dict[str, object]:
    report = validate_system_xml_document(await request.body())
    document = _validated_xml_document_or_422(report)
    project, network = _compile_xml_document_or_422(document)
    return {
        "success": True,
        "validation": report.as_dict(),
        "simulation": pydantic_to_jsonable(project.simulation),
        **network.as_interface_dict(),
    }


@app.post("/api/system-xml/simulate")
async def simulate_system_xml(request: Request) -> dict[str, object]:
    from PythonModels.core.algebraic import AlgebraicSolveError
    from PythonModels.core.solver import SolveIVPConfig
    from PythonModels.core.stream import StreamSolveError
    from PythonModels.systems.generic import (
        GenericFluidSystem,
        SimulationPreparationError,
    )

    report = validate_system_xml_document(await request.body())
    document = _validated_xml_document_or_422(report)
    project, network = _compile_xml_document_or_422(document)
    try:
        system = GenericFluidSystem(network)
        result = system.simulate(
            SolveIVPConfig(
                t_start=project.simulation.t_start,
                t_stop=project.simulation.t_stop,
                method=project.simulation.method,
                max_step=project.simulation.max_step,
            ),
            sample_step=project.simulation.step,
        )
    except SimulationPreparationError as exc:
        raise HTTPException(
            status_code=422,
            detail={
                "message": "The compiled model is not ready for simulation.",
                "issues": [
                    {
                        "severity": "error",
                        "layer": "simulation",
                        **issue.as_dict(),
                    }
                    for issue in exc.issues
                ],
            },
        ) from exc
    except AlgebraicSolveError as exc:
        raise HTTPException(
            status_code=422,
            detail={
                "message": str(exc),
                "issues": [
                    {
                        "severity": "error",
                        "layer": "simulation",
                        "code": "PRESSURE_FLOW_SOLVE_FAILED",
                        "message": str(exc),
                    }
                ],
                "diagnostics": exc.diagnostics.as_dict(),
            },
        ) from exc
    except StreamSolveError as exc:
        raise HTTPException(
            status_code=422,
            detail={
                "message": str(exc),
                "issues": [
                    {
                        "severity": "error",
                        "layer": "simulation",
                        "code": "STREAM_SOLVE_FAILED",
                        "message": str(exc),
                    }
                ],
                "diagnostics": exc.diagnostics.as_dict(),
            },
        ) from exc
    except (RuntimeError, ValueError) as exc:
        raise HTTPException(
            status_code=422,
            detail={
                "message": "Simulation failed while evaluating the compiled model.",
                "issues": [
                    {
                        "severity": "error",
                        "layer": "simulation",
                        "code": "SIMULATION_EXECUTION_FAILED",
                        "message": str(exc),
                    }
                ],
            },
        ) from exc

    return {
        "validation": report.as_dict(),
        "simulation": pydantic_to_jsonable(project.simulation),
        "model": network.as_interface_dict(),
        **result.as_dict(),
    }


def _validated_xml_document_or_422(
    report: SystemXmlValidationReport,
) -> SystemXmlDocument:
    if not report.valid or report.document is None:
        raise HTTPException(
            status_code=422,
            detail={
                "message": "System XML validation failed.",
                "issues": [issue.as_dict() for issue in report.issues],
            },
        )
    return report.document


def _compile_xml_document_or_422(
    document: SystemXmlDocument,
) -> tuple[ReactFlowProjectPayload, "SimulationNetwork"]:
    project = ReactFlowProjectPayload(**document.as_project_data())
    try:
        network = compile_reactflow_network(project)
    except ValueError as exc:
        raise HTTPException(
            status_code=422,
            detail={
                "message": "System XML passed protocol validation but model compilation failed.",
                "issues": [
                    {
                        "severity": "error",
                        "layer": "semantic",
                        "code": "MODEL_COMPILATION_FAILED",
                        "message": str(exc),
                    }
                ],
            },
        ) from exc
    return project, network


def build_reactflow_system_xml(project: ReactFlowProjectPayload) -> bytes:
    system = ET.Element(
        "System",
        {
            "name": project.name,
            "schemaVersion": SYSTEM_XML_SCHEMA_VERSION,
            "unitSystem": SYSTEM_XML_UNIT_SYSTEM,
        },
    )
    ET.SubElement(
        system,
        "Simulation",
        {
            "tStart": str(project.simulation.t_start),
            "tStop": str(project.simulation.t_stop),
            "step": str(project.simulation.step),
            "maxStep": str(project.simulation.max_step),
            "method": project.simulation.method,
        },
    )
    components_node = ET.SubElement(system, "Components")
    connections_node = ET.SubElement(system, "Connections")
    port_index: dict[tuple[str, str], ReactFlowPortDefinition] = {}

    for node in project.nodes:
        component_node = ET.SubElement(
            components_node,
            "Component",
            {
                "id": node.id,
                "name": node.data.label,
                "type": node.data.modelType,
                "componentType": node.data.componentType,
                "x": f"{node.position.x:g}",
                "y": f"{node.position.y:g}",
                "rotation": str(node.data.rotation),
                "mirrored": str(node.data.mirrored).lower(),
            },
        )
        for index, port in enumerate(node.data.ports):
            port_definition = normalize_port_definition(
                port,
                index=index,
                component_type=node.data.componentType,
            )
            port_key = (node.id, port_definition.name)
            if port_key in port_index:
                raise ValueError(
                    f"Component {node.id} contains duplicate port {port_definition.name}."
                )
            port_index[port_key] = port_definition
            port_attributes = {
                "name": port_definition.name,
                "kind": port_definition.kind,
                "domain": port_definition.domain,
                "nominalRole": port_definition.nominalRole,
                "side": port_definition.side,
            }
            if port_definition.kind == "physical":
                port_attributes["positiveFlowDirection"] = (
                    port_definition.positiveFlowDirection or "intoComponent"
                )
            ET.SubElement(component_node, "Port", port_attributes)
        for name, value in node.data.parameters.items():
            ET.SubElement(
                component_node,
                "Parameter",
                {"name": name, "value": str(value)},
            )

    for edge in project.edges:
        first = require_connection_port(
            port_index,
            edge.source,
            edge.sourceHandle,
            edge.id,
        )
        second = require_connection_port(
            port_index,
            edge.target,
            edge.targetHandle,
            edge.id,
        )
        validate_compatible_ports(first, second, edge.id)

        connection_node = ET.SubElement(
            connections_node,
            "Connection",
            {
                "id": edge.id,
                "kind": first.kind,
                "domain": first.domain,
            },
        )
        endpoints = [
            (edge.source, first, None),
            (edge.target, second, None),
        ]
        if first.kind == "signal":
            if first.nominalRole == "input":
                endpoints.reverse()
            endpoints = [
                (endpoints[0][0], endpoints[0][1], "source"),
                (endpoints[1][0], endpoints[1][1], "target"),
            ]
        for component_id, port, role in endpoints:
            attributes = {"component": component_id, "port": port.name}
            if role is not None:
                attributes["role"] = role
            ET.SubElement(connection_node, "Endpoint", attributes)

    ET.indent(system, space="  ")
    return ET.tostring(system, encoding="utf-8", xml_declaration=True)


def normalize_port_definition(
    port: ReactFlowPortDefinition | str,
    *,
    index: int,
    component_type: str,
) -> ReactFlowPortDefinition:
    if isinstance(port, ReactFlowPortDefinition):
        return port

    registered_legacy_ports: dict[str, dict[str, tuple[str, str]]] = {
        "cylinder": {"port_b": ("outlet", "right")},
        "tank": {"port_a": ("inlet", "left")},
        "pipe": {
            "port_a": ("inlet", "left"),
            "port_b": ("outlet", "right"),
        },
        "orifice": {
            "port_a": ("inlet", "left"),
            "port_b": ("outlet", "right"),
        },
        "tee": {
            "port_in": ("bidirectional", "left"),
            "port_out1": ("bidirectional", "right"),
            "port_out2": ("bidirectional", "right"),
        },
    }
    registered = registered_legacy_ports.get(component_type, {}).get(port)
    if registered is not None:
        nominal_role, side = registered
        return ReactFlowPortDefinition(
            name=port,
            nominalRole=nominal_role,
            positiveFlowDirection="intoComponent",
            side=side,
        )

    nominal_role: Literal["inlet", "outlet", "bidirectional"] = "bidirectional"
    if "out" in port or port == "port_b":
        nominal_role = "outlet"
    elif "in" in port or port == "port_a":
        nominal_role = "inlet"
    return ReactFlowPortDefinition(
        name=port,
        nominalRole=nominal_role,
        positiveFlowDirection="intoComponent",
        side="left" if index == 0 else "right",
    )


def require_connection_port(
    port_index: dict[tuple[str, str], ReactFlowPortDefinition],
    component_id: str,
    port_name: str | None,
    connection_id: str,
) -> ReactFlowPortDefinition:
    if port_name is None or (component_id, port_name) not in port_index:
        raise ValueError(
            f"Connection {connection_id} references missing endpoint "
            f"{component_id}.{port_name or '<empty>'}."
        )
    return port_index[(component_id, port_name)]


def validate_compatible_ports(
    first: ReactFlowPortDefinition,
    second: ReactFlowPortDefinition,
    connection_id: str,
) -> None:
    if first.kind != second.kind:
        raise ValueError(f"Connection {connection_id} mixes physical and signal ports.")
    if first.domain != second.domain:
        raise ValueError(f"Connection {connection_id} connects incompatible domains.")
    if first.kind == "signal" and {first.nominalRole, second.nominalRole} != {
        "input",
        "output",
    }:
        raise ValueError(
            f"Signal connection {connection_id} must connect one output to one input."
        )


def compile_reactflow_network(project: ReactFlowProjectPayload) -> "SimulationNetwork":
    from PythonModels.core.medium import IdealGasMedium
    from PythonModels.core.network import SimulationNetwork
    from PythonModels.registry import get_component_model_spec

    medium = IdealGasMedium()
    network = SimulationNetwork(name=project.name)

    for node in project.nodes:
        spec = get_component_model_spec(node.data.modelType)
        parameter_values = {
            parameter.name: parameter_float(node, parameter.name, parameter.default)
            for parameter in spec.parameters
        }
        unknown_parameters = set(node.data.parameters) - set(spec.parameter_by_name)
        if unknown_parameters:
            raise ValueError(
                f"Component '{node.id}' contains unsupported parameters: "
                + ", ".join(sorted(unknown_parameters))
                + "."
            )
        component = spec.create(node.id, medium, parameter_values)
        validate_component_port_interface(node, component.port_definitions)
        network.add_component(component)

    for edge in project.edges:
        network.connect(
            edge.source,
            edge.sourceHandle or "",
            edge.target,
            edge.targetHandle or "",
            connection_id=edge.id,
        )
    return network


def validate_component_port_interface(
    node: ReactFlowNodePayload,
    component_ports: tuple["PortDefinition", ...],
) -> None:
    payload_ports = [
        normalize_port_definition(
            port,
            index=index,
            component_type=node.data.componentType,
        )
        for index, port in enumerate(node.data.ports)
    ]
    expected_by_name = {port.name: port for port in component_ports}
    payload_by_name = {port.name: port for port in payload_ports}
    if len(payload_by_name) != len(payload_ports):
        raise ValueError(f"Component {node.id} contains duplicate port names.")
    if set(payload_by_name) != set(expected_by_name):
        raise ValueError(
            f"Component {node.id} port names do not match model {node.data.modelType}."
        )

    for name, payload_port in payload_by_name.items():
        expected = expected_by_name[name]
        if payload_port.kind != expected.kind or payload_port.domain != expected.domain:
            raise ValueError(f"Component {node.id}.{name} has an incompatible port type.")
        if payload_port.nominalRole != expected.nominal_role:
            raise ValueError(f"Component {node.id}.{name} has an incompatible nominal role.")
        if expected.kind == "physical" and (
            payload_port.positiveFlowDirection or "intoComponent"
        ) != expected.positive_flow_direction:
            raise ValueError(f"Component {node.id}.{name} has an incompatible flow sign.")


def reactflow_project_path(project_id: str) -> Path:
    safe_id = sanitize_project_id(project_id)
    return PROJECT_STORAGE_DIR / f"{safe_id}.json"


def sanitize_project_id(project_id: str) -> str:
    safe_id = "".join(
        character
        for character in project_id.strip()
        if character.isalnum() or character in {"-", "_"}
    )
    if not safe_id:
        raise HTTPException(status_code=400, detail="Project id cannot be empty.")
    return safe_id[:80]


def run_reactflow_testmodel(project: ReactFlowProjectPayload) -> dict[str, object]:
    from PythonModels.core.solver import SolveIVPConfig
    from PythonModels.scripts.run_testmodel import (
        TestModelExecutionConfig,
        TestModelRunConfig,
        TestModelSamplingConfig,
        run_testmodel,
    )
    from PythonModels.systems.testmodel import (
        BranchConfig,
        CylinderConfig,
        OrificeConfig,
        PipeConfig,
        TankConfig,
        TestModelConfig,
    )

    nodes_by_type: dict[str, list[ReactFlowNodePayload]] = {}
    for node in project.nodes:
        nodes_by_type.setdefault(node.data.modelType, []).append(node)

    cylinder = first_node(nodes_by_type, "cylinder")
    tank = first_node(nodes_by_type, "tank")
    orifices = nodes_by_type.get("orifice", [])
    pipes = nodes_by_type.get("pipe", [])

    model_config = TestModelConfig(
        cylinder=CylinderConfig(
            volume=parameter_float(cylinder, "volume", 0.01),
            p0=parameter_float(cylinder, "p0", 35e6),
            T0=parameter_float(cylinder, "T0", 300.0),
        ),
        upper_branch=BranchConfig(
            orifice=OrificeConfig(K=parameter_float(node_at(orifices, 0), "K", 1e-5)),
            pipe=pipe_config_from_node(node_at(pipes, 0), PipeConfig),
        ),
        lower_branch=BranchConfig(
            orifice=OrificeConfig(K=parameter_float(node_at(orifices, 1), "K", 1e-5)),
            pipe=pipe_config_from_node(node_at(pipes, 1), PipeConfig),
        ),
        tank=TankConfig(
            volume=parameter_float(tank, "volume", 0.1),
            p0=parameter_float(tank, "p0", 1e5),
            T0=parameter_float(tank, "T0", 300.0),
        ),
    )
    run_config = TestModelRunConfig(
        model=model_config,
        solver=SolveIVPConfig(
            t_start=project.simulation.t_start,
            t_stop=project.simulation.t_stop,
            method=project.simulation.method,
            max_step=project.simulation.max_step,
        ),
        sampling=TestModelSamplingConfig(step=project.simulation.step),
        execution=TestModelExecutionConfig(use_modelica_reference_if_available=False),
    )
    result = run_testmodel(run_config=run_config)
    series_keys = ("time", "mytank.p", "mytank.T", "mycylinder.p", "mycylinder.T")
    series = {
        key: [float(value) for value in result.series[key]]
        for key in series_keys
        if key in result.series
    }

    return {
        "success": bool(result.solution.success),
        "message": str(result.solution.message),
        "usedModelicaReference": result.used_modelica_reference,
        "final": {
            "time": series["time"][-1],
            "tankPressure": series["mytank.p"][-1],
            "tankTemperature": series["mytank.T"][-1],
            "cylinderPressure": series["mycylinder.p"][-1],
            "cylinderTemperature": series["mycylinder.T"][-1],
        },
        "series": series,
        "artifacts": {
            "primaryCsv": str(result.artifacts.primary_csv_path),
            "temperatureCsv": str(result.artifacts.temperature_csv_path),
            "temperatureSvg": str(result.artifacts.temperature_svg_path),
            "runReport": str(result.artifacts.run_report_path),
        },
        "networkSummary": result.system.network.summary(),
    }


def first_node(
    nodes_by_type: dict[str, list[ReactFlowNodePayload]],
    model_type: str,
) -> ReactFlowNodePayload | None:
    nodes = nodes_by_type.get(model_type, [])
    return nodes[0] if nodes else None


def node_at(
    nodes: list[ReactFlowNodePayload],
    index: int,
) -> ReactFlowNodePayload | None:
    return nodes[index] if index < len(nodes) else None


def parameter_float(
    node: ReactFlowNodePayload | None,
    name: str,
    default: float,
) -> float:
    if node is None:
        return default
    value = node.data.parameters.get(name, default)
    try:
        return float(value)
    except (TypeError, ValueError):
        raise ValueError(f"Parameter '{name}' on component '{node.id}' must be numeric.")


def pipe_config_from_node(node: ReactFlowNodePayload | None, pipe_config_type):
    return pipe_config_type(
        length=parameter_float(node, "length", 5.0),
        diameter=parameter_float(node, "diameter", 0.02),
        lambda_darcy=parameter_float(node, "lambda_darcy", 0.02),
        p0=parameter_float(node, "p0", 1e5),
        T0=parameter_float(node, "T0", 300.0),
    )