from dataclasses import dataclass, field
from typing import Dict, Any, Optional
from runtime_executor.runner import execute_gx, ExecutionError


@dataclass
class XICContext:
    model_path: Optional[str] = None
    mode: str = "chat"
    params: Dict[str, Any] = field(default_factory=dict)
    _state: Dict[str, Any] = field(default_factory=dict)


def op_LOAD_MODEL(ctx: XICContext, *args):
    """LOAD_MODEL <path>: Load a .gx model file."""
    if not args:
        raise ValueError("LOAD_MODEL requires a path argument")
    model_path = args[0]
    ctx.model_path = str(model_path)
    print(f"[XIC] Model loaded: {ctx.model_path}")


def op_SET_MODE(ctx: XICContext, *args):
    """SET_MODE <mode>: Set execution mode (chat, eval, benchmark)."""
    if not args:
        raise ValueError("SET_MODE requires a mode argument")
    ctx.mode = str(args[0])
    print(f"[XIC] Mode set to: {ctx.mode}")


def op_SET_PARAM(ctx: XICContext, *args):
    """SET_PARAM <key> <value>: Set a parameter."""
    if len(args) < 2:
        raise ValueError("SET_PARAM requires key and value arguments")
    key = str(args[0])
    value = args[1]
    ctx.params[key] = value
    print(f"[XIC] Parameter {key} = {value}")


def op_RUN_PROMPT(ctx: XICContext, *args):
    """RUN_PROMPT <prompt>: Execute prompt against loaded model.

    This is the critical operation that wires directly to execute_gx()
    in runtime_executor/runner.py. It reads the .gx binary, decompresses it,
    and executes it with the given prompt.
    """
    if not args:
        raise ValueError("RUN_PROMPT requires a prompt argument")

    if not ctx.model_path:
        raise ValueError("No model loaded. Use LOAD_MODEL first.")

    prompt = str(args[0])

    try:
        # Call the real compressed model runner
        # execute_gx() loads the .gx file, decompresses it via GSZ3,
        # and execs the decompressed Python code
        execution_context = execute_gx(
            path=ctx.model_path,
            trace=ctx.params.get("trace", False),
            profile=ctx.params.get("profile", False)
        )

        print(f"[XIC] Execution complete")
        print(f"[XIC] Result: {execution_context.result if hasattr(execution_context, 'result') else 'OK'}")
        ctx._state["last_result"] = execution_context

    except ExecutionError as e:
        print(f"[XIC] Execution error: {e}")
        raise
    except Exception as e:
        print(f"[XIC] Unexpected error: {e}")
        raise


# Operation dispatch table
OP_TABLE = {
    "LOAD_MODEL": op_LOAD_MODEL,
    "SET_MODE": op_SET_MODE,
    "SET_PARAM": op_SET_PARAM,
    "RUN_PROMPT": op_RUN_PROMPT,
}