import importlib.util
import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import pytest

from mlflow.assistant.providers.claude_code import ClaudeCodeProvider
from mlflow.assistant.types import EventType

REPO_ROOT = Path(__file__).resolve().parents[2]
DEV_STUBS = REPO_ROOT / "dev" / "dev_stubs"
CLAUDE_CLI = DEV_STUBS / "claude_cli.py"


def _load(name: str, path: Path):
    # dev/ is not an installed package; load this module directly by path. It only
    # imports stdlib and references siblings as files, so standalone load works.
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    # Register before exec so @dataclass can resolve the module's globals by name.
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


dev_stubs = _load("mlflow_dev_stubs", DEV_STUBS / "__init__.py")


# --- claude CLI stub ---------------------------------------------------------


def run_claude(*args: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [sys.executable, CLAUDE_CLI, *args], capture_output=True, text=True, timeout=30
    )


def test_claude_auth_probe_exits_zero_with_valid_json():
    result = run_claude("-p", "hi", "--max-turns", "1", "--output-format", "json")
    assert result.returncode == 0
    payload = json.loads(result.stdout)
    assert payload["type"] == "result"
    assert payload["is_error"] is False


def test_claude_stream_json_parses_into_message_and_done_events():
    result = run_claude(
        "-p",
        "hello",
        "--output-format",
        "stream-json",
        "--verbose",
        "--append-system-prompt",
        "ignored prompt with --output-format text and --resume decoys",
    )
    assert result.returncode == 0

    provider = ClaudeCodeProvider()
    events = [
        event
        for line in result.stdout.splitlines()
        if line.strip()
        if (event := provider._parse_message_to_event(json.loads(line))) is not None
    ]
    assert [e.type for e in events] == [EventType.MESSAGE, EventType.DONE]
    assert events[0].data["message"]["content"][0]["text"]
    assert events[1].data["session_id"]


def test_claude_resume_reuses_session_id():
    result = run_claude("-p", "hi", "--output-format", "stream-json", "--resume", "sess-abc")
    assert result.returncode == 0
    events = [json.loads(line) for line in result.stdout.splitlines() if line.strip()]
    assert {e["session_id"] for e in events if "session_id" in e} == {"sess-abc"}


# --- stub registry -----------------------------------------------------------


def test_install_stubs_rejects_unknown_name():
    with pytest.raises(ValueError, match="Unknown stub"):
        dev_stubs.install_stubs(["not-a-stub"])


def test_install_stubs_cleans_up_temp_dirs_on_partial_failure(monkeypatch):
    tmp_root = Path(tempfile.gettempdir())
    before = set(tmp_root.glob("mlflow-dev-stub-bin-*"))

    def boom(_result):
        raise RuntimeError("boom")

    monkeypatch.setitem(dev_stubs._INSTALLERS, "boom", boom)
    with pytest.raises(RuntimeError, match="boom"):
        # claude stages a temp dir, then `boom` fails -- the dir must not leak.
        dev_stubs.install_stubs(["claude", "boom"])

    assert set(tmp_root.glob("mlflow-dev-stub-bin-*")) == before


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="the claude shim is a POSIX shell script; run_dev_server (its consumer) is POSIX-only",
)
def test_install_claude_stages_working_shim():
    result = dev_stubs.install_stubs(["claude"])
    try:
        (shim_dir,) = result.path_prepend
        probe = subprocess.run(
            [shim_dir / "claude", "-p", "hi", "--output-format", "json"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        assert probe.returncode == 0
        assert json.loads(probe.stdout)["type"] == "result"
    finally:
        for path in result.cleanup_paths:
            shutil.rmtree(path, ignore_errors=True)