"""Factory for creating RLM-enhanced agents.

This module provides ``create_rlm_agent()``, a complete standalone agent harness with:
- Planning (todo lists)
- Filesystem access (read, write, edit, search)
- Sub-agents for delegation
- RLM context isolation and evidence tracking (profile-driven tools)
- Memory and skills support
- Auto-summarization for long conversations

The agent is taught via system prompts to use RLM-style reasoning:
load context, search/peek, execute Python for analysis, track evidence,
and produce cited conclusions.
"""

from __future__ import annotations

import inspect
import re
import warnings
from collections.abc import Callable, Sequence
from functools import lru_cache
from importlib.metadata import PackageNotFoundError, version as _get_package_version
from pathlib import Path
from typing import TYPE_CHECKING, Any

from langchain.agents import create_agent
from langchain.agents.middleware import (
    HumanInTheLoopMiddleware,
    InterruptOnConfig,
    TodoListMiddleware,
)
from langchain.agents.middleware.types import AgentMiddleware
from langchain.agents.structured_output import ResponseFormat
from langchain.chat_models import init_chat_model
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import SystemMessage
from langchain_core.tools import BaseTool
from langgraph.cache.base import BaseCache
from langgraph.graph.state import CompiledStateGraph
from langgraph.store.base import BaseStore
from langgraph.types import Checkpointer

from rlmagents.middleware._tools import DEFAULT_RLM_TOOL_PROFILE, RLMToolProfile
from rlmagents.middleware.rlm import RLMMiddleware

if TYPE_CHECKING:
    from rlmagents._harness.backends.protocol import BackendFactory, BackendProtocol
    from rlmagents._harness.subagents import CompiledSubAgent, SubAgent

# Import incorporated harness
from langchain_anthropic import ChatAnthropic
from langchain_anthropic.middleware import AnthropicPromptCachingMiddleware

from rlmagents._harness.backends import StateBackend
from rlmagents._harness.filesystem import FilesystemMiddleware
from rlmagents._harness.memory import MemoryMiddleware
from rlmagents._harness.patch_tool_calls import PatchToolCallsMiddleware
from rlmagents._harness.skills import SkillsMiddleware
from rlmagents._harness.subagents import (
    GENERAL_PURPOSE_SUBAGENT,
    CompiledSubAgent,
    SubAgent,
    SubAgentMiddleware,
)
from rlmagents._harness.summarization import (
    SummarizationMiddleware,
    _compute_summarization_defaults,
)

# Load base prompt
BASE_AGENT_PROMPT = (Path(__file__).resolve().parent / "base_prompt.md").read_text()

_LANGCHAIN_MIN_VERSION = (1, 2, 10)
_LANGCHAIN_MAX_VERSION = (1, 3, 0)
_CODING_RESPONSE_FIELDS = ("plan", "edits", "verification", "risks", "confidence")
_CODING_RESPONSE_FIELD_DESCRIPTIONS = {
    "plan": "Execution plan for requested coding work.",
    "edits": "Concrete file- and line-level edits to apply.",
    "verification": "Commands or checks to verify correctness.",
    "risks": "Caveats and likely regressions.",
    "confidence": "Confidence level from 0 to 1.",
}


def _is_context_size(value: object) -> bool:
    return (
        isinstance(value, tuple)
        and len(value) == 2
        and isinstance(value[0], str)
        and isinstance(value[1], (int, float))
    )


def _normalize_summarization_defaults(defaults: dict[str, Any]) -> dict[str, Any]:
    """Normalize summarization defaults to valid langchain middleware settings."""
    trigger_value = defaults.get("trigger")
    keep_value = defaults.get("keep")

    trigger: tuple[str, int | float] = (
        trigger_value if _is_context_size(trigger_value) else ("messages", 20)
    )
    keep: tuple[str, int | float] = (
        keep_value if _is_context_size(keep_value) else ("messages", 6)
    )

    truncate_args = defaults.get("truncate_args_settings")
    if truncate_args is not None and not isinstance(truncate_args, dict):
        truncate_args = None

    return {
        "trigger": trigger,
        "keep": keep,
        "truncate_args_settings": truncate_args,
    }


def _get_langchain_version() -> str:
    """Return the installed langchain version."""
    return _get_package_version("langchain")


def _parse_version_tuple(version: str) -> tuple[int, int, int]:
    match = re.match(r"^(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)", version)
    if not match:
        raise ValueError(f"Unable to parse langchain version string: {version!r}")
    return (
        int(match.group("major")),
        int(match.group("minor")),
        int(match.group("patch")),
    )


def _is_langchain_version_supported(version: str) -> bool:
    parsed = _parse_version_tuple(version)
    return _LANGCHAIN_MIN_VERSION <= parsed < _LANGCHAIN_MAX_VERSION


def _extract_system_prompt_text(system_prompt: str | SystemMessage | None) -> str:
    if system_prompt is None:
        return ""
    if isinstance(system_prompt, SystemMessage):
        content = system_prompt.content
        if isinstance(content, str):
            return content
        return "".join(
            (
                block
                if isinstance(block, str)
                else str(block.get("text", ""))
                if isinstance(block, dict)
                else str(block)
            )
            for block in content
        )
    return str(system_prompt)


def _is_coding_task(system_prompt: str | SystemMessage | None) -> bool:
    prompt = _extract_system_prompt_text(system_prompt).lower()
    return bool(
        re.search(
            r"\b(code|coding|implement|refactor|patch|bug|function|class)\b",
            prompt,
        )
    )


def _coding_response_format() -> dict[str, Any]:
    return {
        "type": "json_schema",
        "schema": {
            "type": "object",
            "properties": {
                "plan": {
                    "type": "string",
                    "description": _CODING_RESPONSE_FIELD_DESCRIPTIONS["plan"],
                },
                "edits": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": _CODING_RESPONSE_FIELD_DESCRIPTIONS["edits"],
                },
                "verification": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": _CODING_RESPONSE_FIELD_DESCRIPTIONS["verification"],
                },
                "risks": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": _CODING_RESPONSE_FIELD_DESCRIPTIONS["risks"],
                },
                "confidence": {
                    "type": "number",
                    "minimum": 0,
                    "maximum": 1,
                    "description": _CODING_RESPONSE_FIELD_DESCRIPTIONS["confidence"],
                },
            },
            "required": list(_CODING_RESPONSE_FIELDS),
            "additionalProperties": False,
        },
    }


def _model_hint_text(
    model_input: str | BaseChatModel | None,
    resolved_model: BaseChatModel | None = None,
) -> str:
    """Build lowercase text hints for model/provider capability checks."""
    hints: list[str] = []

    for candidate in (model_input, resolved_model):
        if isinstance(candidate, str):
            hints.append(candidate.lower())
            continue
        if candidate is None:
            continue

        hints.append(type(candidate).__name__.lower())
        for attr_name in (
            "model_name",
            "model",
            "name",
            "provider",
            "openai_api_base",
            "base_url",
            "api_base",
        ):
            attr_value = getattr(candidate, attr_name, None)
            if isinstance(attr_value, str):
                hints.append(attr_value.lower())

    return " ".join(hints)


def _supports_default_coding_response_format(
    model_input: str | BaseChatModel | None,
    resolved_model: BaseChatModel | None = None,
) -> bool:
    """Return whether default coding json_schema output should be auto-enabled."""
    model_hints = _model_hint_text(model_input, resolved_model)

    # DeepSeek OpenAI-compatible endpoints currently reject this schema shape.
    if "deepseek" in model_hints:
        return False

    return True


@lru_cache(maxsize=1)
def _assert_langchain_compatibility() -> None:
    """Fail fast if the installed langchain API is unsupported."""
    try:
        version = _get_langchain_version()
    except PackageNotFoundError as exc:
        raise RuntimeError(
            "rlmagents requires langchain to build agents."
        ) from exc

    if not _is_langchain_version_supported(version):
        minimum_version = ".".join(map(str, _LANGCHAIN_MIN_VERSION))
        maximum_version = ".".join(map(str, _LANGCHAIN_MAX_VERSION))
        raise RuntimeError(
            "rlmagents supports langchain version >= "
            f"{minimum_version} and < {maximum_version}; "
            f"installed version is {version}."
        )

    create_agent_signature = inspect.signature(create_agent)
    parameters = create_agent_signature.parameters
    if "backend" not in parameters and not (
        "checkpointer" in parameters and "store" in parameters
    ):
        raise RuntimeError(
            "Installed langchain.agents.create_agent is missing required persistence args. "
            "Expected `backend` or both `checkpointer` and `store`."
        )


def _build_create_agent_kwargs(
    checkpointer: Checkpointer | None,
    store: BaseStore | None,
    backend: BackendProtocol | BackendFactory | None,
    debug: bool,
    name: str | None,
    cache: BaseCache | None,
) -> dict[str, Any]:
    """Build a kwargs mapping for langchain.agents.create_agent."""
    create_agent_parameters = inspect.signature(create_agent).parameters
    kwargs: dict[str, Any] = {}

    if "backend" in create_agent_parameters:
        kwargs["backend"] = backend
    else:
        kwargs["checkpointer"] = checkpointer
        kwargs["store"] = store

    if "debug" in create_agent_parameters:
        kwargs["debug"] = debug
    if "name" in create_agent_parameters:
        kwargs["name"] = name
    if "cache" in create_agent_parameters:
        kwargs["cache"] = cache

    return kwargs


def _build_rlm_middleware(
    *,
    sandbox_timeout: float,
    context_policy: str,
    auto_load_threshold: int,
    auto_load_preview_chars: int,
    rlm_tool_profile: RLMToolProfile,
    rlm_include_tools: Sequence[str],
    rlm_exclude_tools: Sequence[str],
    rlm_system_prompt: str | None,
    sub_query_model: BaseChatModel | None,
    sub_query_timeout: float,
) -> RLMMiddleware:
    """Build a consistently configured RLM middleware instance."""
    return RLMMiddleware(
        sandbox_timeout=sandbox_timeout,
        context_policy=context_policy,
        auto_load_threshold=auto_load_threshold,
        auto_load_preview_chars=auto_load_preview_chars,
        tool_profile=rlm_tool_profile,
        include_tools=rlm_include_tools,
        exclude_tools=rlm_exclude_tools,
        system_prompt=rlm_system_prompt,
        sub_query_model=sub_query_model,
        sub_query_timeout=sub_query_timeout,
    )


def _assert_has_rlm_middleware(
    middleware: Sequence[AgentMiddleware],
    *,
    stack_name: str,
) -> None:
    """Fail fast if a middleware stack is missing `RLMMiddleware`."""
    if any(isinstance(layer, RLMMiddleware) for layer in middleware):
        return
    raise RuntimeError(f"RLM invariant violation: `{stack_name}` middleware stack is missing RLMMiddleware.")


def get_default_model() -> ChatAnthropic:
    """Get the default model for rlmagents."""
    return ChatAnthropic(
        model_name="claude-sonnet-4-5-20250929",
        max_tokens=20000,
    )


def create_rlm_agent(
    model: str | BaseChatModel | None = None,
    tools: Sequence[BaseTool | Callable | dict[str, Any]] | None = None,
    *,
    system_prompt: str | SystemMessage | None = None,
    middleware: Sequence[AgentMiddleware] = (),
    subagents: list[SubAgent | CompiledSubAgent] | None = None,
    skills: list[str] | None = None,
    memory: list[str] | None = None,
    response_format: ResponseFormat | None = None,
    context_schema: type[Any] | None = None,
    checkpointer: Checkpointer | None = None,
    store: BaseStore | None = None,
    backend: BackendProtocol | BackendFactory | None = None,
    interrupt_on: dict[str, bool | InterruptOnConfig] | None = None,
    debug: bool = False,
    name: str | None = None,
    cache: BaseCache | None = None,
    # RLM-specific parameters
    sandbox_timeout: float = 180.0,
    context_policy: str = "trusted",
    auto_load_threshold: int = 10_000,
    auto_load_preview_chars: int = 600,
    rlm_tool_profile: RLMToolProfile = DEFAULT_RLM_TOOL_PROFILE,
    rlm_include_tools: Sequence[str] = (),
    rlm_exclude_tools: Sequence[str] = (),
    sub_query_model: BaseChatModel | None = None,
    sub_query_timeout: float = 120.0,
    rlm_system_prompt: str | None = None,
    enable_rlm_in_subagents: bool = True,
) -> CompiledStateGraph:
    """Create an RLM-enhanced agent.

    This is a complete standalone agent harness with planning, filesystem,
    sub-agents, plus profile-driven RLM tools for context isolation and evidence tracking.

    Args:
        model: The LLM to use. Defaults to claude-sonnet-4-5-20250929.
        tools: Additional tools beyond the built-in ones.
        system_prompt: Custom system instructions.
        middleware: Additional middleware to add.
        subagents: Sub-agent specifications.
        skills: Skill source paths.
        memory: Memory file paths (AGENTS.md).
        response_format: Structured output format.
        context_schema: Context schema.
        checkpointer: Checkpointer for persistence.
        store: Store for persistent storage.
        backend: Backend for file storage/execution.
        interrupt_on: Tool interrupt configurations.
        debug: Enable debug mode.
        name: Agent name.
        cache: Cache for the agent.
        sandbox_timeout: RLM REPL timeout.
        context_policy: RLM context policy.
        auto_load_threshold: Auto-load threshold for large results.
        auto_load_preview_chars: Preview characters shown after auto-load (0 disables previews).
        rlm_tool_profile: RLM tool profile (`full`, `reasoning`, or `core`).
        rlm_include_tools: Additional RLM tools to include on top of the profile.
        rlm_exclude_tools: RLM tools to remove from the selected profile.
        sub_query_model: Optional model used by RLM `sub_query()`/`llm_query()`.
            If omitted, the primary `model` is reused for sub-queries.
        sub_query_timeout: Timeout in seconds for recursive sub-query calls.
        rlm_system_prompt: Custom RLM workflow prompt.
        enable_rlm_in_subagents: Deprecated compatibility flag.

            RLM is always enabled for sub-agents; passing `False` has no effect.

    Returns:
        Compiled LangGraph agent.
    """
    _assert_langchain_compatibility()

    model_input = model

    # Resolve model
    if model is None:
        model = get_default_model()
    elif isinstance(model, str):
        if model.startswith("openai:"):
            model = init_chat_model(model, use_responses_api=True)
        else:
            model = init_chat_model(model)

    # Default recursive sub-query calls to the same model/provider unless
    # explicitly overridden.
    effective_sub_query_model = sub_query_model or model

    # Compute summarization defaults
    summarization_defaults = _normalize_summarization_defaults(
        _compute_summarization_defaults(model)
    )

    # Resolve backend
    resolved_backend = backend if backend is not None else StateBackend

    if not enable_rlm_in_subagents:
        warnings.warn(
            "`enable_rlm_in_subagents=False` is deprecated and ignored; "
            "RLM is always enabled in sub-agents.",
            DeprecationWarning,
            stacklevel=2,
        )

    # Create RLM middleware
    rlm_mw = _build_rlm_middleware(
        sandbox_timeout=sandbox_timeout,
        context_policy=context_policy,
        auto_load_threshold=auto_load_threshold,
        auto_load_preview_chars=auto_load_preview_chars,
        rlm_tool_profile=rlm_tool_profile,
        rlm_include_tools=rlm_include_tools,
        rlm_exclude_tools=rlm_exclude_tools,
        rlm_system_prompt=rlm_system_prompt,
        sub_query_model=effective_sub_query_model,
        sub_query_timeout=sub_query_timeout,
    )

    # Build main agent middleware stack
    agent_middleware: list[AgentMiddleware] = [
        rlm_mw,
        TodoListMiddleware(),
    ]

    # Add memory if provided
    if memory is not None:
        agent_middleware.append(
            MemoryMiddleware(backend=resolved_backend, sources=memory)
        )

    # Add skills if provided
    if skills is not None:
        agent_middleware.append(
            SkillsMiddleware(backend=resolved_backend, sources=skills)
        )

    # Add filesystem
    agent_middleware.append(
        FilesystemMiddleware(backend=resolved_backend)
    )

    # Process sub-agents with RLM integration
    processed_subagents: list[SubAgent | CompiledSubAgent] = []
    for spec in subagents or []:
        if "runnable" in spec:
            processed_subagents.append(spec)
        else:
            subagent_model = spec.get("model", model)
            if isinstance(subagent_model, str):
                subagent_model = init_chat_model(subagent_model)

            subagent_summarization_defaults = _normalize_summarization_defaults(
                _compute_summarization_defaults(subagent_model)
            )
            subagent_middleware: list[AgentMiddleware] = [
                TodoListMiddleware(),
            ]
            subagent_middleware.append(
                _build_rlm_middleware(
                    sandbox_timeout=sandbox_timeout,
                    context_policy=context_policy,
                    auto_load_threshold=auto_load_threshold,
                    auto_load_preview_chars=auto_load_preview_chars,
                    rlm_tool_profile=rlm_tool_profile,
                    rlm_include_tools=rlm_include_tools,
                    rlm_exclude_tools=rlm_exclude_tools,
                    rlm_system_prompt=rlm_system_prompt,
                    sub_query_model=effective_sub_query_model,
                    sub_query_timeout=sub_query_timeout,
                )
            )

            subagent_middleware.extend([
                FilesystemMiddleware(backend=resolved_backend),
                SummarizationMiddleware(
                    model=subagent_model,
                    backend=resolved_backend,
                    trigger=subagent_summarization_defaults["trigger"],
                    keep=subagent_summarization_defaults["keep"],
                    trim_tokens_to_summarize=None,
                    truncate_args_settings=subagent_summarization_defaults["truncate_args_settings"],
                ),
                AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore"),
                PatchToolCallsMiddleware(),
            ])

            if spec.get("skills"):
                subagent_middleware.append(
                    SkillsMiddleware(backend=resolved_backend, sources=spec["skills"])
                )
            subagent_middleware.extend(spec.get("middleware", []))

            processed_spec: SubAgent = {
                **spec,
                "model": subagent_model,
                "tools": spec.get("tools", tools or []),
                "middleware": subagent_middleware,
            }
            if "prompt" in spec:
                processed_spec["system_prompt"] = spec["prompt"]
            elif "system_prompt" not in processed_spec:
                processed_spec["system_prompt"] = spec.get("description", "You are a helpful assistant.")
            _assert_has_rlm_middleware(
                processed_spec["middleware"],
                stack_name=f"subagent:{processed_spec['name']}",
            )
            processed_subagents.append(processed_spec)

    # Build general-purpose sub-agent
    gp_middleware: list[AgentMiddleware] = [
        TodoListMiddleware(),
    ]
    gp_middleware.append(
        _build_rlm_middleware(
            sandbox_timeout=sandbox_timeout,
            context_policy=context_policy,
            auto_load_threshold=auto_load_threshold,
            auto_load_preview_chars=auto_load_preview_chars,
            rlm_tool_profile=rlm_tool_profile,
            rlm_include_tools=rlm_include_tools,
            rlm_exclude_tools=rlm_exclude_tools,
            rlm_system_prompt=rlm_system_prompt,
            sub_query_model=effective_sub_query_model,
            sub_query_timeout=sub_query_timeout,
        )
    )
    gp_middleware.extend([
        FilesystemMiddleware(backend=resolved_backend),
        SummarizationMiddleware(
            model=model,
            backend=resolved_backend,
            trigger=summarization_defaults["trigger"],
            keep=summarization_defaults["keep"],
            trim_tokens_to_summarize=None,
            truncate_args_settings=summarization_defaults["truncate_args_settings"],
        ),
        AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore"),
        PatchToolCallsMiddleware(),
    ])
    if skills:
        gp_middleware.append(SkillsMiddleware(backend=resolved_backend, sources=skills))
    if interrupt_on:
        gp_middleware.append(HumanInTheLoopMiddleware(interrupt_on=interrupt_on))
    _assert_has_rlm_middleware(gp_middleware, stack_name="subagent:general-purpose")

    general_purpose_spec: SubAgent = {
        **GENERAL_PURPOSE_SUBAGENT,
        "model": model,
        "tools": tools or [],
        "middleware": gp_middleware,
    }

    all_subagents: list[SubAgent | CompiledSubAgent] = [general_purpose_spec, *processed_subagents]

    # Add sub-agent middleware
    agent_middleware.append(
        SubAgentMiddleware(
            backend=resolved_backend,
            subagents=all_subagents,
        )
    )

    # Add summarization
    agent_middleware.append(
        SummarizationMiddleware(
            model=model,
            backend=resolved_backend,
            trigger=summarization_defaults["trigger"],
            keep=summarization_defaults["keep"],
            trim_tokens_to_summarize=None,
            truncate_args_settings=summarization_defaults["truncate_args_settings"],
        )
    )

    # Add Anthropic caching
    agent_middleware.append(
        AnthropicPromptCachingMiddleware(unsupported_model_behavior="ignore")
    )

    # Add patch tool calls
    agent_middleware.append(PatchToolCallsMiddleware())

    # Add user middleware
    if middleware:
        agent_middleware.extend(middleware)

    # Add human-in-the-loop
    if interrupt_on is not None:
        agent_middleware.append(HumanInTheLoopMiddleware(interrupt_on=interrupt_on))

    _assert_has_rlm_middleware(agent_middleware, stack_name="main")

    # Combine system prompt
    if system_prompt is None:
        final_prompt: str | SystemMessage = BASE_AGENT_PROMPT
    elif isinstance(system_prompt, SystemMessage):
        new_content = [
            *system_prompt.content_blocks,
            {"type": "text", "text": f"\n\n{BASE_AGENT_PROMPT}"},
        ]
        final_prompt = SystemMessage(content=new_content)
    else:
        final_prompt = system_prompt + "\n\n" + BASE_AGENT_PROMPT

    if (
        response_format is None
        and _is_coding_task(final_prompt)
        and _supports_default_coding_response_format(model_input, model)
    ):
        response_format = _coding_response_format()

    agent_kwargs = _build_create_agent_kwargs(
        checkpointer=checkpointer,
        store=store,
        backend=resolved_backend,
        debug=debug,
        name=name,
        cache=cache,
    )

    return create_agent(
        model,
        system_prompt=final_prompt,
        tools=tools,
        middleware=agent_middleware,
        response_format=response_format,
        context_schema=context_schema,
        **agent_kwargs,
    ).with_config({"recursion_limit": 1000})
