from xic_loader import XICProgram
from xic_ops import XICContext, OP_TABLE


class XICRuntimeError(Exception):
    pass


def run_xic_program(prog: XICProgram) -> XICContext:
    """Execute an XIC program.

    Creates a context, iterates through instructions, dispatches
    each operation to the handler in OP_TABLE, and returns the final context.

    Supports control flow via chain queue (_chain_queue):
    - When a control op (IF, MATCH, LOOP) enqueues a label,
      the executor searches for a CHAIN instruction with that label
      and continues execution from there.
    - Tracks total_steps for guardrail enforcement.
    - Stops if guardrails are triggered (max_total_steps, etc).
    """
    ctx = XICContext()
    max_total_steps = int(ctx.params.get("max_total_steps", 10000))
    instr_ptr = 0
    total_steps = 0

    while instr_ptr < len(prog.instructions) and total_steps < max_total_steps:
        # Check for guardrails
        if "guardrails" in ctx._state and ctx._state["guardrails"]:
            print(f"[XIC-VM] Guardrail triggered, stopping execution")
            break

        # Check for scheduled chain
        next_chain = ctx.pop_next_chain()
        if next_chain:
            # Find CHAIN instruction with matching label
            found = False
            for i, instr in enumerate(prog.instructions):
                if instr.op == "CHAIN" and instr.args and str(instr.args[0]) == next_chain:
                    instr_ptr = i
                    found = True
                    print(f"[XIC-VM] Jumping to chain: {next_chain}")
                    break
            if not found:
                print(f"[XIC-VM] Warning: chain '{next_chain}' not found, continuing")

        # Execute current instruction
        instr = prog.instructions[instr_ptr]
        op_name = instr.op

        if op_name not in OP_TABLE:
            raise XICRuntimeError(f"Unknown operation at instruction {instr_ptr}: {op_name}")

        op_handler = OP_TABLE[op_name]

        try:
            op_handler(ctx, *instr.args)
            total_steps += 1
            ctx._state["total_steps"] = total_steps
        except Exception as e:
            raise XICRuntimeError(f"Operation {op_name} failed at instruction {instr_ptr}: {e}")

        instr_ptr += 1

    if total_steps >= max_total_steps:
        ctx._state.setdefault("guardrails", []).append("max_total_steps_exceeded")
        print(f"[XIC-VM] Max total steps exceeded: {total_steps}")

    return ctx