"""Lower reviewed component contracts to a static C evaluation schedule.

The compact storage-anchored schedule and the extended catalog schedule both
produce standalone C numerics, without Python callbacks or numerical fallback.
"""
from __future__ import annotations

from dataclasses import dataclass
import json
from math import isfinite

from app.simulation.core.metadata import ResultVariableMetadata
from app.simulation.systems.network import SimulationNetwork
from .contracts import SUPPORTED_TYPES, SUPPORTED_VERSIONS
from .tolerances import state_absolute_tolerance


class NativeCapabilityError(ValueError):
    """The complete model cannot be represented by this native backend."""


_STORAGE_ANCHORED_TYPES = frozenset(
    "amesim_" + name for name in (
        "pnch023", "pnch012", "pnvo001", "pnpl01", "step0", "ud00",
        "forc", "pnrp17", "mecmas21", "f000", "lstp00a",
    )
)


@dataclass(frozen=True)
class NativeProgram:
    source: str
    header: str
    state_keys: tuple[str, ...]
    variables: tuple[ResultVariableMetadata, ...]
    component_types: tuple[str, ...]
    evaluation_schedule: dict | None = None

    def manifest(self) -> dict:
        return {
            "abiVersion": 1, "stateKeys": self.state_keys,
            "variables": [v.as_dict() for v in self.variables],
            "componentTypes": self.component_types,
            "componentVersions": {name: SUPPORTED_VERSIONS[name] for name in self.component_types},
            "jacobianPolicy": "CVODE default; no custom Jacobian",
            "evaluationSchedule": self.evaluation_schedule or {"strategy": "storage-anchored", "cyclicBlockCount": 0},
        }


class _Groups:
    def __init__(self, items):
        self.parent = {x: x for x in items}

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        self.parent[self.find(b)] = self.find(a)


def _number(value):
    value = float(value)
    if not isfinite(value):
        raise NativeCapabilityError("Native constants must be finite.")
    return repr(value)


def compile_native_program(network: SimulationNetwork) -> NativeProgram:
    network.validate_port_supplies()
    from .extended import catalog_contracts
    contracts = catalog_contracts()
    for component in network.components.values():
        if type(component) is not contracts.get(component.model_type):
            raise NativeCapabilityError(f'{component.name}: no native contract for {component.model_type}')
        if component.MODEL_VERSION != SUPPORTED_VERSIONS.get(component.model_type):
            raise NativeCapabilityError(f'{component.name}: native kernel version does not match the component contract')
    # Preserve the compact, verified schedule for its supported topology.
    # Both paths emit C; this never falls back to Python numerical execution.
    try:
        return _compile_storage_anchored_program(network)
    except NativeCapabilityError:
        from .extended import compile_extended_program
        return compile_extended_program(network)


def _compile_storage_anchored_program(network: SimulationNetwork) -> NativeProgram:
    components = list(network.components.values())
    if not components:
        raise NativeCapabilityError("Native simulation requires a dynamic model.")
    for c in components:
        if c.model_type not in _STORAGE_ANCHORED_TYPES:
            raise NativeCapabilityError(f"{c.name}: unsupported native component {c.model_type}.")
        if c.MODEL_VERSION != SUPPORTED_VERSIONS[c.model_type]:
            raise NativeCapabilityError(f"{c.name}: native kernel version does not match the component contract.")
        if c.model_type == "amesim_pnch023" and c.kth*c.sth != 0:
            raise NativeCapabilityError(f"{c.name}: native v1 supports adiabatic PNCH023 only.")
        medium = getattr(c, "medium", None)
        if medium is not None and (
            getattr(medium, "SUBSTANCE_ID", None),
            getattr(medium, "PROPERTY_METHOD_ID", None),
        ) != ("helium", "peng_robinson"):
            raise NativeCapabilityError(f"{c.name}: native v1 requires helium / Peng-Robinson.")
        if c.model_type == "amesim_mecmas21":
            if int(c.stoptype) not in (1, 4) or any(
                getattr(c, key) != 0 for key in ("fcoul", "fstick", "rvisc", "wind", "theta")
            ):
                raise NativeCapabilityError(f"{c.name}: native v1 supports friction-free masses with stoptype 1 or 4.")
        if c.model_type == "amesim_lstp00a" and int(c.stiffmode) != 1:
            raise NativeCapabilityError(f"{c.name}: native v1 requires explicit contact stiffness.")
    ports = {
        (c.name, p.name): p for c in components for p in c.active_port_definitions
    }
    adjacent = {}
    groups = _Groups(ports)
    for edge in network.connections:
        a, b = (p.key for p in edge.endpoints)
        # Signal fan-out is allowed; physical ports have one external connection.
        if edge.kind == "physical":
            if a in adjacent or b in adjacent:
                raise NativeCapabilityError("Native physical ports require one connection each.")
            adjacent[a], adjacent[b] = b, a
            groups.union(a, b)
        else:
            source, target = (a, b) if ports[a].nominal_role == "output" else (b, a)
            adjacent[target] = source
    for c in components:
        for name in c.required_connection_ports:
            if (c.name, name) not in adjacent:
                raise NativeCapabilityError(f"{c.name}.{name}: unconnected physical port.")
        names = [p.name for p in c.active_port_definitions if p.domain == "pneumatic"]
        if c.model_type in ("amesim_pnch023", "amesim_pnch012"):
            for name in names[1:]:
                groups.union((c.name, names[0]), (c.name, name))
        if c.model_type == "amesim_mecmas21":
            groups.union((c.name, "port_1"), (c.name, "port_2"))
        if c.model_type == "amesim_pnrp17":
            groups.union((c.name, "port_2"), (c.name, "port_5"))
            groups.union((c.name, "port_3"), (c.name, "port_4"))

    chambers = [c for c in components if c.model_type in ("amesim_pnch023", "amesim_pnch012")]
    masses = [c for c in components if c.model_type == "amesim_mecmas21"]
    chamber_by_group, mass_by_group = {}, {}
    for items, mapping in ((chambers, chamber_by_group), (masses, mass_by_group)):
        for c in items:
            root = groups.find((c.name, "port_1"))
            if root in mapping:
                raise NativeCapabilityError(f"{c.name}: coupled storage/mass reduction is outside native v1.")
            mapping[root] = c
    for ep, port in ports.items():
        mapping = chamber_by_group if port.domain == "pneumatic" else mass_by_group
        if port.kind == "physical" and groups.find(ep) not in mapping:
            raise NativeCapabilityError(f"{ep}: no unique storage/mass anchor; native v1 cannot close this block.")

    variables = tuple(v for c in components for v in c.result_variable_metadata())
    slots = {v.key: i for i, v in enumerate(variables)}
    assigned = set()
    lines, init, declarations = [], [], []
    state_keys, state_index = [], {}
    for c in components:
        fields = ("m", "U") if c in chambers else (("v", "x") if c in masses else ())
        for field in fields:
            key = f"{c.name}.{field}"
            state_index[key] = len(state_keys)
            state_keys.append(key)
    if not state_keys:
        raise NativeCapabilityError("Native v1 requires continuous states.")
    if len(state_keys) > 256 or len(variables) > 8192:
        raise NativeCapabilityError("Native v1 supports at most 256 states and 8192 outputs per model.")

    def w(key):
        return f"w[{slots[key]}]"

    def key(c, field):
        return f"{c.name}.{field}"

    def get(c, field):
        return w(key(c, field))

    def put(c, field, expression):
        k = key(c, field)
        lines.append(f"{w(k)} = {expression};")
        assigned.add(k)

    def si(c, field):
        return state_index[key(c, field)]

    def mass_at(c, port):
        return mass_by_group[groups.find((c.name, port))]

    def chamber_at(c, port):
        return chamber_by_group[groups.find((c.name, port))]

    for c in masses:
        init.extend([f"y[{si(c, 'v')}] = {_number(c.v0)};", f"y[{si(c, 'x')}] = {_number(c.x0)};"])
        for field in ("v", "x"):
            put(c, field, f"y[{si(c, field)}]")
    for (cid, pname), port in ports.items():
        if port.domain == "mechanical":
            m = mass_by_group[groups.find((cid, pname))]
            for field in ("v", "x"):
                k = f"{cid}.{pname}.{field}"
                lines.append(f"{w(k)} = y[{si(m, field)}];")
                assigned.add(k)

    signal_specs = []
    for c in components:
        if c.model_type == "amesim_step0":
            put(c, "y", f"t < {_number(c.time)} ? {_number(c.initial)} : {_number(c.final)}")
            signal_specs.append(("step", c))
        elif c.model_type == "amesim_ud00":
            index = len(signal_specs)
            declarations.append(f"static const double signal_{index}[24] = {{" + ",".join(
                _number(v) for values in (c.starts, c.ends, c.durations) for v in values
            ) + "};")
            put(c, "y", f"native_signal(t, {_number(c.tstart)}, {c.nstages}, {int(c.iscyclic)}, signal_{index})")
            signal_specs.append((index, c))
        else:
            continue
        put(c, "out.signal", get(c, "y"))
    for c in components:
        for port in c.active_port_definitions:
            if port.kind == "signal" and port.nominal_role == "input":
                ep = (c.name, port.name)
                if ep not in adjacent:
                    # PNVO's explicit unconnected opening is a supported default.
                    if c.model_type == "amesim_pnvo001":
                        put(c, port.name + ".signal", _number(c.opening0))
                        continue
                    raise NativeCapabilityError(f"{ep}: missing signal input.")
                source = adjacent[ep]
                source_key = f"{source[0]}.{source[1]}.signal"
                if source_key not in assigned:
                    raise NativeCapabilityError(f"{ep}: unsupported signal dependency.")
                put(c, port.name + ".signal", w(source_key))

    pistons = [c for c in components if c.model_type == "amesim_pnrp17"]
    for c in pistons:
        put(c, "length", f"{_number(c.x0)} + {get(c, 'port_5.x')} - {get(c, 'port_4.x')}")
        put(c, "volume", f"{_number(c.effective_area)} * {get(c, 'length')}")
        put(c, "volume_flow", f"{_number(c.effective_area)} * ({get(c, 'port_5.v')} - {get(c, 'port_4.v')})")
        if chamber_at(c, "port_1").model_type != "amesim_pnch012":
            raise NativeCapabilityError(f"{c.name}: moving volume requires PNCH012.")
    volume_expressions = {}
    for gi, c in enumerate(chambers):
        connected = [p for p in pistons if chamber_at(p, "port_1") is c]
        base = c.cvol if c.model_type == "amesim_pnch023" else c.cvol0 + sum(c.external_volumes.values())
        volume = _number(base)
        rate = "0.0"
        if c.model_type == "amesim_pnch012":
            volume += "".join(" + " + get(p, "volume") for p in connected)
            put(c, "vol", f"fmax({_number(c.cvol0 / 100)}, {volume})")
            rate = _number(sum(c.external_volume_rates.values())) + "".join(" + " + get(p, "volume_flow") for p in connected)
            put(c, "dvol", f"{get(c, 'vol')} <= {_number(c.cvol0 / 100)} ? 0.0 : ({rate})")
            volume, rate = get(c, "vol"), get(c, "dvol")
        volume_expressions[c.name] = (volume, rate)
        # Match the Python constructor: m/U use configured storage volume;
        # connected moving volumes subsequently change recovered p/T.
        initial_volume = max(base, c.cvol0 / 100) if c.model_type == "amesim_pnch012" else base
        init.append(f"if (!native_gas_init({_number(c.p0)}, {_number(c.T0)}, {_number(initial_volume)}, &y[{si(c, 'm')}])) return 0;")
        lines.append(f"NativeGas gas_{gi};")
        lines.append(f"if (!native_gas(y[{si(c, 'm')}], y[{si(c, 'U')}], {volume}, &gas_{gi})) return 0;")
        for field in ("m", "U"):
            put(c, field, f"y[{si(c, field)}]")
        for field in ("p", "T", "rho", "u", "h"):
            put(c, field, f"gas_{gi}.{field}")
    for (cid, pname), port in ports.items():
        if port.domain == "pneumatic":
            c = chamber_by_group[groups.find((cid, pname))]
            for field, expr in (("p", get(c, "p")), ("h_outflow", get(c, "h")), ("m_flow", "0.0")):
                k = f"{cid}.{pname}.{field}"
                lines.append(f"{w(k)} = {expr};")
                assigned.add(k)
    for c in components:
        if c.model_type != "amesim_pnvo001":
            continue
        a, b = chamber_at(c, "port_2"), chamber_at(c, "port_3")
        put(c, "xv", f"fmax(0.0, fmin(1.0, {get(c, 'res.signal')}))")
        lines.append(f"if (!native_orifice({get(a, 'p')}, {get(b, 'p')}, {get(a, 'h')}, {get(b, 'h')}, {_number(c.effective_cq * c.maximum_area)}, {get(c, 'xv')}, &{get(c, 'port_2.m_flow')}, &{get(c, 'cm')}, &{get(c, 'gasvel')})) return 0;")
        assigned.update((key(c, "cm"), key(c, "gasvel")))
        put(c, "port_3.m_flow", f"-{get(c, 'port_2.m_flow')}")
        put(c, "port_2.h_outflow", get(b, "h"))
        put(c, "port_3.h_outflow", get(a, "h"))
        for pname in ("port_2", "port_3"):
            other = adjacent[c.name, pname]
            # Resistance-to-resistance streams need the extended pressure/enthalpy closure.
            if network.components[other[0]] not in chambers:
                raise NativeCapabilityError(f"{c.name}: native v1 requires valve ports directly connected to storage.")
            lines.append(f"{w(f'{other[0]}.{other[1]}.m_flow')} = -{get(c, pname + '.m_flow')};")

    force_known = set()
    def force(c, port, expr):
        put(c, port + ".f", expr)
        force_known.add((c.name, port))
    equations = []
    for c in components:
        if c.model_type == "amesim_forc":
            put(c, "force", f"{_number(c.direction)} * {get(c, 'res.signal')}")
            force(c, "port_2", f"-{get(c, 'force')}")
        elif c.model_type == "amesim_f000":
            force(c, "port_1", "0.0")
        elif c.model_type == "amesim_lstp00a":
            put(c, "gap", f"{_number(c.gap0)} + ({get(c, 'port_2.x')} - {get(c, 'port_1.x')})")
            put(c, "penetration", f"fmax(-{get(c, 'gap')}, 0.0)")
            put(c, "force", f"native_contact({get(c, 'penetration')}, {get(c, 'port_1.v')} - {get(c, 'port_2.v')}, {_number(c.kcont)}, {_number(c.rcont)}, {_number(c.Pdis)}, {int(c.discContactOption)})")
            force(c, "port_1", get(c, "force"))
            force(c, "port_2", f"-{get(c, 'force')}")
        elif c.model_type == "amesim_pnrp17":
            put(c, "pressure_force", f"({get(c, 'port_1.p')} - 101300.0) * {_number(c.effective_area)}")
            equations.extend([
                ((c.name, "port_2"), (c.name, "port_5"), f"-{get(c, 'pressure_force')}"),
                ((c.name, "port_3"), (c.name, "port_4"), get(c, "pressure_force")),
            ])
    for edge in network.connections:
        if edge.domain == "mechanical":
            equations.append((*[p.key for p in edge.endpoints], "0.0"))
    pending = equations
    while pending:
        remaining = []
        for a, b, total in pending:
            if a in force_known and b in force_known:
                raise NativeCapabilityError("Overconstrained native force balance.")
            if a not in force_known and b not in force_known:
                remaining.append((a, b, total))
                continue
            target, source = (b, a) if a in force_known else (a, b)
            c = network.components[target[0]]
            force(c, target[1], f"({total}) - {w(f'{source[0]}.{source[1]}.f')}")
        if len(remaining) == len(pending):
            raise NativeCapabilityError("Native v1 cannot resolve this mechanical force loop.")
        pending = remaining
    for c in masses:
        expression = f"({get(c, 'port_1.f')} + {get(c, 'port_2.f')}) / {_number(c.mass)}"
        put(c, "a", expression)
        lines.append(f"dy[{si(c, 'v')}] = {get(c, 'a')}; dy[{si(c, 'x')}] = {get(c, 'v')};")
        if int(c.stoptype) == 1:
            lines.append(f"native_stop_motion({get(c, 'x')}, {get(c, 'v')}, {_number(c.xmin)}, {_number(c.xmax)}, &dy[{si(c, 'v')}], &dy[{si(c, 'x')}]);")
            put(c, "a", f"dy[{si(c, 'v')}]")
        for field in ("Fvisc", "Ffric", "Fmin", "Fmax"):
            put(c, field, "0.0")
    for c in chambers:
        mass_terms, energy_terms = [], []
        for port in c.active_port_definitions:
            q = get(c, port.name + ".m_flow")
            other = adjacent[c.name, port.name]
            inlet_h = w(f"{other[0]}.{other[1]}.h_outflow")
            mass_terms.append(q)
            energy_terms.append(f"{q} * ({q} > 0.0 ? {inlet_h} : {get(c, 'h')})")
        volume, rate = volume_expressions[c.name]
        energy_terms.extend([f"{_number(c.kth*c.sth)} * ({_number(c.extemp)} - {get(c, 'T')})", f"-{get(c, 'p')} * ({rate})"])
        lines.extend([f"dy[{si(c, 'm')}] = " + " + ".join(mass_terms) + ";", f"dy[{si(c, 'U')}] = " + " + ".join(energy_terms) + ";"])
    missing = set(slots) - assigned
    if missing:
        raise NativeCapabilityError(f"Native output coverage is incomplete: {sorted(missing)}")

    stops = [(si(c, "v"), c.xmin, c.xmax) for c in masses if int(c.stoptype) == 1]
    stop_c = ",".join(f"{{{v},{_number(lo)},{_number(hi)},0,0,0,0}}" for v, lo, hi in stops) or "{0,0,0,0,0,0,0}"
    next_event = []
    for index, c in signal_specs:
        if index == "step":
            next_event.append(f"if (t < {_number(c.time)}) result = fmin(result, {_number(c.time)});")
        else:
            next_event.append(f"result = fmin(result, native_signal_break(t, end, {_number(c.tstart)}, {c.nstages}, {int(c.iscyclic)}, signal_{index}));")
    source = '\n'.join([
        '#include "model.h"', '#include <math.h>', *declarations,
        f"const NativeStop model_stops[{max(1,len(stops))}] = {{{stop_c}}};",
        f"const double model_atol[NSTATES] = {{{','.join(map(state_absolute_tolerance, state_keys))}}};",
        "const char *const model_output_keys[NOUTPUTS] = {" + ",".join(json.dumps(v.key, ensure_ascii=True) for v in variables) + "};",
        "int model_init(double *y) {", *init, "return 1; }",
        "int model_eval(double t, const double *y, double *dy, double *w) {", "(void)t;", *lines,
        "for (int i=0;i<NSTATES;i++) if (!isfinite(dy[i])) return 0;",
        "for (int i=0;i<NOUTPUTS;i++) if (!isfinite(w[i])) return 0;",
        "return 1; }",
        "double model_next_break(double t, double end) { (void)t; double result=end;", *next_event,
        "return result; }", "",
    ])
    header = f'''#ifndef GENERATED_NATIVE_MODEL_H
#define GENERATED_NATIVE_MODEL_H
#include "kernels.h"
#define NSTATES {len(state_keys)}
#define NOUTPUTS {len(variables)}
#define NSTOPS {len(stops)}
extern const NativeStop model_stops[{max(1,len(stops))}];
extern const double model_atol[NSTATES];
extern const char *const model_output_keys[NOUTPUTS];
int model_init(double *y);
int model_eval(double t, const double *y, double *dy, double *w);
double model_next_break(double t, double end);
#endif
'''
    return NativeProgram(source, header, tuple(state_keys), variables, tuple(sorted({c.model_type for c in components})))