import bpy
from . import utils
from ..utils import caches

solver_object_types = {'LAGRANGIAN_SOLVER', 'MPM_SOLVER'}
simulation_object_types = {
            'AFFINE_BODY', 'CLOTH', 'ELASTIC_ROD', 'DEFORMABLE_BODY', 'MPM_MESHLESS_PARTICLES', 'COLLIDER', 'SOFT_COLLIDER'
        }
def draw_sim_control(layout, context):
    obj = context.object
    lay = layout

    # Create Solver Buttons (always show if in VIEW_3D)
    if context.area.type == 'VIEW_3D':
        row = lay.row(align=True)
        row.operator("hi_phy.create_lagrangian_solver", text="New Lagrangian", icon='PLUS')
        row.operator("hi_phy.create_mpm_solver", text="New MPM", icon='PLUS')
        lay.separator()

    # Object List (Show all objects with solvers attached)
    box = lay.box()
    box.label(text="Objects with Solvers", icon='INFO')

    found_any = False
    for o in context.scene.objects:
        if o.hi_phy.is_active and o.hi_phy.solver:
            found_any = True
            row = box.row(align=True)
            # Use a split to show object and its solver
            split = row.split(factor=0.5)
            split.label(text=o.name, icon='OBJECT_DATA')
            split.label(text=f"-> {o.hi_phy.solver.name}", icon='PHYSICS')
            # Select button
            op = row.operator("hi_phy.select_object", text="", icon='RESTRICT_SELECT_OFF')
            op.object_name = o.name

    if not found_any:
        box.label(text="No objects with solvers attached")

    lay.separator()

    # check if any simulation is running
    active_solver = None
    target_solver = None
    for solver_name in caches:
        if (caches[solver_name].IsSimulationRunning()):
            active_solver = context.scene.objects[solver_name]

    # Controls for active object (if it's a participant or solver)
    if obj and hasattr(obj, "hi_phy") and obj.hi_phy.is_active:
        hi_phy = obj.hi_phy

        is_solver = hi_phy.object_type in solver_object_types
        is_participant = hi_phy.object_type in simulation_object_types

        if is_participant or is_solver:
            lay.label(text=f"Active: {obj.name} ({hi_phy.object_type})")

            # Always show solver selection for participants
            if is_participant:
                utils.draw_prop(lay, hi_phy, 'solver', 'Solver', expand=False, use_column=True)

            lay.separator()

            target_solver = obj if is_solver else hi_phy.solver

    col = lay.column()
    if active_solver or target_solver:
        if not target_solver:
            target_solver = active_solver
            # Frame Range (from target solver)
        row = col.row(align=True)
        row.prop(target_solver.hi_phy, "frame_range_simulation_start", text="Start")
        row.prop(target_solver.hi_phy, "frame_range_simulation_end", text="End")
        col.separator()

        # Simulation Buttons
        is_mpm = target_solver and target_solver.hi_phy.object_type == 'MPM_SOLVER'

        row = col.row(align=True)
        if is_mpm:
            op = row.column(align=True).operator('hi_phy.run_mpm_sim', text='Run Sim')
            op.active_solver = target_solver.name
        else:
            op = row.column(align=True).operator('hi_phy.run_sim', text='Run Sim')
            if target_solver:
                op.active_solver = target_solver.name
            op = row.column(align=True).operator('hi_phy.run_sim_interactive', text='Run Sim Interactive')
            if target_solver:
                op.active_solver = target_solver.name

        row = col.row(align=True)
        op_stop = row.column(align=True).operator('hi_phy.stop_sim', text='Stop Sim')
        if target_solver:
            op_stop.active_solver = target_solver.name

        row = col.row(align=True)
        op_clear = row.column(align=True).operator('hi_phy.clear_cache', text='Clear Cache')
        if target_solver:
            op_clear.solver_name = target_solver.name
        row = col.row(align=True)
        op_apply = row.column(align=True).operator('hi_phy.apply_sim_result', text='Clear Cache and Apply Simulation Result')
        if target_solver:
            op_apply.active_solver = target_solver.name

        # Progress Bar (Show only if solver exists and is running/has progress)
        if target_solver and target_solver.name in caches:
            cache = caches[target_solver.name]
            [finished_frames, total_frames] = cache.GetProgress()
            if total_frames > 0:
                row = lay.row(align=True)
                progress = finished_frames / total_frames
                row.progress(text=f'Progress: {finished_frames}/{total_frames}', factor=progress)
            elif cache.IsSimulationRunning():
                row = lay.row(align=True)
                row.label(text="Simulation running...", icon='PLAY')
    else:
        # Main controls container
        col.enabled = False
        lay.label(text="Please select a solver to enable controls", icon='ERROR')
        row = col.row(align=True)
        row.label(text="Start: -")
        row.label(text="End: -")

class HI_PHY_PT_SimulationControl_Properties(bpy.types.Panel):
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "physics"
    bl_label = "Hi Phy: Simulation Control"
    bl_options = {'DEFAULT_CLOSED'}

    @classmethod
    def poll(cls, context):
        obj = context.object
        if not obj or not hasattr(obj, "hi_phy"):
            return False
        hi_phy = obj.hi_phy
        return hi_phy.is_active and hi_phy.object_type in simulation_object_types

    def draw(self, context):
        draw_sim_control(self.layout, context)

class HI_PHY_PT_SimulationControl_Viewport(bpy.types.Panel):
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "HiPhy"
    bl_label = "Simulation Control"

    @classmethod
    def poll(cls, context):
        return True

    def draw(self, context):
        draw_sim_control(self.layout, context)
