"""Benchmark: a ``for_each`` fan-out over N files, re-run after small changes.

    uv run maturin develop --release                      # the numbers mean little in a debug build
    uv run python examples/cocoindex_agents/benchmarks/fan_out.py --files 10000

One orchestrator opens a localfs source and fans out one closed worker per
file; every worker reads its file and answers. A scripted model stands in for
inference (no API key), so what is measured is this package and the engine:
wall time, tool executions and model calls for a cold run and then five
incremental cases: an unchanged rerun, one file touched (mtime only), one
file edited, one file added, one file removed.

Two shapes (DESIGN.md §7, *Open vs closed*):

- ``S1``: the orchestrator has no open child. Its own records (the walk, the
  source card) cover everything beneath it, so an unchanged rerun skips it
  whole.
- ``S2``: the orchestrator first calls a planner child that samples the source
  (an open child, as in the ``table`` flow). The planner records a read, so
  the orchestrator's memo is flagged ``has_open_descendants``.

Not part of the test suite: it takes minutes at 10,000 files. The numbers on
record are in DESIGN.md §7.
"""

from __future__ import annotations

import argparse
import os
import pathlib
import shutil
import sys
import tempfile
import time
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass

import cocoindex as coco

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from cocoindex_agents import (  # noqa: E402
    AGENTS_CONFIG,
    MODEL_CLIENT,
    AgentsConfig,
    RoleSpec,
    SessionResult,
    run_open_session,
)
from cocoindex_agents.hints import RunHints  # noqa: E402
from cocoindex_agents.spec import SessionSpec  # noqa: E402
from cocoindex_agents.testing import CountingRegistry, ScriptedModel  # noqa: E402
from cocoindex_agents.tools import agents, sources  # noqa: E402
from cocoindex_agents.transcript import ModelOutput, ToolCall, Transcript  # noqa: E402

ORCHESTRATOR = RoleSpec(
    "orchestrator", "bench", ("open_source", "call", "for_each", "call_agent")
)
WORKER = RoleSpec("worker", "bench", ("read_input",))
PLANNER = RoleSpec("planner", "bench", ("call",))

SHAPES = ("S1", "S2")


def _call(tool: str, **args: object) -> ToolCall:
    return ToolCall(tool, dict(args))


def _script(docs: pathlib.Path, shape: str) -> ScriptedModel:
    """The orchestrator's plan, as the model would make it, with no inference."""

    def script(spec: SessionSpec, prefix: Transcript, hints: RunHints) -> ModelOutput:
        n = len(prefix.steps)
        if spec.role.name == "worker":
            if n == 0:
                return ModelOutput(calls=(_call("read_input"),))
            text = prefix.steps[0].results[0].value["text"]
            return ModelOutput(answer=f"{len(text)} chars")
        if spec.role.name == "planner":
            if n == 0:
                return ModelOutput(
                    calls=(
                        _call(
                            "call", handle="src:docs", method="sample", args={"k": 3}
                        ),
                    )
                )
            return ModelOutput(answer="one worker per file")
        plan: list[ToolCall] = [
            _call(
                "open_source",
                kind="localfs",
                name="docs",
                path=str(docs),
                patterns=["*.md"],
            )
        ]
        if shape == "S2":
            plan.append(
                _call(
                    "call_agent",
                    key="plan",
                    task="look at a sample and propose the task",
                    role="planner",
                    inputs=["src:docs"],
                )
            )
        plan.append(
            _call(
                "for_each",
                source="src:docs",
                key="items",
                task="count the characters",
            )
        )
        if n < len(plan):
            return ModelOutput(calls=(plan[n],))
        return ModelOutput(answer=str(prefix.steps[-1].results[0].value))

    return ScriptedModel(script)


@dataclass
class Measurement:
    shape: str
    case: str
    seconds: float
    model_calls: int
    tool_executions: dict[str, int]

    def row(self) -> str:
        tools = ", ".join(f"{k}={v}" for k, v in sorted(self.tool_executions.items()))
        return (
            f"| {self.shape} | {self.case} | {self.seconds:.2f} | {self.model_calls} "
            f"| {tools or '—'} |"
        )


def _bump_mtime(path: pathlib.Path) -> None:
    """Edits within one process can be faster than the filesystem's mtime tick."""
    st = os.stat(path)
    os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns + 1_000_000))


def _write(path: pathlib.Path, text: str) -> None:
    path.write_text(text)
    _bump_mtime(path)


@contextmanager
def _workspace(files: int) -> Iterator[tuple[pathlib.Path, pathlib.Path]]:
    root = pathlib.Path(tempfile.mkdtemp(prefix="cocoindex_agents_bench_"))
    try:
        docs = root / "docs"
        docs.mkdir()
        for i in range(files):
            (docs / f"{i:06d}.md").write_text(f"document {i}\n" * 4)
        yield docs, root / "coco_db"
    finally:
        shutil.rmtree(root, ignore_errors=True)


def run_shape(shape: str, files: int) -> None:
    with _workspace(files) as (docs, db_path):
        env = coco.Environment(coco.Settings.from_env(db_path=db_path))
        try:
            _run_cases(env, docs, shape, files)
        finally:
            env.close()


def _run_cases(
    env: coco.Environment, docs: pathlib.Path, shape: str, files: int
) -> None:
    model = _script(docs, shape)
    tools = CountingRegistry(
        [
            sources.open_source,
            sources.call_method,
            sources.read_input,
            agents.for_each,
            agents.call_agent,
        ]
    )
    config = AgentsConfig(
        roles={r.name: r for r in (ORCHESTRATOR, WORKER, PLANNER)}, tools=tools
    )
    env.context_provider.provide(MODEL_CLIENT, model)
    env.context_provider.provide(AGENTS_CONFIG, config)
    app = coco.App(
        coco.AppConfig(name=f"fan_out_{shape}", environment=env),
        run_open_session,
        config.spec("orchestrator", "count the characters of every doc"),
    )

    def measure(case: str, mutate: Callable[[], None]) -> None:
        mutate()
        model.reset()
        tools.reset()
        started = time.perf_counter()
        result = app.update_blocking()
        seconds = time.perf_counter() - started
        assert isinstance(result, SessionResult)
        m = Measurement(shape, case, seconds, model.calls, dict(tools.executed))
        print(m.row(), flush=True)

    def nothing() -> None:
        pass

    # In order: a cold run, then one mutation per incremental case.
    cases: dict[str, Callable[[], None]] = {
        "cold": nothing,
        "unchanged": nothing,
        "touch one": lambda: _bump_mtime(docs / "000001.md"),
        "edit one": lambda: _write(docs / "000002.md", "edited\n"),
        "add one": lambda: _write(docs / f"{files:06d}.md", "added\n"),
        "remove one": lambda: (docs / "000003.md").unlink(),
    }
    for case, mutate in cases.items():
        measure(case, mutate)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--files", type=int, default=10_000)
    parser.add_argument("--shapes", nargs="+", choices=SHAPES, default=list(SHAPES))
    args = parser.parse_args(argv)
    print(f"files: {args.files}; cocoindex {coco.__version__}", flush=True)
    print("| shape | case | seconds | model calls | tool executions |")
    print("|---|---|---|---|---|")
    for shape in args.shapes:
        run_shape(shape, args.files)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())