# Use the existing connected PyFluent variable solver.
import math
import json

task_directory = r"C:\Users\13977\AppData\Local\FluentAGI\data\runs\20261017-150300-e37701"
s = solver.settings
prefix = task_directory + r"\local_jet_9fd5359e"
velocity = 0.8
dt = 0.001
completed_steps = 0
initial_volume = 1.584e-5
animation_outputs = []

def checked(label, result):
    print(label, result)
    if isinstance(result, dict):
        if result.get("success") is False or result.get("ok") is False:
            raise RuntimeError(label + ": " + str(result))
        if result.get("status") in ["error", "failed"]:
            raise RuntimeError(label + ": " + str(result))
        if result.get("error") or result.get("errors"):
            raise RuntimeError(label + ": " + str(result))
    return result

def object_for(collection, name):
    names = list(collection.get_object_names())
    if name not in names:
        collection.create(name=name)
    if name not in collection.get_object_names():
        raise RuntimeError("Object creation failed: " + name)
    return collection[name]

def save(label):
    path = prefix + "_" + label + ".cas.h5"
    s.file.write_case_data(file_name=path)
    print("SAVED_CASE_DATA", path)

def number(raw, name):
    if isinstance(raw, dict):
        if name in raw:
            value = raw[name]
            while isinstance(value, (list, tuple)) and len(value):
                value = value[0]
            if isinstance(value, (float, int)):
                return float(value)
        for value in raw.values():
            result = number(value, name)
            if result is not None:
                return result
    elif isinstance(raw, (list, tuple)):
        for value in raw:
            result = number(value, name)
            if result is not None:
                return result
    return None

def water_volume():
    raw = s.solution.report_definitions.compute(
        report_defs=["local_water_volume"]
    )
    print("WATER_VOLUME_REPORT", raw)
    result = number(raw, "local_water_volume")
    if result is None or not math.isfinite(result):
        raise RuntimeError("Invalid water-volume report")
    return result

gmsh_code = r"""
xs = [0.0, 0.012, 0.042, 0.060]
ys = [0.0, 0.002, 0.008, 0.014, 0.044, 0.050, 0.056, 0.060]
zs = [0.0, 0.010, 0.020, 0.030]
boxes = []
for i in range(len(xs)-1):
    for j in range(len(ys)-1):
        for k in range(len(zs)-1):
            tag = gmsh.model.occ.addBox(
                xs[i], ys[j], zs[k],
                xs[i+1]-xs[i], ys[j+1]-ys[j], zs[k+1]-zs[k])
            boxes.append((3, tag))
tag = gmsh.model.occ.addBox(-0.012, 0.050, 0.010, 0.012, 0.006, 0.010)
boxes.append((3, tag))
tag = gmsh.model.occ.addBox(0.060, 0.002, 0.010, 0.012, 0.006, 0.010)
boxes.append((3, tag))
gmsh.model.occ.fragment([boxes[0]], boxes[1:])
gmsh.model.occ.synchronize()
tol = 1e-6
groups = {"inlet": [], "outlet": [], "walls": []}
top_count = 0
drain_count = 0
for dim, tag in gmsh.model.getEntities(2):
    up, down = gmsh.model.getAdjacencies(2, tag)
    if len(up) != 1:
        continue
    x, y, z = gmsh.model.occ.getCenterOfMass(2, tag)
    if abs(x+0.012) < tol and 0.050 < y < 0.056 and 0.010 < z < 0.020:
        groups["inlet"].append(tag)
    elif abs(x-0.072) < tol and 0.002 < y < 0.008 and 0.010 < z < 0.020:
        groups["outlet"].append(tag)
        drain_count += 1
    elif abs(y-0.060) < tol and 0.0 < x < 0.060:
        groups["outlet"].append(tag)
        top_count += 1
    else:
        groups["walls"].append(tag)
if top_count == 0 or drain_count == 0:
    raise RuntimeError("Top vent or lower-right drain is missing")
for name in ["inlet", "outlet", "walls"]:
    if not groups[name]:
        raise RuntimeError("Missing physical boundary: " + name)
    group = gmsh.model.addPhysicalGroup(2, groups[name])
    gmsh.model.setPhysicalName(2, group, name)
volumes = [tag for dim, tag in gmsh.model.getEntities(3)]
group = gmsh.model.addPhysicalGroup(3, volumes)
gmsh.model.setPhysicalName(3, group, "fluid")
for dim, tag in gmsh.model.getEntities(1):
    ends = gmsh.model.getBoundary([(1, tag)], oriented=False)
    a = gmsh.model.getValue(0, ends[0][1], [])
    b = gmsh.model.getValue(0, ends[1][1], [])
    distance = math.sqrt(sum((a[i]-b[i])**2 for i in range(3)))
    midpoint = [(a[i]+b[i])/2.0 for i in range(3)]
    spacing = 0.0025
    if abs(a[0]-b[0]) > tol:
        if midpoint[0] < 0.012 or midpoint[0] > 0.042:
            spacing = 0.0015
        else:
            spacing = 0.003
    elif abs(a[1]-b[1]) > tol:
        if midpoint[1] < 0.008 or midpoint[1] > 0.050:
            spacing = 0.0015
        elif midpoint[1] < 0.014 or midpoint[1] > 0.044:
            spacing = 0.002
        else:
            spacing = 0.003
    elif abs(a[2]-b[2]) > tol:
        if 0.010 < midpoint[2] < 0.020:
            spacing = 0.00125
    intervals = max(1, int(round(distance/spacing)))
    gmsh.model.mesh.setTransfiniteCurve(tag, intervals+1)
for dim, tag in gmsh.model.getEntities(2):
    gmsh.model.mesh.setTransfiniteSurface(tag)
    gmsh.model.mesh.setRecombine(2, tag)
for tag in volumes:
    gmsh.model.mesh.setTransfiniteVolume(tag)
gmsh.option.setNumber("Mesh.ElementOrder", 1)
gmsh.write(output_dir + "/tank_inlet_drain.step")
gmsh.model.mesh.generate(3)
types, element_tags, nodes = gmsh.model.mesh.getElements(3)
count = 0
for i in range(len(types)):
    if int(types[i]) != 5:
        raise RuntimeError("Non-hexahedral volume cell detected")
    count += len(element_tags[i])
expected_volume = 0.060*0.060*0.030 + 2*0.012*0.006*0.010
actual_volume = sum(gmsh.model.occ.getMass(3, tag) for tag in volumes)
if abs(actual_volume-expected_volume)/expected_volume > 1e-6:
    raise RuntimeError("Fluid geometry volume validation failed")
print("VERIFIED_HEX8_CELLS", count)
print("VERIFIED_FLUID_VOLUME_M3", actual_volume)
print("TOP_VENT_FACES", top_count)
print("LOWER_RIGHT_DRAIN_FACES", drain_count)
gmsh.write(output_dir + "/local_hex_native.msh")
"""

if s.setup.models.is_active():
    print("PREVIOUS_MODEL", s.setup.models.multiphase.get_state())
    print("PREVIOUS_RUN_STATE", s.solution.run_calculation.get_state())
    save("before_geometry_change")

mesh = checked("LOCAL_HEX_MESH", fluentagi.geometry_mesh(
    directory=task_directory, code=gmsh_code
))
if not isinstance(mesh, dict) or not mesh.get("mesh_file"):
    raise RuntimeError("Geometry adapter did not return a Fluent mesh_file")
s.file.read_mesh(file_name=mesh["mesh_file"])
s.mesh.check()

checked("WATER_AIR_SETUP", fluentagi.configure_water_air(
    directory=task_directory,
    velocity=velocity,
    time_step=dt,
    gravity_y=-9.81,
    inlet="inlet",
    outlet="outlet"
))
checked("SURFACE_TENSION", fluentagi.configure_surface_tension(
    directory=task_directory, coefficient=0.072,
    phase_a="phase-1", phase_b="phase-2"
))
checked("PHASES", fluentagi.inspect_phases())
print("BOUNDARIES", s.setup.boundary_conditions.get_state())

scheme = s.solution.methods.p_v_coupling.flow_scheme
for candidate in ["PISO", "piso"]:
    if candidate in scheme.allowed_values():
        scheme.set_state(candidate)
        break
print("SOLUTION_METHODS", s.solution.methods.get_state())

checked("INITIALIZE_NEW_GEOMETRY", fluentagi.initialize_air(
    directory=task_directory
))
zones = list(s.setup.cell_zone_conditions.fluid.get_object_names())
report = object_for(s.solution.report_definitions.volume, "local_water_volume")
report.set_state({
    "report_type": "volume-integral",
    "field": "phase-2-vof",
    "cell_zones": zones
})
print("INITIAL_WATER_PATCH", s.solution.initialization.patch.calculate_patch(
    domain="phase-2", cell_zones=zones, registers=[],
    variable="mp", reference_frame="Relative to Cell Zone",
    use_custom_field_function=False, custom_field_function_name="",
    value="IF(x < 0 [m], 1, IF(y < 0.008 [m], 1, 0))"
))
actual = water_volume()
if abs(actual-initial_volume)/initial_volume > 0.03:
    save("initial_water_check_failed")
    raise RuntimeError("Initial tank and pipe water-volume validation failed")
save("initial")

contour = object_for(s.results.graphics.contour, "local_vof_check")
if "phase-2-vof" not in contour.field.allowed_values():
    raise RuntimeError("Water volume-fraction field is unavailable")
contour.field = "phase-2-vof"
iso = s.results.surfaces.iso_surface
object_for(iso, "local_midplane").set_state({
    "field": "z-coordinate", "iso_values": [0.015]
})
object_for(iso, "local_free_surface").set_state({
    "field": "phase-2-vof", "iso_values": [0.5]
})
surfaces = ["local_midplane", "local_free_surface"]
s.solution.run_calculation.transient_controls.time_step_size = dt
s.solution.run_calculation.transient_controls.max_iter_per_time_step = 15

try:
    for label, frames in [("pre_0p02s", 3), ("final_0p25s", 24)]:
        result = fluentagi.record_animation(
            directory=task_directory, field="phase-2-vof",
            frames=frames, steps_per_frame=10,
            max_iter_per_step=15, surfaces=surfaces
        )
        checked("ANIMATION_" + label, result)
        animation_outputs.append(result)
        completed_steps += (frames-1)*10
        save(label)
        actual = water_volume()
        upper_bound = initial_volume + 0.000048*completed_steps*dt
        print("WATER_INVENTORY", {
            "nominal_time_s": completed_steps*dt,
            "water_volume_m3": actual,
            "initial_plus_injected_volume_m3": upper_bound,
            "difference_includes_outlet_transport_m3": upper_bound-actual
        })
        print("RUN_STATE", s.solution.run_calculation.get_state())
        if actual < 0 or actual > upper_bound*1.05:
            raise RuntimeError("Water inventory failed basic bounds check")
    for field in ["phase-2-vof", "velocity-magnitude"]:
        checked("FINAL_EXPORT_" + field, fluentagi.export_results(
            directory=task_directory, fields=[field], surfaces=surfaces
        ))
except Exception as exc:
    print("STOPPED", str(exc))
    print("CONFIRMED_BATCH_STEPS", completed_steps)
    try:
        save("interrupted")
    except Exception as save_error:
        print("SAVE_ERROR", str(save_error))
    raise

print("ANIMATION_OUTPUTS", animation_outputs)
print("FINAL_SUMMARY", {
    "tank_dimensions_m": [0.060, 0.060, 0.030],
    "inlet_pipe_dimensions_m": [0.012, 0.006, 0.010],
    "drain_pipe_dimensions_m": [0.012, 0.006, 0.010],
    "drain_y_range_m": [0.002, 0.008],
    "inlet_velocity_m_s": velocity,
    "time_step_s": dt,
    "nominal_physical_time_s": completed_steps*dt,
    "water_volume_m3": actual,
    "task_directory": task_directory
})
print("Outlet physical group combines the atmospheric top vent and drain end.")
print("Inventory bounds do not constitute a measured phase-flux balance.")
print("Verify residuals, actual flow time, wall impact and drainage in animation.")
print("Wall-clock duration is machine-dependent; no hard timeout is configured.")