"""Воркер одного LLM-тика (запускается сабпроцессом из isolated_runner).

Читает payload из stdin, гоняет create_harness() один раз и умирает —
каждый тик получает чистый asyncio-рантайм Strands.
Трасса пишется напрямую в event store (es доступен).
Результат — в JSON-файл (path в payload['result_path']); stdout — только лог.
"""
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def build_level_model_from_payload(payload: dict):
    """Модель для воркера — та же сборка, что и в runtime.build_level_model:
    9router понимает chat/completions (usage там же), Responses API — шим без usage."""
    from strands.models.openai import OpenAIModel
    import os
    params = {}
    effort = payload.get("effort") or "auto"
    if effort and effort != "auto":
        params["reasoning_effort"] = effort
    model_id = payload.get("model") or "auto"
    if model_id == "none":
        model_id = "auto"
    return OpenAIModel(
        client_args={"api_key": os.environ.get("OPENAI_API_KEY", ""),
                     "base_url": os.environ.get("OPENAI_BASE_URL", "http://127.0.0.1:20128/v1")},
        model_id=model_id,
        **({"params": params} if params else {}),
    )


def _emit(payload, obj):
    rp = payload.get("result_path")
    if rp:
        try:
            with open(rp, "w") as f:
                json.dump(obj, f, ensure_ascii=False)
        except OSError:
            pass


def _extract_usage(agent) -> dict:
    """Usage из истории сообщений агента.

    Strands прикрепляет usage к каждой assistant-сообщении в
    message["metadata"]["usage"] (event_loop.py) — по одному словарю
    на модельный вызов (шаг тика). Накапливаем по шагам: тик может
    состоять из многих вызовов, бюджет интересует суммарный.

    Старый разбор m["message"]["usage"] не находил ничего (структура
    другая) — usage всегда возвращался пустым, spend не рос.
    """
    usage = {"prompt_tokens": 0, "completion_tokens": 0, "model_calls": 0}
    try:
        msgs = getattr(agent, "messages", None) or []
        for m in msgs:
            u = ((m or {}).get("metadata") or {}).get("usage") or {}
            p = u.get("inputTokens") or 0
            c = u.get("outputTokens") or 0
            if isinstance(p, (int, float)):
                usage["prompt_tokens"] += p
            if isinstance(c, (int, float)):
                usage["completion_tokens"] += c
            if p or c:
                usage["model_calls"] += 1
    except Exception:
        pass
    if not usage["model_calls"]:
        return {}
    return usage


def _emit_budget_incremental(level: str, trace_id: str, usage_delta: dict):
    """Записать BudgetSpent за один модельный вызов.

    Воркер — единственный автор spend-событий, но scheduler может убить
    его по таймауту в любой момент (SIGKILL). Поэтому usage пишем
    ИНКРЕМЕНТАЛЬНО: каждый завершённый модельный вызов — своё событие.
    Хвост недописанного вызова теряется, но весь тик — никогда.
    """
    if not usage_delta:
        return
    try:
        from harness import es
        es.append_event("harness", "budget", "BudgetSpent", {
            "level": level, "trace_id": trace_id, "source": "worker",
            **usage_delta,
        })
    except Exception:
        pass  # трекинг не должен валить тик


def _wire_budget_hooks(agent, level: str, trace_id: str):
    """После каждого модельного вызова — usage-дельта в стор.

    Strands вешает usage на assistant-сообщение в
    message["metadata"]["usage"] (одна запись на модельный вызов).
    MessageAddedEvent стреляет на добавление КАЖДОЙ записи истории,
    нас интересуют только assistant с usage.
    """
    from strands.hooks import MessageAddedEvent

    seen = set()

    def on_message_added(event):
        msg = getattr(event, "message", None) or {}
        if msg.get("role") != "assistant":
            return
        usage = ((msg.get("metadata") or {}).get("usage")) or {}
        p = usage.get("inputTokens") or 0
        c = usage.get("outputTokens") or 0
        if not (p or c):
            return
        key = id(msg)
        if key in seen:  # MessageUpdatedEvent может переигрывать — дедуп по объекту
            return
        seen.add(key)
        _emit_budget_incremental(level, trace_id, {
            "prompt_tokens": p, "completion_tokens": c,
        })

    agent.add_hook(on_message_added, MessageAddedEvent)
    return agent


def main():
    payload = json.loads(sys.stdin.read())

    from harness import es
    from harness.runtime import setup_provider, _wire_trace_hooks

    setup_provider()
    from strands_harness import create_harness

    trace_id = payload.get("trace_prefix") or "tick"
    level = payload.get("level") or "?"
    try:
        agent = create_harness(
            model=build_level_model_from_payload(payload),
            effort="auto",  # effort уже в параметрах модели
            instructions=payload.get("instructions", ""),
            session=None,
            builtin_tools=("read", "write", "edit", "shell") if payload.get("can_shell") else [],
            memory=False,
            skills=False,
        )
        agent = _wire_trace_hooks(agent, es, trace_id)
        agent = _wire_budget_hooks(agent, level, trace_id)
        result = agent(payload.get("prompt", ""))
        usage = _extract_usage(agent)
        _emit(payload, {"ok": True, "text": str(result), "usage": usage, "trace_id": trace_id})
    except Exception as e:
        import traceback
        _emit(payload, {"ok": False, "error": f"{e}\n{traceback.format_exc()[-800:]}"})
        # орфан-воркер после убийства родителя может умереть на печати — не критично


if __name__ == "__main__":
    main()