"""Middleware to inject live flow context into subagent prompts."""

import json
from collections.abc import Callable
from typing import Any

from langchain.agents.middleware.types import AgentMiddleware, ModelRequest, ModelResponse

from agent_flows_builder.utils import load_flow_json


class ContextInjectionMiddleware(AgentMiddleware):
    """Injects current /flow.json context into the system prompt."""

    def __init__(self, *, formatter: Callable[[dict[str, Any]], str] | None = None) -> None:
        super().__init__()
        self.formatter = formatter or _default_formatter

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        context_block = self._build_context_block(request)
        if context_block:
            prompt = request.system_prompt or ""
            prompt = f"{prompt}\n\n{context_block}".strip()
            request = request.override(system_prompt=prompt)
        return handler(request)

    def _build_context_block(self, request: ModelRequest) -> str:
        runtime = getattr(request, "runtime", None)
        if runtime is None:
            return ""

        files = runtime.state.get("files", {})
        flow_raw = files.get("/flow.json")
        if not flow_raw:
            return ""

        try:
            flow = load_flow_json(flow_raw)
        except Exception:
            return ""

        available_vars = _extract_variables(flow)
        step_ids = _extract_step_ids(flow)
        payload = {
            "available_variables": available_vars,
            "existing_step_ids": step_ids,
        }
        return self.formatter(payload)


def _extract_variables(flow: dict[str, Any]) -> list[str]:
    """Extract variable names from the flow_variables step."""
    steps = flow.get("steps", [])
    flow_vars_step = next((s for s in steps if s.get("type") == "flow_variables"), None)
    if not flow_vars_step:
        return []
    return [
        v.get("name", "")
        for v in flow_vars_step.get("config", {}).get("variables", [])
        if v.get("name")
    ]


def _extract_step_ids(flow: dict[str, Any]) -> list[str]:
    """Extract all step IDs for collision avoidance."""
    steps = flow.get("steps", [])
    return [s.get("id", "") for s in steps if s.get("id")]


def _default_formatter(payload: dict[str, Any]) -> str:
    """Render a concise context block for injection."""
    return (
        "## Current Flow Context (injected)\n"
        f"- available_variables: {json.dumps(payload.get('available_variables', []), ensure_ascii=False)}\n"
        f"- existing_step_ids: {json.dumps(payload.get('existing_step_ids', []), ensure_ascii=False)}"
    )
