from __future__ import annotations

import argparse
import json
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any

import pytest

import pyro_mcp.ollama_demo as ollama_demo
from pyro_mcp.api import Pyro as RealPyro
from pyro_mcp.vm_manager import VmManager as RealVmManager


@pytest.fixture(autouse=True)
def _mock_vm_manager_for_tests(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    class TestPyro(RealPyro):
        def __init__(self) -> None:
            super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))

    monkeypatch.setattr(ollama_demo, "Pyro", TestPyro)


def _stepwise_model_response(payload: dict[str, Any], step: int) -> dict[str, Any]:
    if step == 1:
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "",
                        "tool_calls": [{"id": "1", "function": {"name": "vm_list_profiles"}}],
                    }
                }
            ]
        }
    if step == 2:
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "",
                        "tool_calls": [
                            {
                                "id": "2",
                                "function": {
                                    "name": "vm_run",
                                    "arguments": json.dumps(
                                        {
                                            "profile": "debian-git",
                                            "command": "printf 'true\\n'",
                                            "vcpu_count": 1,
                                            "mem_mib": 512,
                                            "network": True,
                                        }
                                    ),
                                },
                            }
                        ],
                    }
                }
            ]
        }
    return {
        "choices": [
            {"message": {"role": "assistant", "content": "Executed git command in ephemeral VM."}}
        ]
    }


def test_run_ollama_tool_demo_happy_path(monkeypatch: pytest.MonkeyPatch) -> None:
    requests: list[dict[str, Any]] = []

    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        assert base_url == "http://localhost:11434/v1"
        requests.append(payload)
        return _stepwise_model_response(payload, len(requests))

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)

    logs: list[str] = []
    result = ollama_demo.run_ollama_tool_demo(log=logs.append)

    assert result["fallback_used"] is False
    assert str(result["exec_result"]["stdout"]).strip() == "true"
    assert result["final_response"] == "Executed git command in ephemeral VM."
    assert len(result["tool_events"]) == 2
    assert any(line == "[model] input user" for line in logs)
    assert any(line == "[model] output assistant" for line in logs)
    assert any("[model] tool_call vm_run" in line for line in logs)
    assert any(line == "[tool] calling vm_run" for line in logs)
    assert any(line == "[tool] result vm_run" for line in logs)


def test_run_ollama_tool_demo_recovers_from_bad_vm_id(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    requests: list[dict[str, Any]] = []

    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        assert base_url == "http://localhost:11434/v1"
        requests.append(payload)
        step = len(requests)
        if step == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "role": "assistant",
                            "tool_calls": [
                                {
                                    "id": "1",
                                    "function": {
                                        "name": "vm_exec",
                                        "arguments": json.dumps(
                                            {
                                                "vm_id": "vm_list_profiles",
                                                "command": ollama_demo.NETWORK_PROOF_COMMAND,
                                            }
                                        ),
                                    },
                                }
                            ],
                        }
                    }
                ]
            }
        return _stepwise_model_response(payload, step - 1)

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)

    result = ollama_demo.run_ollama_tool_demo()
    first_event = result["tool_events"][0]
    assert first_event["tool_name"] == "vm_exec"
    assert first_event["success"] is False
    assert "does not exist" in str(first_event["result"]["error"])
    assert int(result["exec_result"]["exit_code"]) == 0


def test_run_ollama_tool_demo_resolves_vm_id_placeholder(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    requests: list[dict[str, Any]] = []

    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        assert base_url == "http://localhost:11434/v1"
        requests.append(payload)
        step = len(requests)
        if step == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "role": "assistant",
                            "content": "",
                            "tool_calls": [
                                {"id": "1", "function": {"name": "vm_list_profiles"}},
                                {"id": "2", "function": {"name": "vm_list_profiles"}},
                                {
                                    "id": "3",
                                    "function": {
                                        "name": "vm_create",
                                        "arguments": json.dumps(
                                            {
                                                "profile": "debian-git",
                                                "vcpu_count": "2",
                                                "mem_mib": "2048",
                                            }
                                        ),
                                    },
                                },
                                {
                                    "id": "4",
                                    "function": {
                                        "name": "vm_exec",
                                        "arguments": json.dumps(
                                            {
                                                "vm_id": "<vm_id_returned_by_vm_create>",
                                                "command": "printf 'true\\n'",
                                                "timeout_seconds": "300",
                                            }
                                        ),
                                    },
                                },
                            ],
                        }
                    }
                ]
            }
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "Executed git command in ephemeral VM.",
                    }
                }
            ]
        }

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)

    logs: list[str] = []
    result = ollama_demo.run_ollama_tool_demo(log=logs.append)

    assert result["fallback_used"] is False
    assert int(result["exec_result"]["exit_code"]) == 0
    assert any("resolved vm_id placeholder" in line for line in logs)
    exec_event = next(
        event for event in result["tool_events"] if event["tool_name"] == "vm_exec"
    )
    assert exec_event["success"] is True


def test_dispatch_tool_call_vm_exec_autostarts_created_vm(tmp_path: Path) -> None:
    pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
    created = pyro.create_vm(profile="debian-base", vcpu_count=1, mem_mib=512, ttl_seconds=60)
    vm_id = str(created["vm_id"])

    executed = ollama_demo._dispatch_tool_call(
        pyro,
        "vm_exec",
        {"vm_id": vm_id, "command": "printf 'git version\\n'", "timeout_seconds": "30"},
    )

    assert int(executed["exit_code"]) == 0


def test_run_ollama_tool_demo_raises_without_vm_exec(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url, payload
        return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]}

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
    with pytest.raises(RuntimeError, match="did not execute a successful vm_run or vm_exec"):
        ollama_demo.run_ollama_tool_demo()


def test_run_ollama_tool_demo_uses_fallback_when_not_strict(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url, payload
        return {"choices": [{"message": {"role": "assistant", "content": "No tools"}}]}

    class TestPyro(RealPyro):
        def __init__(self) -> None:
            super().__init__(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
    monkeypatch.setattr(ollama_demo, "Pyro", TestPyro)
    monkeypatch.setattr(
        ollama_demo,
        "_run_direct_lifecycle_fallback",
        lambda pyro: {
            "vm_id": "vm-1",
            "command": ollama_demo.NETWORK_PROOF_COMMAND,
            "stdout": "true\n",
            "stderr": "",
            "exit_code": 0,
            "duration_ms": 5,
            "execution_mode": "host_compat",
            "cleanup": {"deleted": True, "reason": "post_exec_cleanup", "vm_id": "vm-1"},
        },
    )
    logs: list[str] = []
    result = ollama_demo.run_ollama_tool_demo(strict=False, log=logs.append)
    assert result["fallback_used"] is True
    assert int(result["exec_result"]["exit_code"]) == 0
    assert any(line == "[model] output assistant" for line in logs)
    assert any("[fallback]" in line for line in logs)


def test_run_ollama_tool_demo_verbose_logs_values(monkeypatch: pytest.MonkeyPatch) -> None:
    requests = 0

    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url
        nonlocal requests
        requests += 1
        return _stepwise_model_response(payload, requests)

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)

    logs: list[str] = []
    result = ollama_demo.run_ollama_tool_demo(verbose=True, log=logs.append)

    assert result["fallback_used"] is False
    assert str(result["exec_result"]["stdout"]).strip() == "true"
    assert any("[model] input user:" in line for line in logs)
    assert any("[model] tool_call vm_list_profiles args={}" in line for line in logs)
    assert any("[tool] result vm_run " in line for line in logs)


@pytest.mark.parametrize(
    ("tool_call", "error"),
    [
        (1, "invalid tool call entry"),
        ({"id": "", "function": {"name": "vm_list_profiles"}}, "valid call id"),
        ({"id": "1"}, "function metadata"),
        ({"id": "1", "function": {"name": 3}}, "name is invalid"),
    ],
)
def test_run_ollama_tool_demo_tool_call_validation(
    monkeypatch: pytest.MonkeyPatch,
    tool_call: Any,
    error: str,
) -> None:
    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url, payload
        return {"choices": [{"message": {"role": "assistant", "tool_calls": [tool_call]}}]}

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
    with pytest.raises(RuntimeError, match=error):
        ollama_demo.run_ollama_tool_demo()


def test_run_ollama_tool_demo_max_rounds(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url, payload
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "tool_calls": [{"id": "1", "function": {"name": "vm_list_profiles"}}],
                    }
                }
            ]
        }

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
    with pytest.raises(RuntimeError, match="exceeded maximum rounds"):
        ollama_demo.run_ollama_tool_demo()


@pytest.mark.parametrize(
    ("exec_result", "error"),
    [
        ("bad", "result shape is invalid"),
        ({"exit_code": 1, "stdout": "true"}, "expected exit_code=0"),
        ({"exit_code": 0, "stdout": "false"}, "did not confirm repository clone success"),
    ],
)
def test_run_ollama_tool_demo_exec_result_validation(
    monkeypatch: pytest.MonkeyPatch,
    exec_result: Any,
    error: str,
) -> None:
    responses: list[dict[str, Any]] = [
        {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "tool_calls": [
                            {"id": "1", "function": {"name": "vm_run", "arguments": "{}"}}
                        ],
                    }
                }
            ]
        },
        {"choices": [{"message": {"role": "assistant", "content": "done"}}]},
    ]

    def fake_post_chat_completion(base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
        del base_url, payload
        return responses.pop(0)

    def fake_dispatch(pyro: Any, tool_name: str, arguments: dict[str, Any]) -> Any:
        del pyro, arguments
        if tool_name == "vm_run":
            return exec_result
        return {"ok": True}

    monkeypatch.setattr(ollama_demo, "_post_chat_completion", fake_post_chat_completion)
    monkeypatch.setattr(ollama_demo, "_dispatch_tool_call", fake_dispatch)
    with pytest.raises(RuntimeError, match=error):
        ollama_demo.run_ollama_tool_demo()


def test_dispatch_tool_call_coverage(tmp_path: Path) -> None:
    pyro = RealPyro(manager=RealVmManager(backend_name="mock", base_dir=tmp_path / "vms"))
    profiles = ollama_demo._dispatch_tool_call(pyro, "vm_list_profiles", {})
    assert "profiles" in profiles
    created = ollama_demo._dispatch_tool_call(
        pyro,
        "vm_create",
        {
            "profile": "debian-base",
            "vcpu_count": "1",
            "mem_mib": "512",
            "ttl_seconds": "60",
            "network": False,
        },
    )
    vm_id = str(created["vm_id"])
    started = ollama_demo._dispatch_tool_call(pyro, "vm_start", {"vm_id": vm_id})
    assert started["state"] == "started"
    status = ollama_demo._dispatch_tool_call(pyro, "vm_status", {"vm_id": vm_id})
    assert status["vm_id"] == vm_id
    executed = ollama_demo._dispatch_tool_call(
        pyro,
        "vm_exec",
        {"vm_id": vm_id, "command": "printf 'true\\n'", "timeout_seconds": "30"},
    )
    assert int(executed["exit_code"]) == 0
    executed_run = ollama_demo._dispatch_tool_call(
        pyro,
        "vm_run",
        {
            "profile": "debian-base",
            "command": "printf 'true\\n'",
            "vcpu_count": "1",
            "mem_mib": "512",
            "timeout_seconds": "30",
            "network": False,
        },
    )
    assert int(executed_run["exit_code"]) == 0
    with pytest.raises(RuntimeError, match="unexpected tool requested by model"):
        ollama_demo._dispatch_tool_call(pyro, "nope", {})


def test_format_tool_error() -> None:
    error = ValueError("bad args")
    result = ollama_demo._format_tool_error("vm_exec", {"vm_id": "x"}, error)
    assert result["ok"] is False
    assert result["error_type"] == "ValueError"


@pytest.mark.parametrize(
    ("arguments", "error"),
    [
        ({}, "must be a non-empty string"),
        ({"k": 3}, "must be a non-empty string"),
    ],
)
def test_require_str(arguments: dict[str, Any], error: str) -> None:
    with pytest.raises(ValueError, match=error):
        ollama_demo._require_str(arguments, "k")


@pytest.mark.parametrize(
    ("arguments", "expected"),
    [
        ({"k": 1}, 1),
        ({"k": "1"}, 1),
        ({"k": " 42 "}, 42),
    ],
)
def test_require_int_accepts_numeric_strings(arguments: dict[str, Any], expected: int) -> None:
    assert ollama_demo._require_int(arguments, "k") == expected


@pytest.mark.parametrize("value", ["", "abc", "2.5", True])
def test_require_int_validation(value: Any) -> None:
    with pytest.raises(ValueError, match="must be an integer"):
        ollama_demo._require_int({"k": value}, "k")


@pytest.mark.parametrize(("arguments", "expected"), [({}, False), ({"k": True}, True)])
def test_require_bool(arguments: dict[str, Any], expected: bool) -> None:
    assert ollama_demo._require_bool(arguments, "k", default=False) is expected


def test_require_bool_validation() -> None:
    with pytest.raises(ValueError, match="must be a boolean"):
        ollama_demo._require_bool({"k": "true"}, "k")


def test_post_chat_completion_success(monkeypatch: pytest.MonkeyPatch) -> None:
    class StubResponse:
        def __enter__(self) -> StubResponse:
            return self

        def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
            del exc_type, exc, tb

        def read(self) -> bytes:
            return b'{"ok": true}'

    def fake_urlopen(request: Any, timeout: int) -> StubResponse:
        assert timeout == 90
        assert request.full_url == "http://localhost:11434/v1/chat/completions"
        return StubResponse()

    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
    assert ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1}) == {"ok": True}


def test_post_chat_completion_errors(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_urlopen_error(request: Any, timeout: int) -> Any:
        del request, timeout
        raise urllib.error.URLError("boom")

    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen_error)
    with pytest.raises(RuntimeError, match="failed to call Ollama"):
        ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1})

    class StubResponse:
        def __enter__(self) -> StubResponse:
            return self

        def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
            del exc_type, exc, tb

        def read(self) -> bytes:
            return b'["bad"]'

    def fake_urlopen_non_object(request: Any, timeout: int) -> StubResponse:
        del request, timeout
        return StubResponse()

    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen_non_object)
    with pytest.raises(TypeError, match="unexpected Ollama response shape"):
        ollama_demo._post_chat_completion("http://localhost:11434/v1", {"x": 1})


@pytest.mark.parametrize(
    ("raw", "expected"),
    [(None, {}), ({}, {}), ("", {}), ('{"a":1}', {"a": 1})],
)
def test_parse_tool_arguments(raw: Any, expected: dict[str, Any]) -> None:
    assert ollama_demo._parse_tool_arguments(raw) == expected


def test_parse_tool_arguments_invalid() -> None:
    with pytest.raises(TypeError, match="decode to an object"):
        ollama_demo._parse_tool_arguments("[]")
    with pytest.raises(TypeError, match="dictionary or JSON object string"):
        ollama_demo._parse_tool_arguments(3)


@pytest.mark.parametrize(
    ("response", "msg"),
    [
        ({}, "did not contain completion choices"),
        ({"choices": [1]}, "unexpected completion choice format"),
        ({"choices": [{"message": "bad"}]}, "did not contain a message"),
    ],
)
def test_extract_message_validation(response: dict[str, Any], msg: str) -> None:
    with pytest.raises(RuntimeError, match=msg):
        ollama_demo._extract_message(response)


def test_build_parser_defaults() -> None:
    parser = ollama_demo._build_parser()
    args = parser.parse_args([])
    assert args.model == ollama_demo.DEFAULT_OLLAMA_MODEL
    assert args.base_url == ollama_demo.DEFAULT_OLLAMA_BASE_URL
    assert args.verbose is False


def test_build_parser_verbose_flag() -> None:
    parser = ollama_demo._build_parser()
    args = parser.parse_args(["-v"])
    assert args.verbose is True


def test_main_uses_parser_and_prints_logs(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    class StubParser:
        def parse_args(self) -> argparse.Namespace:
            return argparse.Namespace(base_url="http://x", model="m", verbose=False)

    monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())
    monkeypatch.setattr(
        ollama_demo,
        "run_ollama_tool_demo",
        lambda base_url, model, strict=True, verbose=False, log=None: {
            "exec_result": {
                "exit_code": 0,
                "stdout": "true\n",
                "execution_mode": "host_compat",
            },
            "fallback_used": False,
        },
    )
    ollama_demo.main()
    output = capsys.readouterr().out
    assert "[summary] exit_code=0 fallback_used=False execution_mode=host_compat" in output
    assert "[summary] stdout=" not in output


def test_main_verbose_prints_stdout(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    class StubParser:
        def parse_args(self) -> argparse.Namespace:
            return argparse.Namespace(base_url="http://x", model="m", verbose=True)

    monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())
    monkeypatch.setattr(
        ollama_demo,
        "run_ollama_tool_demo",
        lambda base_url, model, strict=True, verbose=False, log=None: {
            "exec_result": {
                "exit_code": 0,
                "stdout": "true\n",
                "execution_mode": "host_compat",
            },
            "fallback_used": False,
        },
    )
    ollama_demo.main()
    output = capsys.readouterr().out
    assert "[summary] stdout=true" in output


def test_main_logs_error_and_exits_nonzero(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    class StubParser:
        def parse_args(self) -> argparse.Namespace:
            return argparse.Namespace(base_url="http://x", model="m", verbose=False)

    monkeypatch.setattr(ollama_demo, "_build_parser", lambda: StubParser())

    def fake_run(
        base_url: str,
        model: str,
        strict: bool = True,
        verbose: bool = False,
        log: Any = None,
    ) -> dict[str, Any]:
        del base_url, model, strict, verbose, log
        raise RuntimeError("demo did not execute a successful vm_exec")

    monkeypatch.setattr(ollama_demo, "run_ollama_tool_demo", fake_run)

    with pytest.raises(SystemExit, match="1"):
        ollama_demo.main()

    output = capsys.readouterr().out
    assert "[error] demo did not execute a successful vm_exec" in output