"""
Colony Sound Block GPU Bridge — stdin/stdout JSON pipe.

Reads JSON world models from stdin, runs through SBColonyBlockKernel (GPU),
writes JSON results to stdout. One line per request/response.

This bridge is spawned by the Node.js MCP server as a persistent subprocess.
The GPU kernel stays warm between requests — no re-initialization overhead.

Usage:
    # Single request (pipe in, get result):
    echo '{"factors":[...], "relationships":[...], "options":[...]}' | python3 -m afs.gpu.colony_bridge

    # Persistent mode (newline-delimited JSON):
    python3 -m afs.gpu.colony_bridge --persistent

Protocol:
    Request:  {"factors": [...], "relationships": [...], "options": [...], "config": {...}}
    Response: {"strategies": [...], "summary": {...}, "total_ms": float, ...}
    Error:    {"error": "message"}
"""

import sys
import json
import time
import traceback

# GPU kernel import
from .sb_colony_block_kernel import SBColonyBlockKernel, gpu_available, DEFAULT_CONFIG


def process_request(request: dict, kernel_cache: dict) -> dict:
    """Process a single colony request through the GPU kernel.

    Reuses kernel if dimensions match, otherwise creates new one.
    """
    factors = request.get("factors", [])
    relationships = request.get("relationships", [])
    options = request.get("options", [])
    config = request.get("config", {})

    n_dims = len(factors)
    if n_dims == 0:
        return {"error": "No factors provided"}

    # Normalize options
    normalized_options = []
    for opt in options:
        if isinstance(opt, str):
            normalized_options.append({"name": opt})
        else:
            normalized_options.append(opt)

    n_attractors = max(len(normalized_options), 1)

    # Merge config with defaults
    merged_config = {**DEFAULT_CONFIG, **config}
    probe_count = merged_config.get("probe_count", 32)

    # Scale probes for high-D
    if n_dims >= 8 and probe_count < n_dims * 4:
        probe_count = max(probe_count, n_dims * 4)
        merged_config["probe_count"] = probe_count

    # Get or create kernel (reuse if dims match)
    cache_key = (n_dims, n_attractors, probe_count)
    kernel = kernel_cache.get("kernel")
    if kernel is None or kernel_cache.get("key") != cache_key:
        kernel = SBColonyBlockKernel(
            num_dims=n_dims,
            max_probes=probe_count,
            max_attractors=max(n_attractors, 8),
        )
        kernel_cache["kernel"] = kernel
        kernel_cache["key"] = cache_key

    # Run
    input_dict = {
        "factors": factors,
        "relationships": relationships,
        "options": normalized_options,
    }

    result = kernel.run(input_dict, merged_config)

    # Ensure JSON-serializable (numpy → list)
    return _to_json_safe(result)


def _to_json_safe(obj):
    """Convert numpy arrays and types to JSON-safe Python types."""
    import numpy as np

    if isinstance(obj, dict):
        return {k: _to_json_safe(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_to_json_safe(v) for v in obj]
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, (np.integer,)):
        return int(obj)
    elif isinstance(obj, (np.floating,)):
        return float(obj)
    elif isinstance(obj, np.bool_):
        return bool(obj)
    return obj


def main():
    """Main entry point — persistent or single-shot mode."""
    persistent = "--persistent" in sys.argv

    has_gpu = gpu_available()
    sys.stderr.write(f"[colony_bridge] GPU={'yes' if has_gpu else 'no'}, persistent={persistent}\n")
    sys.stderr.flush()

    if not has_gpu:
        sys.stderr.write("[colony_bridge] WARNING: No GPU available. Colony will fail.\n")
        sys.stderr.flush()

    kernel_cache = {}

    if persistent:
        # Persistent mode: read newline-delimited JSON from stdin
        sys.stderr.write("[colony_bridge] Persistent mode — waiting for requests on stdin\n")
        sys.stderr.flush()

        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue

            try:
                request = json.loads(line)
                result = process_request(request, kernel_cache)
                sys.stdout.write(json.dumps(result) + "\n")
                sys.stdout.flush()
            except Exception as e:
                err = {"error": str(e), "traceback": traceback.format_exc()}
                sys.stdout.write(json.dumps(err) + "\n")
                sys.stdout.flush()
    else:
        # Single-shot: read entire stdin as one JSON request
        try:
            raw = sys.stdin.read()
            request = json.loads(raw)
            result = process_request(request, kernel_cache)
            sys.stdout.write(json.dumps(result))
            sys.stdout.flush()
        except Exception as e:
            err = {"error": str(e), "traceback": traceback.format_exc()}
            sys.stdout.write(json.dumps(err))
            sys.stdout.flush()
            sys.exit(1)


if __name__ == "__main__":
    main()
