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)
    symbolic_mode: bool = False


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, symbolic)."""
    if not args:
        raise ValueError("SET_MODE requires a mode argument")
    ctx.mode = str(args[0])
    if ctx.mode == "symbolic":
        ctx.symbolic_mode = True
    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 or symbolic cognition.

    If ctx.symbolic_mode is True, routes through glyphos/cognitive_kernel.py.
    Otherwise, routes to execute_gx() with optional GPU acceleration.
    """
    if not args:
        raise ValueError("RUN_PROMPT requires a prompt argument")

    prompt = str(args[0])

    if ctx.symbolic_mode:
        from glyphos.cognitive_kernel import run_symbolic_prompt
        result = run_symbolic_prompt(prompt, context=ctx.params.get("context"))
        print(f"[XIC-SYMBOLIC] {result}")
        ctx._state["last_symbolic_result"] = result
        return

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

    try:
        if ctx.params.get("use_gpu"):
            from xic_extensions.gpu_runtime import has_gpu, run_on_gpu
            if has_gpu():
                print(f"[XIC-GPU] Running on GPU: {ctx.model_path}")
                execution_context = run_on_gpu(ctx.model_path, ctx.params)
            else:
                print(f"[XIC-GPU] No GPU detected, falling back to CPU")
                execution_context = execute_gx(
                    ctx.model_path,
                    trace=ctx.params.get("trace", False),
                    profile=ctx.params.get("profile", False)
                )
        else:
            execution_context = execute_gx(
                ctx.model_path,
                trace=ctx.params.get("trace", False),
                profile=ctx.params.get("profile", False)
            )

        print(f"[XIC] Execution complete")
        print(f"[XIC] Result: {getattr(execution_context, 'result', '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


def op_STREAM(ctx: XICContext, *args):
    """STREAM <prompt>: Execute and stream output line by line."""
    if not args:
        raise ValueError("STREAM requires a prompt argument")
    prompt = str(args[0])

    if ctx.symbolic_mode:
        from glyphos.cognitive_kernel import run_symbolic_prompt
        result = run_symbolic_prompt(prompt, context=ctx.params.get("context"))
        for chunk in result.split("\n"):
            if chunk.strip():
                print(f"[XIC-STREAM] {chunk}")
        ctx._state["last_symbolic_result"] = result
        return

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

    use_gpu = ctx.params.get("use_gpu")
    if use_gpu:
        from xic_extensions.gpu_runtime import has_gpu, run_on_gpu
        if has_gpu():
            print(f"[XIC-GPU] Streaming on GPU: {ctx.model_path}")
            exec_ctx = run_on_gpu(ctx.model_path, ctx.params)
        else:
            print(f"[XIC-GPU] No GPU detected, falling back to CPU")
            exec_ctx = execute_gx(ctx.model_path,
                                  trace=ctx.params.get("trace", False),
                                  profile=ctx.params.get("profile", False))
    else:
        exec_ctx = execute_gx(ctx.model_path,
                              trace=ctx.params.get("trace", False),
                              profile=ctx.params.get("profile", False))

    result_text = str(getattr(exec_ctx, "result", "OK"))
    for chunk in result_text.split("\n"):
        if chunk.strip():
            print(f"[XIC-STREAM] {chunk}")
    ctx._state["last_result"] = exec_ctx


def op_CHAIN(ctx: XICContext, *args):
    """CHAIN <label>: Mark start of a named chain; passes context forward."""
    if not args:
        raise ValueError("CHAIN requires a label argument")
    label = str(args[0])
    ctx.params["chain_label"] = label
    print(f"[XIC-CHAIN] Entering chain: {label}")


def op_CALL_GLYPH(ctx: XICContext, *args):
    """CALL_GLYPH <glyph_id> <payload>: Invoke cognition with a glyph context."""
    if not args:
        raise ValueError("CALL_GLYPH requires glyph_id argument")
    glyph_id = str(args[0])
    payload = str(args[1]) if len(args) > 1 else ""

    from glyphos.cognitive_kernel import run_symbolic_prompt
    glyph_context = dict(ctx.params.get("context", {}))
    glyph_context["glyph_id"] = glyph_id
    result = run_symbolic_prompt(payload, context=glyph_context)
    print(f"[XIC-GLYPH] {result}")
    ctx._state[f"glyph_{glyph_id}"] = result


def op_SET_CONTEXT(ctx: XICContext, *args):
    """SET_CONTEXT <key> <value>: Set symbolic/cognitive context key."""
    if len(args) < 2:
        raise ValueError("SET_CONTEXT requires key and value")
    if "context" not in ctx.params:
        ctx.params["context"] = {}
    ctx.params["context"][str(args[0])] = args[1]
    print(f"[XIC] Context {args[0]} = {args[1]}")


def op_LOG(ctx: XICContext, *args):
    """LOG <message>: Structured log from XIC program."""
    message = str(args[0]) if args else ""
    print(f"[XIC-LOG] {message}")


# Operation dispatch table
OP_TABLE = {
    "LOAD_MODEL": op_LOAD_MODEL,
    "SET_MODE": op_SET_MODE,
    "SET_PARAM": op_SET_PARAM,
    "SET_CONTEXT": op_SET_CONTEXT,
    "RUN_PROMPT": op_RUN_PROMPT,
    "STREAM": op_STREAM,
    "CHAIN": op_CHAIN,
    "CALL_GLYPH": op_CALL_GLYPH,
    "LOG": op_LOG,
}