import argparse
import gc
import statistics
import time
from dataclasses import dataclass
from typing import Callable, Dict, List, Tuple

from langchain_core.runnables import RunnableLambda
import langchain_core.runnables.utils as utils


# ----------------------------
# Test case builders
# ----------------------------

def case_direct_invoke() -> RunnableLambda:
    """agent.invoke(x) where agent is a captured RunnableLambda."""
    agent = RunnableLambda(lambda x: x)

    def my_func(x: str) -> str:
        return agent.invoke(x)

    return RunnableLambda(my_func)


def case_nested_attr_invoke() -> RunnableLambda:
    """box.agent.invoke(x) where box is captured and agent is an attribute."""
    agent = RunnableLambda(lambda x: x)

    class Box:
        def __init__(self, a):
            self.agent = a

    box = Box(agent)

    def my_func(x: str) -> str:
        return box.agent.invoke(x)

    return RunnableLambda(my_func)


def case_global_invoke() -> RunnableLambda:
    """GLOBAL_AGENT.invoke(x) where GLOBAL_AGENT is a module global."""
    global GLOBAL_AGENT
    GLOBAL_AGENT = RunnableLambda(lambda x: x)  # noqa: PLW0603

    def my_func(x: str) -> str:
        return GLOBAL_AGENT.invoke(x)

    return RunnableLambda(my_func)


def case_nested_two_levels() -> RunnableLambda:
    """outer.inner.agent.invoke(x) (deeper attribute chain)."""
    agent = RunnableLambda(lambda x: x)

    class Inner:
        def __init__(self, a):
            self.agent = a

    class Outer:
        def __init__(self, inner):
            self.inner = inner

    outer = Outer(Inner(agent))

    def my_func(x: str) -> str:
        return outer.inner.agent.invoke(x)

    return RunnableLambda(my_func)


CASES: Dict[str, Callable[[], RunnableLambda]] = {
    "direct_invoke": case_direct_invoke,
    "nested_attr_invoke": case_nested_attr_invoke,
    "global_invoke": case_global_invoke,
    "nested_two_levels": case_nested_two_levels,
}


# ----------------------------
# Helpers
# ----------------------------

def ns_now() -> int:
    return time.perf_counter_ns()


def time_one_call_ns(fn: Callable[[], None]) -> int:
    t0 = ns_now()
    fn()
    t1 = ns_now()
    return t1 - t0


def find_deps_cache_keys(r: RunnableLambda) -> List[str]:
    """Find which __dict__ keys appear after first computing deps."""
    before = set(r.__dict__.keys())
    _ = r.deps
    after = set(r.__dict__.keys())
    return sorted(after - before)


def clear_cache_keys(r: RunnableLambda, keys: List[str]) -> None:
    for k in keys:
        r.__dict__.pop(k, None)


def detect_getsource_called_when_computing_deps(maker: Callable[[], RunnableLambda]) -> Tuple[bool, str]:
    """
    Returns (called, error_message).
    - called=True if utils.inspect.getsource was invoked during r.deps
    - error_message is non-empty if r.deps raised some other error
    """
    r = maker()
    original_getsource = utils.inspect.getsource
    called = False

    def wrapped_getsource(*args, **kwargs):
        nonlocal called
        called = True
        return original_getsource(*args, **kwargs)

    utils.inspect.getsource = wrapped_getsource
    try:
        _ = r.deps
        return called, ""
    except Exception as e:
        # still restore, then report error
        return called, f"{type(e).__name__}: {e}"
    finally:
        utils.inspect.getsource = original_getsource


@dataclass
class Result:
    case: str
    getsource_called: str  # "yes"/"no"/"error"
    cold_ms: float
    fresh_us: float
    recompute_us: float


def bench_case(
    name: str,
    maker: Callable[[], RunnableLambda],
    n_fresh: int,
    n_recompute: int,
    repeats: int,
    do_gc: bool,
) -> Result:
    """
    Benchmarks:
      - cold_ms: first deps on a fresh runnable
      - fresh_us: avg deps cost when building a fresh runnable each iteration
      - recompute_us: avg deps cost when forcing recompute on same runnable by clearing cache keys
    Reports medians over repeats.
    """
    called, err = detect_getsource_called_when_computing_deps(maker)
    if err:
        getsource_flag = "error"
    else:
        getsource_flag = "yes" if called else "no"

    # Determine cache keys used by deps on this version (using this case)
    probe = maker()
    cache_keys = find_deps_cache_keys(probe)

    cold_samples_ms: List[float] = []
    fresh_samples_us: List[float] = []
    recompute_samples_us: List[float] = []

    for _ in range(repeats):
        if do_gc:
            gc.collect()

        # Cold
        r = maker()
        cold_ns = time_one_call_ns(lambda: getattr(r, "deps"))
        cold_samples_ms.append(cold_ns / 1e6)

        # Fresh instances
        def fresh_loop():
            for __ in range(n_fresh):
                rr = maker()
                _ = rr.deps

        fresh_ns = time_one_call_ns(fresh_loop)
        fresh_samples_us.append((fresh_ns / n_fresh) / 1e3)

        # Forced recompute on same instance
        r = maker()
        _ = r.deps  # establish cache
        def recompute_loop():
            for __ in range(n_recompute):
                clear_cache_keys(r, cache_keys)
                _ = r.deps

        recompute_ns = time_one_call_ns(recompute_loop)
        recompute_samples_us.append((recompute_ns / n_recompute) / 1e3)

    return Result(
        case=name,
        getsource_called=getsource_flag,
        cold_ms=statistics.median(cold_samples_ms),
        fresh_us=statistics.median(fresh_samples_us),
        recompute_us=statistics.median(recompute_samples_us),
    )


# ----------------------------
# Main
# ----------------------------

def main() -> None:
    p = argparse.ArgumentParser(description="RunnableLambda.deps: detect inspect.getsource usage + benchmark")
    p.add_argument("--cases", default="all", help="Comma-separated case names or 'all'")
    p.add_argument("--n-fresh", type=int, default=5000, help="Iterations for fresh-instance benchmark")
    p.add_argument("--n-recompute", type=int, default=2000, help="Iterations for forced recompute benchmark")
    p.add_argument("--repeats", type=int, default=7, help="Repeat count; median reported")
    p.add_argument("--no-gc", action="store_true", help="Disable gc.collect() between repeats")
    args = p.parse_args()

    if args.cases.strip().lower() == "all":
        selected = list(CASES.keys())
    else:
        selected = [c.strip() for c in args.cases.split(",") if c.strip()]
        unknown = [c for c in selected if c not in CASES]
        if unknown:
            raise SystemExit(f"Unknown case(s): {unknown}. Options: {list(CASES.keys())}")

    # Show cache keys for a representative case
    probe = case_nested_attr_invoke()
    print("deps cache keys added after first compute:", find_deps_cache_keys(probe))
    print()

    results: List[Result] = []
    for name in selected:
        results.append(
            bench_case(
                name=name,
                maker=CASES[name],
                n_fresh=args.n_fresh,
                n_recompute=args.n_recompute,
                repeats=args.repeats,
                do_gc=not args.no_gc,
            )
        )

    header = f"{'case':20} {'getsource?':>10} {'cold(ms)':>10} {'fresh(us)':>12} {'recompute(us)':>14}"
    print(header)
    print("-" * len(header))
    for r in results:
        print(f"{r.case:20} {r.getsource_called:>10} {r.cold_ms:10.3f} {r.fresh_us:12.2f} {r.recompute_us:14.2f}")


if __name__ == "__main__":
    main()