"""Agent runner: build a tool-augmented agent and run one eval case.

The agent gets read-only ``read_file`` / ``list_directory`` tools scoped to the
suite's ``file_root``. Depending on the artifact contract it may also get a single
``write_report`` tool (configurable name); in ``final_message`` mode there is no
report tool and the agent's last message is graded instead.
"""

import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from langchain.agents import create_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import BaseTool, StructuredTool, tool

from skill_eval_runner.config import ArtifactConfig, EvalSuiteConfig, PromptsConfig
from skill_eval_runner.files import expand_files
from skill_eval_runner.suite import EvalCase
from skill_eval_runner.usage import sum_usage

# Key under which the report tool stashes its content for the runner to read back.
_ARTIFACT_KEY = "artifact"

DEFAULT_WITH_SKILL_PROMPT = (
    "You are a software engineering assistant. "
    "Follow the skill instructions below exactly before doing anything else.\n\n"
    "<skill>\n{skill}\n</skill>\n\n"
    "The fixture files are available via the read_file and list_directory tools.\n"
    "{artifact}"
)

DEFAULT_WITHOUT_SKILL_PROMPT = (
    "You are a software engineering assistant. "
    "The fixture files are available via the read_file and list_directory tools.\n"
    "{artifact}"
)


@dataclass
class AgentResult:
    """The graded artifact plus token/latency accounting for one agent run."""

    response: str
    input_tokens: int
    output_tokens: int
    duration_ms: int


def _artifact_instructions(artifact: ArtifactConfig) -> str:
    """Return the system-prompt sentence describing how to emit the final artifact."""
    if artifact.mode == "tool":
        return (
            f"When you have finished, call {artifact.tool_name} once with the full "
            "content. Do not print it as your text response — write it via the tool."
        )
    return "When you have finished, give the full result as your final message."


def build_system_prompt(
    *,
    with_skill: bool,
    skill_md: str,
    artifact: ArtifactConfig,
    prompts: PromptsConfig,
) -> str:
    """Render the system prompt for a run, filling the ``{skill}``/``{artifact}`` slots."""
    if with_skill:
        template = prompts.with_skill or DEFAULT_WITH_SKILL_PROMPT
    else:
        template = prompts.without_skill or DEFAULT_WITHOUT_SKILL_PROMPT
    # Literal replacement (not str.format) so braces in skill text are safe.
    return template.replace("{skill}", skill_md).replace(
        "{artifact}", _artifact_instructions(artifact)
    )


def make_tools(
    file_root: Path, artifact: ArtifactConfig, written: dict[str, str]
) -> list[BaseTool]:
    """Build the read-only file tools plus, in ``tool`` mode, the report tool."""
    tools = _make_file_tools(file_root.resolve())

    if artifact.mode == "tool":
        tools.append(_make_report_tool(artifact, written))

    return tools


def _make_file_tools(base: Path) -> list[BaseTool]:
    """Build the read-only ``read_file`` and ``list_directory`` tools, sandboxed to ``base``."""

    @tool
    def read_file(path: str) -> str:
        """Read a text file. Path is relative to the fixture root directory."""
        target = (base / path).resolve()
        if not target.is_relative_to(base):
            return "Error: path is outside the allowed directory"
        if not target.exists():
            return f"Error: file not found: {path}"
        if target.is_dir():
            return f"Error: {path} is a directory — use list_directory instead"
        return target.read_text()

    @tool
    def list_directory(path: str) -> str:
        """List the contents of a directory. Path is relative to the fixture root."""
        target = (base / path).resolve()
        if not target.is_relative_to(base):
            return "Error: path is outside the allowed directory"
        if not target.exists():
            return f"Error: path not found: {path}"
        if not target.is_dir():
            return f"Error: {path} is not a directory — use read_file instead"
        lines = []
        for entry in sorted(target.iterdir()):
            kind = "dir " if entry.is_dir() else "file"
            lines.append(f"{kind}  {entry.relative_to(base)}")
        return "\n".join(lines)

    return [read_file, list_directory]


def _make_report_tool(artifact: ArtifactConfig, written: dict[str, str]) -> BaseTool:
    """Create the single configurable ``write_report`` tool for ``tool`` mode."""

    def _write(content: str) -> str:
        written[_ARTIFACT_KEY] = content
        return f"{artifact.filename} written ({len(content)} characters)."

    return StructuredTool.from_function(
        func=_write,
        name=artifact.tool_name,
        description=artifact.tool_description,
    )


def _extract_response(written: dict[str, str], messages: list[Any]) -> str:
    """Prefer the artifact written via the report tool; else the last AI message."""
    if _ARTIFACT_KEY in written:
        return written[_ARTIFACT_KEY]
    for msg in reversed(messages):
        if isinstance(msg, AIMessage) and not msg.tool_calls:
            content = msg.content
            return content if isinstance(content, str) else str(content)
    return ""


def run_agent(
    case: EvalCase,
    *,
    with_skill: bool,
    model: BaseChatModel,
    config: EvalSuiteConfig,
) -> AgentResult:
    """Run one eval case in one configuration and return its artifact + accounting."""
    skill_md = config.skill_file.read_text() if with_skill else ""
    system = build_system_prompt(
        with_skill=with_skill,
        skill_md=skill_md,
        artifact=config.artifact,
        prompts=config.prompts,
    )

    written: dict[str, str] = {}
    tools = make_tools(config.file_root, config.artifact, written)
    agent = create_agent(model, tools, system_prompt=system)

    user_content = case.prompt
    if case.files:
        injected = expand_files(case.files, config.file_root)
        if injected:
            sections = "\n\n".join(
                f"### {path}\n```\n{content}\n```" for path, content in injected
            )
            user_content = f"{user_content}\n\n<context>\n{sections}\n</context>"

    start = time.time()
    result = agent.invoke(
        {"messages": [HumanMessage(content=user_content)]},
        config={"recursion_limit": 50},
    )
    duration_ms = int((time.time() - start) * 1000)

    messages = list(result["messages"])
    response = _extract_response(written, messages)
    input_tokens, output_tokens = sum_usage(messages)

    return AgentResult(
        response=response,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        duration_ms=duration_ms,
    )