from __future__ import annotations

from langchain.agents.middleware.types import ModelRequest
from langchain_core.messages import HumanMessage

from EvoScientist.middleware.skill_context import BudgetedSkillsMiddleware


def _skill(name: str, description: str) -> dict[str, object]:
    return {
        "name": name,
        "description": description,
        "path": f"/skills/{name}/SKILL.md",
        "allowed_tools": [],
    }


def _request(query: str, skills: list[dict[str, object]]) -> ModelRequest:
    return ModelRequest(
        model=object(),
        messages=[HumanMessage(content=query)],
        system_prompt="Base system prompt.",
        state={"skills_metadata": skills},
    )


def test_skill_context_prefers_relevant_skill_and_omits_irrelevant_catalog():
    middleware = BudgetedSkillsMiddleware(
        backend=object(), sources=["/skills/"], max_skills=2, max_skills_bytes=1024
    )
    result = middleware.modify_request(
        _request(
            "请分析蛋白质结构预测结果",
            [
                _skill("protein-structure", "蛋白质结构预测和结果分析。"),
                _skill("frontend-design", "Build polished user interfaces."),
            ],
        )
    )

    prompt = str(result.system_message.content)
    assert "protein-structure" in prompt
    assert "frontend-design" not in prompt
    assert "query-relevant subset" in prompt


def test_skill_context_enforces_count_and_utf8_budget():
    middleware = BudgetedSkillsMiddleware(
        backend=object(),
        sources=["/skills/"],
        max_skills=16,
        max_skills_bytes=2048,
        max_description_bytes=128,
    )
    skills = [
        _skill(f"analysis-{index}", "analysis " + "x" * 1_000) for index in range(300)
    ]

    selected = middleware._select_skills(skills, "analysis")
    rendered = middleware._format_budgeted_skills(selected)

    assert len(selected) == 16
    assert len(rendered.encode("utf-8")) <= 2048
    assert "analysis-299" not in rendered


def test_skill_context_does_not_fall_back_to_all_skills_without_a_match():
    middleware = BudgetedSkillsMiddleware(backend=object(), sources=["/skills/"])
    result = middleware.modify_request(
        _request(
            "unrelated request",
            [_skill("protein-structure", "Protein folding workflow.")],
        )
    )

    prompt = str(result.system_message.content)
    assert "protein-structure" not in prompt
    assert "query-relevant subset" in prompt


def test_skill_context_accepts_a_single_skill_source():
    middleware = BudgetedSkillsMiddleware(backend=object(), sources="/skills/")

    assert middleware.sources == ["/skills/"]