import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch

import pytest
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient

from mlflow.assistant.config import AssistantConfig, ProjectConfig
from mlflow.assistant.config import ProviderConfig as AssistantProviderConfig
from mlflow.assistant.providers import OllamaProvider
from mlflow.assistant.providers.base import (
    AssistantProvider,
    CLINotInstalledError,
    NotAuthenticatedError,
    ProviderConfig,
    ProviderNotConfiguredError,
)
from mlflow.assistant.types import Event, Message, ToolUseBlock
from mlflow.server.assistant.api import (
    PermissionDecision,
    _require_localhost,
    assistant_router,
)
from mlflow.server.assistant.session import SESSION_DIR, SessionManager, save_process_pid
from mlflow.utils.os import is_windows


class MockProvider(AssistantProvider):
    """Mock provider for testing."""

    @property
    def name(self) -> str:
        return "mock_provider"

    @property
    def display_name(self) -> str:
        return "Mock Provider"

    @property
    def description(self) -> str:
        return "Mock provider for testing"

    @property
    def config_path(self) -> Path:
        return Path.home() / ".mlflow" / "assistant" / "mock-config.json"

    def is_available(self) -> bool:
        return True

    def load_config(self) -> ProviderConfig:
        return ProviderConfig()

    def check_connection(self, echo=print) -> None:
        pass

    def resolve_skills_path(self, base_directory: Path) -> Path:
        return base_directory / ".mock" / "skills"

    def list_models(self, base_url: str | None = None, api_key: str | None = None) -> list[str]:
        raise NotImplementedError

    async def astream(
        self,
        prompt: str,
        tracking_uri: str,
        session_id: str | None = None,
        mlflow_session_id: str | None = None,
        cwd: Path | None = None,
        context: dict[str, Any] | None = None,
    ):
        yield Event.from_message(message=Message(role="user", content="Hello from mock"))
        yield Event.from_result(result="complete", session_id="mock-session-123")


@pytest.fixture(autouse=True)
def isolated_config(tmp_path, monkeypatch):
    """Redirect config to tmp_path to avoid modifying real user config."""
    import mlflow.assistant.config as config_module

    config_home = tmp_path / ".mlflow" / "assistant"
    config_path = config_home / "config.json"

    monkeypatch.setattr(config_module, "MLFLOW_ASSISTANT_HOME", config_home)
    monkeypatch.setattr(config_module, "CONFIG_PATH", config_path)

    return config_home


@pytest.fixture(autouse=True)
def clear_sessions():
    """Clear session storage before each test."""
    if SESSION_DIR.exists():
        shutil.rmtree(SESSION_DIR)
    yield
    if SESSION_DIR.exists():
        shutil.rmtree(SESSION_DIR)


@pytest.fixture
def client():
    """Create test client with mock provider and bypassed localhost check."""
    app = FastAPI()
    app.include_router(assistant_router)

    # Override localhost dependency to allow TestClient requests
    async def mock_require_localhost():
        pass

    app.dependency_overrides[_require_localhost] = mock_require_localhost

    mock_provider = MockProvider()
    with (
        patch("mlflow.server.assistant.api.list_providers", return_value=[mock_provider]),
        patch("mlflow.server.assistant.api._get_selected_provider", return_value=mock_provider),
    ):
        yield TestClient(app)


def test_message(client):
    response = client.post(
        "/ajax-api/3.0/mlflow/assistant/message",
        json={
            "message": "Hello",
            "context": {"trace_id": "tr-123", "experiment_id": "exp-456"},
        },
    )

    assert response.status_code == 200
    data = response.json()
    session_id = data["session_id"]
    assert session_id is not None
    assert data["stream_url"] == f"/ajax-api/3.0/mlflow/assistant/stream/{data['session_id']}"

    # continue the conversation
    response = client.post(
        "/ajax-api/3.0/mlflow/assistant/message",
        json={"message": "Second message", "session_id": session_id},
    )

    assert response.status_code == 200
    assert response.json()["session_id"] == session_id


def test_stream_not_found_for_invalid_session(client):
    response = client.get("/ajax-api/3.0/mlflow/assistant/sessions/invalid-session-id/stream")
    assert response.status_code == 404
    assert "Session not found" in response.json()["detail"]


def test_stream_bad_request_when_no_pending_message(client):
    # Create session and consume the pending message
    r = client.post("/ajax-api/3.0/mlflow/assistant/message", json={"message": "Hi"})
    session_id = r.json()["session_id"]
    client.get(f"/ajax-api/3.0/mlflow/assistant/sessions/{session_id}/stream")

    # Try to stream again without a new message
    response = client.get(f"/ajax-api/3.0/mlflow/assistant/sessions/{session_id}/stream")

    assert response.status_code == 400
    assert "No pending message" in response.json()["detail"]


def test_stream_returns_sse_events(client):
    r = client.post("/ajax-api/3.0/mlflow/assistant/message", json={"message": "Hi"})
    session_id = r.json()["session_id"]

    response = client.get(f"/ajax-api/3.0/mlflow/assistant/sessions/{session_id}/stream")

    assert response.status_code == 200
    assert "text/event-stream" in response.headers["content-type"]

    content = response.text
    assert "event: message" in content
    assert "event: done" in content
    assert "Hello from mock" in content


def test_health_check_returns_ok_when_healthy(client):
    response = client.get("/ajax-api/3.0/mlflow/assistant/providers/mock_provider/health")
    assert response.status_code == 200
    assert response.json() == {"status": "ok"}


def test_health_check_returns_404_for_unknown_provider(client):
    response = client.get("/ajax-api/3.0/mlflow/assistant/providers/unknown_provider/health")
    assert response.status_code == 404
    assert "not found" in response.json()["detail"]


def test_health_check_returns_412_when_cli_not_installed():
    app = FastAPI()
    app.include_router(assistant_router)

    async def mock_require_localhost():
        pass

    app.dependency_overrides[_require_localhost] = mock_require_localhost

    class CLINotInstalledProvider(MockProvider):
        def check_connection(self, echo=None):
            raise CLINotInstalledError("CLI not installed")

    provider = CLINotInstalledProvider()
    with patch("mlflow.server.assistant.api.list_providers", return_value=[provider]):
        client = TestClient(app)
        response = client.get("/ajax-api/3.0/mlflow/assistant/providers/mock_provider/health")
        assert response.status_code == 412
        assert "CLI not installed" in response.json()["detail"]


def test_health_check_returns_401_when_not_authenticated():
    app = FastAPI()
    app.include_router(assistant_router)

    async def mock_require_localhost():
        pass

    app.dependency_overrides[_require_localhost] = mock_require_localhost

    class NotAuthenticatedProvider(MockProvider):
        def check_connection(self, echo=None):
            raise NotAuthenticatedError("Not authenticated")

    provider = NotAuthenticatedProvider()
    with patch("mlflow.server.assistant.api.list_providers", return_value=[provider]):
        client = TestClient(app)
        response = client.get("/ajax-api/3.0/mlflow/assistant/providers/mock_provider/health")
        assert response.status_code == 401
        assert "Not authenticated" in response.json()["detail"]


def test_get_config_returns_empty_config(client):
    response = client.get("/ajax-api/3.0/mlflow/assistant/config")
    assert response.status_code == 200
    data = response.json()
    assert data["providers"] == {}
    assert data["projects"] == {}


def test_get_config_returns_existing_config(client, tmp_path):
    # Set up existing config by saving it first
    project_dir = tmp_path / "project"
    project_dir.mkdir()

    config = AssistantConfig(
        providers={"claude_code": AssistantProviderConfig(model="default", selected=True)},
        projects={"exp-123": ProjectConfig(type="local", location=str(project_dir))},
    )
    config.save()

    response = client.get("/ajax-api/3.0/mlflow/assistant/config")
    assert response.status_code == 200
    data = response.json()
    assert data["providers"]["claude_code"]["model"] == "default"
    assert data["providers"]["claude_code"]["selected"] is True
    assert data["projects"]["exp-123"]["location"] == str(project_dir)


def test_update_config_sets_provider(client):
    response = client.put(
        "/ajax-api/3.0/mlflow/assistant/config",
        json={"providers": {"claude_code": {"model": "opus", "selected": True}}},
    )
    assert response.status_code == 200
    data = response.json()
    assert data["providers"]["claude_code"]["selected"] is True


def test_update_config_sets_project(client, tmp_path):
    project_dir = tmp_path / "my_project"
    project_dir.mkdir()

    response = client.put(
        "/ajax-api/3.0/mlflow/assistant/config",
        json={"projects": {"exp-456": {"type": "local", "location": str(project_dir)}}},
    )
    assert response.status_code == 200
    data = response.json()
    assert data["projects"]["exp-456"]["location"] == str(project_dir)


def test_update_config_expand_user_home(client, tmp_path):
    # Create a directory under a "fake home" structure to test ~ expansion
    fake_home = tmp_path / "home" / "user"
    project_dir = fake_home / "my_project"
    project_dir.mkdir(parents=True)

    with patch("mlflow.server.assistant.api.Path.expanduser") as mock_expanduser:
        # Make expanduser return our tmp_path directory
        mock_expanduser.return_value = project_dir

        response = client.put(
            "/ajax-api/3.0/mlflow/assistant/config",
            json={"projects": {"exp-456": {"type": "local", "location": "~/my_project"}}},
        )
        assert response.status_code == 200
        data = response.json()
        assert data["projects"]["exp-456"]["location"] == str(project_dir)


@pytest.mark.asyncio
async def test_localhost_allows_ipv4():
    mock_request = MagicMock()
    mock_request.client.host = "127.0.0.1"
    await _require_localhost(mock_request)


@pytest.mark.asyncio
async def test_localhost_allows_ipv6():
    mock_request = MagicMock()
    mock_request.client.host = "::1"
    await _require_localhost(mock_request)


@pytest.mark.asyncio
async def test_localhost_blocks_external_ip():
    mock_request = MagicMock()
    mock_request.client.host = "192.168.1.100"

    with pytest.raises(HTTPException, match="same host"):
        await _require_localhost(mock_request)


@pytest.mark.asyncio
async def test_localhost_blocks_external_hostname():
    mock_request = MagicMock()
    mock_request.client.host = "external.example.com"

    with pytest.raises(HTTPException, match="same host"):
        await _require_localhost(mock_request)


@pytest.mark.asyncio
async def test_localhost_blocks_when_no_client():
    mock_request = MagicMock()
    mock_request.client = None

    with pytest.raises(HTTPException, match="same host"):
        await _require_localhost(mock_request)


def test_validate_session_id_accepts_valid_uuid():
    valid_uuid = "f5f28c66-5ec6-46a1-9a2e-ca55fb64bf47"
    SessionManager.validate_session_id(valid_uuid)  # Should not raise


def test_validate_session_id_rejects_invalid_format():
    with pytest.raises(ValueError, match="Invalid session ID format"):
        SessionManager.validate_session_id("invalid-session-id")


def test_validate_session_id_rejects_path_traversal():
    with pytest.raises(ValueError, match="Invalid session ID format"):
        SessionManager.validate_session_id("../../../etc/passwd")


def _is_process_running(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except (OSError, ValueError):  # ValueError is raised on Windows
        return False


def test_patch_session_cancel_with_process(client):
    r = client.post("/ajax-api/3.0/mlflow/assistant/message", json={"message": "Hi"})
    session_id = r.json()["session_id"]

    # Start a real subprocess and register it with the session
    with subprocess.Popen([sys.executable, "-c", "import time; time.sleep(10)"]) as proc:
        save_process_pid(session_id, proc.pid)

        assert _is_process_running(proc.pid)

        response = client.patch(
            f"/ajax-api/3.0/mlflow/assistant/sessions/{session_id}",
            json={"status": "cancelled"},
        )

        assert response.status_code == 200
        data = response.json()
        assert "terminated" in data["message"]

        # Wait for the process to actually terminate
        proc.wait(timeout=5)
        assert proc.returncode is not None
        # On non-Windows, verify the process is no longer running via PID check.
        # Skip on Windows because PIDs are reused more aggressively.
        if not is_windows():
            assert not _is_process_running(proc.pid)


class _DeferredProvider(MockProvider):
    """Pauses at the prompt on the first turn; resumes from the decision in context."""

    async def astream(
        self,
        prompt,
        tracking_uri,
        session_id=None,
        mlflow_session_id=None,
        cwd=None,
        context=None,
    ):
        decisions = (context or {}).get("tool_decisions") or {}
        if not decisions:
            # First turn: surface the tool, emit the prompt, and end the turn.
            yield Event.from_message(
                Message(
                    role="assistant",
                    content=[ToolUseBlock(id="t1", name="Bash", input={"command": "echo hi"})],
                )
            )
            yield Event.from_permission_request("t1", "Bash", {"command": "echo hi"})
            yield Event.from_result(result=None, session_id="prov-paused")
            return
        # Resume: apply the delivered decision.
        allowed = decisions.get("t1") == "allow"
        yield Event.from_message(Message(role="assistant", content="ran" if allowed else "denied"))
        yield Event.from_result(result=None, session_id="prov-done")


@pytest.mark.parametrize(("decision", "expected_text"), [("allow", "ran"), ("deny", "denied")])
@pytest.mark.asyncio
async def test_stream_pauses_then_resumes(decision, expected_text):
    """The turn ENDS at the permission prompt (no hang, no cross-request state);
    a resume request then drives a fresh stream to completion.
    """
    from mlflow.server.assistant.api import resolve_permission, stream_response

    session_id = "f5f28c66-5ec6-46a1-9a2e-ca55fb64bf47"
    session = SessionManager.create()
    session.set_pending_message(role="user", content="hi")
    SessionManager.save(session_id, session)

    mock_request = MagicMock()
    mock_request.base_url = "http://localhost:5000/"
    provider = _DeferredProvider()

    # First turn: the stream completes immediately at the prompt (no await).
    with patch("mlflow.server.assistant.api._get_selected_provider", return_value=provider):
        response = await stream_response(mock_request, session_id)
        first = "".join([c async for c in response.body_iterator])
    assert "permission_request" in first
    assert "event: done" in first

    # Deliver the decision, then a fresh stream resumes to completion.
    res = await resolve_permission(
        session_id, PermissionDecision(request_id="t1", decision=decision)
    )
    assert res.session_id == session_id

    with patch("mlflow.server.assistant.api._get_selected_provider", return_value=provider):
        response2 = await stream_response(mock_request, session_id)
        second = "".join([c async for c in response2.body_iterator])
    assert expected_text in second
    assert "event: done" in second


class _CaptureProvider(MockProvider):
    """Records the prompt and context astream is called with, then completes."""

    def __init__(self):
        self.captured: dict[str, Any] = {}

    async def astream(
        self,
        prompt,
        tracking_uri,
        session_id=None,
        mlflow_session_id=None,
        cwd=None,
        context=None,
    ):
        self.captured = {"prompt": prompt, "context": context or {}}
        yield Event.from_result(result=None, session_id="prov-done")


@pytest.mark.asyncio
async def test_stream_prefers_new_message_over_stale_tool_decision():
    """A pending message and a stale decision can coexist if a resume stream never
    consumed the decision and the user typed again. The new message must win: the
    provider sees the prompt and NOT the stale tool_decisions (which would otherwise
    resume the abandoned turn and silently drop the message).
    """
    from mlflow.server.assistant.api import stream_response

    session_id = "f5f28c66-5ec6-46a1-9a2e-ca55fb64bf47"
    session = SessionManager.create()
    session.set_pending_message(role="user", content="what is 2+2")
    session.pending_tool_decisions = {"t1": "allow"}
    SessionManager.save(session_id, session)

    mock_request = MagicMock()
    mock_request.base_url = "http://localhost:5000/"
    provider = _CaptureProvider()

    with patch("mlflow.server.assistant.api._get_selected_provider", return_value=provider):
        response = await stream_response(mock_request, session_id)
        _ = "".join([c async for c in response.body_iterator])

    assert provider.captured["prompt"] == "what is 2+2"
    assert "tool_decisions" not in provider.captured["context"]


@pytest.mark.asyncio
async def test_stream_forwards_tool_decision_when_no_pending_message():
    """A genuine resume (decision delivered, no new message) still forwards the
    tool_decisions so the provider can continue the paused turn.
    """
    from mlflow.server.assistant.api import stream_response

    session_id = "f5f28c66-5ec6-46a1-9a2e-ca55fb64bf47"
    session = SessionManager.create()
    session.pending_tool_decisions = {"t1": "allow"}
    SessionManager.save(session_id, session)

    mock_request = MagicMock()
    mock_request.base_url = "http://localhost:5000/"
    provider = _CaptureProvider()

    with patch("mlflow.server.assistant.api._get_selected_provider", return_value=provider):
        response = await stream_response(mock_request, session_id)
        _ = "".join([c async for c in response.body_iterator])

    assert provider.captured["prompt"] == ""
    assert provider.captured["context"]["tool_decisions"] == {"t1": "allow"}


def test_resolve_permission_rejects_invalid_session(client):
    response = client.post(
        "/ajax-api/3.0/mlflow/assistant/sessions/not-a-uuid/permission",
        json={"request_id": "t1", "decision": "allow"},
    )
    assert response.status_code == 400


def test_resolve_permission_returns_404_for_unknown_session(client):
    # Well-formed UUID, but no such session exists.
    response = client.post(
        "/ajax-api/3.0/mlflow/assistant/sessions/f5f28c66-5ec6-46a1-9a2e-ca55fb64bf47/permission",
        json={"request_id": "t1", "decision": "allow"},
    )
    assert response.status_code == 404


def test_install_skills_success(client):
    with patch(
        "mlflow.server.assistant.api.install_skills", return_value=["skill1", "skill2"]
    ) as mock_install:
        response = client.post(
            "/ajax-api/3.0/mlflow/assistant/skills/install",
            json={"type": "custom", "custom_path": "/tmp/test-skills"},
        )

        assert response.status_code == 200
        data = response.json()
        assert data["installed_skills"] == ["skill1", "skill2"]
        expected_path = os.path.join(os.sep, "tmp", "test-skills")
        assert data["skills_directory"] == expected_path
        mock_install.assert_called_once_with(Path(expected_path))


def test_install_skills_skips_when_already_installed(client):
    with (
        patch("mlflow.server.assistant.api.Path.exists", return_value=True),
        patch(
            "mlflow.server.assistant.api.list_installed_skills",
            return_value=["existing_skill"],
        ) as mock_list,
        patch("mlflow.server.assistant.api.install_skills") as mock_install,
    ):
        response = client.post(
            "/ajax-api/3.0/mlflow/assistant/skills/install",
            json={"type": "custom", "custom_path": "/tmp/test-skills"},
        )

        assert response.status_code == 200
        data = response.json()
        assert data["installed_skills"] == ["existing_skill"]
        mock_install.assert_not_called()
        mock_list.assert_called_once()


def test_update_config_partial_update_preserves_selected_provider(client):
    ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
    # Pre-populate config: claude_code selected, ollama exists but not selected
    config = AssistantConfig(
        providers={
            "claude_code": AssistantProviderConfig(model="opus", selected=True),
            ollama: AssistantProviderConfig(
                model="llama3", selected=False, base_url="http://localhost:11434"
            ),
        }
    )
    config.save()

    # Partially update ollama base_url without a selected flag
    response = client.put(
        "/ajax-api/3.0/mlflow/assistant/config",
        json={"providers": {ollama: {"base_url": "http://localhost:12345"}}},
    )

    assert response.status_code == 200
    data = response.json()
    assert data["providers"]["claude_code"]["selected"] is True
    assert data["providers"][ollama]["selected"] is False
    assert data["providers"][ollama]["base_url"] == "http://localhost:12345"


def test_list_ollama_models_returns_model_list(client):
    mock_provider = MockProvider()
    mock_provider.list_models = MagicMock(return_value=["llama3"])

    with patch("mlflow.server.assistant.api.list_providers", return_value=[mock_provider]):
        response = client.get(
            "/ajax-api/3.0/mlflow/assistant/providers/mock_provider/models",
            params={"base_url": "http://localhost:11434"},
        )

    assert response.status_code == 200
    assert response.json() == {"models": ["llama3"]}
    mock_provider.list_models.assert_called_once_with("http://localhost:11434", None)


def test_list_models_reads_api_key_from_header_not_query(client):
    """api_key must travel as the X-API-Key header so it stays out of access
    logs, browser history, and referer headers. This test pins that
    contract and verifies the value reaches the provider unchanged.
    """
    mock_provider = MockProvider()
    mock_provider.list_models = MagicMock(return_value=["llama3"])

    with patch("mlflow.server.assistant.api.list_providers", return_value=[mock_provider]):
        response = client.get(
            "/ajax-api/3.0/mlflow/assistant/providers/mock_provider/models",
            params={"base_url": "http://localhost:11434"},
            headers={"X-API-Key": "sk-test-secret"},
        )

    assert response.status_code == 200
    assert response.json() == {"models": ["llama3"]}
    mock_provider.list_models.assert_called_once_with("http://localhost:11434", "sk-test-secret")


def test_list_models_ignores_api_key_query_param(client):
    """Defense in depth: even if a caller passes api_key as a query param,
    the endpoint must not forward it to the provider — that would re-enable
    the access-log leak the header migration was meant to prevent.
    """
    mock_provider = MockProvider()
    mock_provider.list_models = MagicMock(return_value=["llama3"])

    with patch("mlflow.server.assistant.api.list_providers", return_value=[mock_provider]):
        response = client.get(
            "/ajax-api/3.0/mlflow/assistant/providers/mock_provider/models",
            params={"base_url": "http://localhost:11434", "api_key": "sk-leaked"},
        )

    assert response.status_code == 200
    mock_provider.list_models.assert_called_once_with("http://localhost:11434", None)


def test_list_ollama_models_returns_412_when_not_installed(client):
    class MissingDependencyProvider(MockProvider):
        def list_models(self, base_url: str | None = None, api_key: str | None = None) -> list[str]:
            raise CLINotInstalledError("ollama package missing")

    with patch(
        "mlflow.server.assistant.api.list_providers",
        return_value=[MissingDependencyProvider()],
    ):
        response = client.get("/ajax-api/3.0/mlflow/assistant/providers/mock_provider/models")

    assert response.status_code == 412
    assert "ollama" in response.json()["detail"].lower()


def test_list_ollama_models_returns_503_on_connection_failure(client):
    class UnreachableProvider(MockProvider):
        def list_models(self, base_url: str | None = None, api_key: str | None = None) -> list[str]:
            raise ProviderNotConfiguredError("Cannot connect to Ollama server")

    with patch("mlflow.server.assistant.api.list_providers", return_value=[UnreachableProvider()]):
        response = client.get(
            "/ajax-api/3.0/mlflow/assistant/providers/mock_provider/models",
            params={"base_url": "http://localhost:11434"},
        )

    assert response.status_code == 503
    assert "Cannot connect" in response.json()["detail"]


def test_list_provider_models_returns_404_for_unsupported_provider(client):
    class UnsupportedProvider(MockProvider):
        @property
        def name(self) -> str:
            return "unsupported_provider"

    with patch(
        "mlflow.server.assistant.api.list_providers",
        return_value=[UnsupportedProvider()],
    ):
        response = client.get(
            "/ajax-api/3.0/mlflow/assistant/providers/unsupported_provider/models"
        )

    assert response.status_code == 404
    assert "not supported" in response.json()["detail"]