from agent.schemas import Answer
from eval.checks import run_checks

GOOD = Answer(
    answer_md=(
        "Use `torch.optim.SGD`:\n"
        "```python\n"
        "import torch\n"
        "opt = torch.optim.SGD(model.parameters(), lr=0.01)\n"
        "```\n"
    ),
    symbols_used=["torch.optim.SGD"],
    torch_version="2.12",
)


def test_good_answer_passes_all():
    assert all(v is None for v in run_checks(GOOD).values())


def test_syntax_error_caught():
    bad = GOOD.model_copy(update={"answer_md": "```python\ndef broken(:\n```"})
    assert "block 0" in run_checks(bad)["parses"]


def test_disallowed_import_caught():
    bad = GOOD.model_copy(update={"answer_md": "```python\nimport requests\n```"})
    assert "requests" in run_checks(bad)["imports"]


def test_from_import_of_torch_allowed():
    ok = GOOD.model_copy(update={"answer_md": "```python\nfrom torch import nn\n```"})
    assert run_checks(ok)["imports"] is None


def test_missing_symbol_caught():
    bad = GOOD.model_copy(update={"symbols_used": ["torch.nn.LSTM"]})
    assert "torch.nn.LSTM" in run_checks(bad)["symbols"]


def test_aliased_symbol_accepted():
    ok = GOOD.model_copy(
        update={
            "answer_md": "```python\nfrom torch import nn\nlayer = nn.Linear(3, 4)\n```",
            "symbols_used": ["torch.nn.Linear"],
        }
    )
    assert run_checks(ok)["symbols"] is None


def test_indented_code_block_parses():
    ok = GOOD.model_copy(
        update={"answer_md": "```python\n    import torch\n    x = torch.ones(2)\n```"}
    )
    assert run_checks(ok)["parses"] is None


def test_untagged_fence_not_parsed_as_python():
    ok = GOOD.model_copy(
        update={"answer_md": "```\ndata/train/\n    class1/\n        img1.jpg\n```"}
    )
    assert run_checks(ok)["parses"] is None


def test_symbol_not_matched_as_substring_of_prose():
    # regression: `nn.Linear` must not be "present" just because the English
    # word "linear" appears; word-boundary matching, not substring
    bad = GOOD.model_copy(
        update={
            "answer_md": "SGD applies a linear update to the parameters.",
            "symbols_used": ["torch.nn.Linear"],
        }
    )
    assert "torch.nn.Linear" in run_checks(bad)["symbols"]


def test_capitalised_python_fence_still_checked_for_imports():
    # a bad import must not hide behind a ```Python (capital P) tag
    bad = GOOD.model_copy(update={"answer_md": "```Python\nimport requests\n```"})
    assert "requests" in run_checks(bad)["imports"]


def test_relative_import_is_disallowed():
    bad = GOOD.model_copy(update={"answer_md": "```python\nfrom . import utils\n```"})
    assert "relative" in run_checks(bad)["imports"]