"""
Worker-side test runner for the KVStore PBT tutorial.

Hypothesis runs here, off the UI thread, so the page can keep painting
the progress bar while a test is in flight. A TESTCASE_CALLBACKS hook
pushes per-example progress to the main thread via the sync channel.
"""

import sys
import types
import importlib
import io
import time
import traceback
import dataclasses
import ast

# Bootstrap a fake kvstore package mirroring the cross-imports the test
# files use (e.g. `from kvstore.test_message import inserts`).
_pkg = types.ModuleType("kvstore")
sys.modules.setdefault("kvstore", _pkg)

import message as _message  # noqa: E402

sys.modules["kvstore.message"] = _message
import store as _store  # noqa: E402

sys.modules["kvstore.store"] = _store
import client as _client  # noqa: E402

sys.modules["kvstore.client"] = _client
import trace as _trace  # noqa: E402

sys.modules["kvstore.trace"] = _trace

for _tname in ("test_message", "test_trace"):
    try:
        _tmod = importlib.import_module(_tname)
        sys.modules[f"kvstore.{_tname}"] = _tmod
    except Exception as _e:
        print(f"Worker: could not pre-register kvstore.{_tname}: {_e}")


from pyscript import sync  # noqa: E402
import js  # noqa: E402
import json  # noqa: E402


_PUBLISH_INTERVAL_MS = 50.0
_LAST_PUBLISH_MS = [0.0]

try:
    _channel = js.BroadcastChannel.new("hyp-progress")
except Exception as _e:
    _channel = None
    print(f"Worker: BroadcastChannel unavailable ({_e})")


def _post(payload: dict) -> None:
    """Send a JSON-encoded message to the main thread via BroadcastChannel."""
    if _channel is None:
        return
    try:
        _channel.postMessage(json.dumps(payload))
    except Exception:
        pass


def _publish_state(force: bool = False) -> None:
    """Stream progress to main. Throttled to ~20 Hz."""
    now_ms = time.time() * 1000.0
    if not force and (now_ms - _LAST_PUBLISH_MS[0]) < _PUBLISH_INTERVAL_MS:
        return
    _LAST_PUBLISH_MS[0] = now_ms
    _post(
        {
            "type": "__hyp_progress__",
            "examples": _state["examples"],
            "shrinks": _state["shrinks"],
            "phase": _state["phase"],
            "repr": _state.get("last_repr", ""),
            "status": _state.get("last_status", ""),
        }
    )


def _publish_start(max_examples: int) -> None:
    _LAST_PUBLISH_MS[0] = 0.0
    _post(
        {
            "type": "__hyp_start__",
            "max_examples": int(max_examples),
            "observability": bool(_OBSERVABILITY_OK),
        }
    )


def _publish_end() -> None:
    _post(
        {
            "type": "__hyp_end__",
            "examples": _state["examples"],
            "shrinks": _state["shrinks"],
        }
    )


_MAX_LOG = 5000  # cap so a runaway 1M-example run doesn't blow memory


_state = {
    "examples": 0,
    "shrinks": 0,
    "phase": "generate",
    "last_repr": "",
    "last_status": "",
    "log": [],  # list of {"repr","status","phase"} for post-hoc browsing
}


def _ast_to_value(node, depth: int = 0):
    """Recursive ast → JSON-safe Python. Constructor calls become dicts."""
    if depth > 8:
        return None
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        v = _ast_to_value(node.operand, depth + 1)
        return -v if isinstance(v, (int, float)) else v
    if isinstance(node, (ast.List, ast.Tuple)):
        return [_ast_to_value(e, depth + 1) for e in node.elts]
    if isinstance(node, ast.Dict):
        out = {}
        for k_node, v_node in zip(node.keys, node.values):
            try:
                k = _ast_to_value(k_node, depth + 1)
            except Exception:
                continue
            out[str(k)] = _ast_to_value(v_node, depth + 1)
        return out
    if isinstance(node, ast.Call):
        name = node.func.id if isinstance(node.func, ast.Name) else "?"
        out = {"__type__": name}
        for kw in node.keywords:
            if kw.arg is None:
                continue
            out[kw.arg] = _ast_to_value(kw.value, depth + 1)
        return out
    if isinstance(node, ast.Name):
        return node.id
    try:
        return ast.literal_eval(node)
    except Exception:
        return None


def _parse_representation(repr_str: str):
    """Parse 'test_xxx(ball=Ball(x=..., ...), duration=0.5)' into a kwargs dict.
    Returns {} on failure (so callers can move on without a replay button)."""
    if not repr_str or "(" not in repr_str:
        return {}
    try:
        tree = ast.parse(repr_str.strip(), mode="eval")
    except Exception:
        return {}
    call = tree.body
    if not isinstance(call, ast.Call):
        return {}
    out = {}
    for kw in call.keywords:
        if kw.arg is None:
            continue
        try:
            out[kw.arg] = _ast_to_value(kw.value)
        except Exception:
            pass
    return out


def _serialize_arg(arg, depth: int = 0):
    """Reduce a test-case argument to a JSON-safe form for replay."""
    if depth > 6:
        return repr(arg)[:120]
    try:
        if arg is None or isinstance(arg, (bool, int, float, str)):
            return arg
        if dataclasses.is_dataclass(arg) and not isinstance(arg, type):
            return {
                "__dataclass__": type(arg).__name__,
                **{f.name: _serialize_arg(getattr(arg, f.name), depth + 1)
                   for f in dataclasses.fields(arg)},
            }
        if isinstance(arg, (list, tuple)):
            return [_serialize_arg(x, depth + 1) for x in arg]
        if isinstance(arg, dict):
            return {str(k): _serialize_arg(v, depth + 1) for k, v in arg.items()}
    except Exception:
        pass
    return repr(arg)[:120]


def _detect_phase(data: dict) -> str:
    """Hypothesis observability events label phase under different keys
    depending on the version — check all the places it might live."""
    for key in ("phase", "how_generated"):
        v = data.get(key)
        if isinstance(v, str) and v:
            return v
    metadata = data.get("metadata")
    if isinstance(metadata, dict):
        for key in ("phase", "how_generated"):
            v = metadata.get(key)
            if isinstance(v, str) and v:
                return v
    return "generate"


def _observability_callback(testcase) -> None:
    """Called by Hypothesis after each test case."""
    try:
        data = testcase if isinstance(testcase, dict) else {}
        if data.get("type") not in (None, "test_case"):
            return
        phase = _detect_phase(data)
        if phase == "shrink":
            _state["shrinks"] += 1
            _state["phase"] = "shrink"
        else:
            _state["examples"] += 1
            _state["phase"] = phase
        repr_str = data.get("representation") or ""
        if not repr_str:
            args = data.get("arguments") or {}
            if isinstance(args, dict) and args:
                repr_str = ", ".join(f"{k}={v!r}" for k, v in args.items())
        repr_str = str(repr_str)[:240]
        status = str(data.get("status") or "")
        if repr_str:
            _state["last_repr"] = repr_str
        _state["last_status"] = status
        if len(_state["log"]) < _MAX_LOG:
            args_raw = data.get("arguments")
            args_clean = {}
            if isinstance(args_raw, dict) and args_raw:
                for k, v in args_raw.items():
                    args_clean[str(k)] = _serialize_arg(v)
            else:
                # Hypothesis often leaves `arguments` empty for positional
                # @given signatures; fall back to parsing the call string.
                args_clean = _parse_representation(repr_str)
            _state["log"].append({
                "repr": repr_str,
                "status": status,
                "phase": phase,
                "args": args_clean,
            })
        _publish_state()
    except Exception:
        pass


_OBSERVABILITY_OK = False
try:
    from hypothesis.internal.observability import TESTCASE_CALLBACKS

    TESTCASE_CALLBACKS.append(_observability_callback)
    _OBSERVABILITY_OK = True
except Exception as _e:
    print(f"Worker: observability hook unavailable ({_e}). "
          f"Progress will fall back to time-based estimate.")


def update_source(name: str, code: str) -> str:
    """Write a source file to the worker's FS and invalidate the module cache."""
    try:
        with open(f"{name}.py", "w") as f:
            f.write(code)
        for mod_key in (name, f"kvstore.{name}"):
            if mod_key in sys.modules:
                del sys.modules[mod_key]
        return "OK"
    except Exception as e:
        return f"ERR {e}"


def _get_max_examples(fn) -> int:
    for attr in (
        "_hypothesis_internal_use_settings",
        "hypothesis_internal_use_settings",
    ):
        s = getattr(fn, attr, None)
        if s is not None:
            v = getattr(s, "max_examples", None)
            if isinstance(v, int):
                return v
    return 100


def run_test(module_name: str, func_name: str):
    _state["examples"] = 0
    _state["shrinks"] = 0
    _state["phase"] = "generate"
    _state["last_repr"] = ""
    _state["last_status"] = ""
    _state["log"] = []

    buf = io.StringIO()
    old_out, old_err = sys.stdout, sys.stderr
    sys.stdout = sys.stderr = buf

    status = "PASS"
    rendered_error = ""
    t0 = time.time()
    max_examples = 100
    try:
        if module_name in sys.modules:
            del sys.modules[module_name]
        mod = importlib.import_module(module_name)
        fn = getattr(mod, func_name)
        max_examples = _get_max_examples(fn)
        _publish_start(max_examples)
        fn()
    except Exception as e:
        rendered_error = "".join(
            traceback.format_exception_only(type(e), e)
        ).strip()
        status = "FAIL"
    finally:
        sys.stdout, sys.stderr = old_out, old_err
        _publish_end()

    elapsed = time.time() - t0
    bug_threshold = 0.0
    try:
        import physics as _physics_mod  # already imported above if test used it
        bug_threshold = float(getattr(_physics_mod, "BUG_THRESHOLD", 0.0) or 0.0)
    except Exception:
        pass
    return {
        "status": status,
        "error": rendered_error,
        "examples": _state["examples"],
        "shrinks": _state["shrinks"],
        "max_examples": max_examples,
        "elapsed": elapsed,
        "output": buf.getvalue(),
        "log": list(_state["log"]),
        "log_truncated": _state["examples"] + _state["shrinks"] > _MAX_LOG,
        "bug_threshold": bug_threshold,
    }


def ping() -> str:
    return "pong"


def has_observability() -> bool:
    return _OBSERVABILITY_OK


sync.run_test = run_test
sync.update_source = update_source
sync.ping = ping
sync.has_observability = has_observability

# Tell main about observability state at startup via a one-shot message.
_post({"type": "__hyp_ready__", "observability": bool(_OBSERVABILITY_OK)})