core-python/tests/test_space_agent.py

"""Tests Maris AI projektu aģenta helperiem."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from maris_core.space_agent import (
    SPACE_AGENT_MODEL_DEFAULT,
    SPACE_AGENT_SPACE_REPO_DEFAULT,
    SpaceAgentChatRequest,
    SpaceAgentToolCall,
    build_space_agent_messages,
    execute_space_agent_tool,
    generate_space_agent_reply,
    get_space_agent_runtime_info,
    resolve_space_agent_models,
)


def test_space_agent_runtime_defaults_are_stable(monkeypatch) -> None:
    monkeypatch.delenv("HF_SPACE_ASSISTANT_MODEL", raising=False)
    monkeypatch.delenv("HF_SPACE_REPO", raising=False)
    monkeypatch.delenv("HF_SPACE_ASSISTANT_MODELS", raising=False)

    runtime = get_space_agent_runtime_info()

    assert runtime.model == SPACE_AGENT_MODEL_DEFAULT
    assert runtime.default_model == SPACE_AGENT_MODEL_DEFAULT
    assert runtime.space_repo == SPACE_AGENT_SPACE_REPO_DEFAULT
    assert SPACE_AGENT_MODEL_DEFAULT in runtime.available_models
    assert "model_dataset_playbook" in runtime.tool_names
    assert "browser_capabilities" in runtime.tool_names
    assert "persona_catalog" in runtime.tool_names
    assert "read_workspace_file" in runtime.tool_names
    assert "write_workspace_file" in runtime.tool_names
    assert "workspace_command_catalog" in runtime.tool_names
    assert any(item["title"] == "Model & dataset fixer" for item in runtime.capabilities)
    assert runtime.command_presets
    assert runtime.command_presets[0]["items"]


def test_build_space_agent_messages_includes_system_prompt_and_history() -> None:
    request = SpaceAgentChatRequest(
        message="Palīdzi ar manu Maris darba telpu",
        model="MarisUK/Codex",
        task_mode="code",
        history=[
            {"role": "user", "content": "Sveiks"},
            {"role": "assistant", "content": "Sveiks!"},
        ],
    )

    messages = build_space_agent_messages(request)

    assert messages[0]["role"] == "system"
    assert "MarisUK/Codex" in messages[0]["content"]
    assert "Maris AI Project Operator" in messages[0]["content"]
    assert "model_dataset_playbook" in messages[0]["content"]
    assert "audit → validate → evaluate → fix → train → sync" in messages[0]["content"]
    assert "`code`" in messages[0]["content"]
    assert messages[-1] == {"role": "user", "content": "Palīdzi ar manu Maris darba telpu"}
    assert messages[1]["content"] == "Sveiks"


class _DummyClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.models: list[str] = []

    def chat_completion(self, **kwargs: object) -> dict[str, object]:
        self.models.append(str(kwargs["model"]))
        return {
            "choices": [
                {
                    "message": {
                        "content": "Profesionāla atbilde par Maris darba telpas konfigurāciju.",
                    }
                }
            ]
        }


def test_generate_space_agent_reply_uses_hf_client_when_available() -> None:
    request = SpaceAgentChatRequest(
        message="Sakārto manu Space",
        model="MarisUK/Codex",
    )

    response = generate_space_agent_reply(request, client_factory=_DummyClient, token="hf_test")

    assert response.model == "MarisUK/Codex"
    assert response.used_fallback is False
    assert "Profesionāla atbilde" in response.response
    assert response.task_mode == "chat"


def test_generate_space_agent_reply_uses_hf_inference_space_runtime_config(
    monkeypatch,
) -> None:
    captured_kwargs: dict[str, object] = {}

    class _ConfiguredClient:
        def __init__(self, **kwargs: object) -> None:
            captured_kwargs.update(kwargs)

        def chat_completion(self, **kwargs: object) -> dict[str, object]:
            return {
                "choices": [
                    {
                        "message": {
                            "content": f"Atbilde no {kwargs['model']}",
                        }
                    }
                ]
            }

    monkeypatch.setenv("HF_INFERENCE_API_KEY", "hf_inference_secret")

    response = generate_space_agent_reply(
        SpaceAgentChatRequest(message="Sakārto manu Space"),
        client_factory=_ConfiguredClient,
    )

    assert response.model == SPACE_AGENT_MODEL_DEFAULT
    assert captured_kwargs == {
        "provider": "hf-inference",
        "base_url": "https://api-inference.huggingface.co",
        "token": "hf_inference_secret",
    }


class _ToolCallingClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.calls = 0

    def chat_completion(self, **_: object) -> dict[str, object]:
        self.calls += 1
        if self.calls == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "content": (
                                '{"mode":"tool","tool_calls":['
                                '{"name":"project_runtime","arguments":{}},'
                                '{"name":"sync_commands","arguments":{}}'
                                "]}"
                            )
                        }
                    }
                ]
            }
        return {
            "choices": [
                {
                    "message": {
                        "content": '{"mode":"final","response":"Izmanto sync komandu un publicē Space profesionāli."}'
                    }
                }
            ]
        }


class _MultiStepToolCallingClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.calls = 0

    def chat_completion(self, **_: object) -> dict[str, object]:
        self.calls += 1
        if self.calls == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "content": (
                                '{"mode":"tool","tool_calls":['
                                '{"name":"list_huggingface_repo_files","arguments":'
                                '{"repo_id":"MarisUK/maris-ai-master","repo_type":"model"}}'
                                "]}"
                            )
                        }
                    }
                ]
            }
        if self.calls == 2:
            return {
                "choices": [
                    {
                        "message": {
                            "content": (
                                '{"mode":"tool","tool_calls":['
                                '{"name":"read_huggingface_repo_file","arguments":'
                                '{"repo_id":"MarisUK/maris-ai-master","repo_type":"model","path":"README.md"}},'
                                '{"name":"write_huggingface_repo_file","arguments":'
                                '{"repo_id":"MarisUK/maris-ai-master","repo_type":"model","path":"README.md",'
                                '"content":"salabots modelis","commit_message":"Fix model README"}}'
                                "]}"
                            )
                        }
                    }
                ]
            }
        return {
            "choices": [
                {
                    "message": {
                        "content": (
                            '{"mode":"final","response":"Pārbaudīju modeli, salaboju README un saglabāju izmaiņas."}'
                        )
                    }
                }
            ]
        }


def test_execute_space_agent_tool_returns_runtime_payload() -> None:
    result = execute_space_agent_tool(SpaceAgentToolCall(name="project_runtime", arguments={}))

    assert result["model"] == SPACE_AGENT_MODEL_DEFAULT
    assert result["space_repo"] == SPACE_AGENT_SPACE_REPO_DEFAULT


def test_execute_space_agent_tool_returns_model_dataset_playbook() -> None:
    result = execute_space_agent_tool(
        SpaceAgentToolCall(name="model_dataset_playbook", arguments={})
    )

    assert result["dataset_repo"]
    assert result["model_repo"]
    assert result["space_repo"]
    assert "latest_agent_principles" in result
    assert "recommended_loop" in result
    assert "validate_dataset" in result["repo_commands"]
    assert any(
        "HF_TOKEN" in item or "MARIS_REPO_TOKEN" in item for item in result["required_setup"]
    )


def test_execute_space_agent_tool_returns_browser_capabilities() -> None:
    result = execute_space_agent_tool(SpaceAgentToolCall(name="browser_capabilities", arguments={}))

    assert result["provider"] == "playwright"
    assert "extract_text" in result["supported_actions"]


def test_execute_space_agent_tool_returns_workspace_command_catalog() -> None:
    result = execute_space_agent_tool(
        SpaceAgentToolCall(name="workspace_command_catalog", arguments={})
    )

    assert result["presets"]
    assert any(group["category"] == "python" for group in result["presets"])
    assert any(
        item["id"] == "frontend-build" for group in result["presets"] for item in group["items"]
    )


def test_execute_space_agent_tool_returns_persona_catalog() -> None:
    result = execute_space_agent_tool(SpaceAgentToolCall(name="persona_catalog", arguments={}))

    assert result["default_persona_id"] == "assistant"
    assert any(persona["id"] == "teacher" for persona in result["personas"])


def test_execute_space_agent_tool_can_list_huggingface_repos(monkeypatch) -> None:
    class FakeRepo:
        def __init__(self, repo_id: str) -> None:
            self.id = repo_id

    class FakeApi:
        def list_models(self, **_: object):
            return [FakeRepo("MarisUK/maris-ai-master")]

        def list_datasets(self, **_: object):
            return []

        def list_spaces(self, **_: object):
            return [FakeRepo("MarisUK/maris.ai.agent")]

    monkeypatch.setattr("maris_core.space_agent._get_hf_api_client", lambda: FakeApi())

    result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="list_huggingface_repos", arguments={"repo_type": "all", "limit": 5}
        )
    )

    assert result["owner"] == "MarisUK"
    assert any(entry["id"] == "MarisUK/maris-ai-master" for entry in result["entries"])
    assert any(entry["repo_type"] == "space" for entry in result["entries"])


def test_execute_space_agent_tool_can_read_and_write_huggingface_repo_files(
    monkeypatch, tmp_path: Path
) -> None:
    uploaded: list[dict[str, object]] = []
    downloaded = tmp_path / "README.md"
    downloaded.write_text("hf saturs", encoding="utf-8")

    class FakeApi:
        def upload_file(self, **kwargs: object) -> None:
            uploaded.append(kwargs)

        def list_repo_files(self, **_: object):
            return ["README.md", "app.py"]

    monkeypatch.setattr("maris_core.space_agent._get_hf_api_client", lambda: FakeApi())
    monkeypatch.setattr(
        "maris_core.space_agent._download_hf_repo_file",
        lambda **_: str(downloaded),
    )

    listing = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="list_huggingface_repo_files",
            arguments={"repo_id": "MarisUK/maris.ai.agent", "repo_type": "space"},
        )
    )
    assert "README.md" in listing["entries"]

    read_result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="read_huggingface_repo_file",
            arguments={
                "repo_id": "MarisUK/maris.ai.agent",
                "repo_type": "space",
                "path": "README.md",
            },
        )
    )
    assert read_result["content"] == "hf saturs"

    write_result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="write_huggingface_repo_file",
            arguments={
                "repo_id": "MarisUK/maris.ai.agent",
                "repo_type": "space",
                "path": "README.md",
                "content": "jauns hf saturs",
                "commit_message": "Update README",
            },
        )
    )
    assert write_result["saved"] is True
    assert uploaded[0]["repo_id"] == "MarisUK/maris.ai.agent"
    assert uploaded[0]["path_in_repo"] == "README.md"


def test_execute_space_agent_tool_can_list_read_and_write_workspace_files(tmp_path: Path) -> None:
    workspace = tmp_path / "workspace"
    workspace.mkdir()
    existing_file = workspace / "notes.txt"
    existing_file.write_text("sākotnējais saturs", encoding="utf-8")

    listing = execute_space_agent_tool(
        SpaceAgentToolCall(name="list_workspace", arguments={"path": "."}),
        context={"workspace_root": str(workspace)},
    )
    assert listing["path"] == "."
    assert any(entry["path"] == "notes.txt" for entry in listing["entries"])

    read_result = execute_space_agent_tool(
        SpaceAgentToolCall(name="read_workspace_file", arguments={"path": "notes.txt"}),
        context={"workspace_root": str(workspace)},
    )
    assert read_result["content"] == "sākotnējais saturs"

    write_result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="write_workspace_file",
            arguments={"path": "nested/todo.txt", "content": "jauns saturs"},
        ),
        context={"workspace_root": str(workspace)},
    )
    assert write_result["saved"] is True
    assert write_result["operation"] == "create"
    assert "+++ b/nested/todo.txt" in write_result["diff"]
    assert (workspace / "nested" / "todo.txt").read_text(encoding="utf-8") == "jauns saturs"


def test_execute_space_agent_tool_stages_workspace_write_when_approval_required(
    tmp_path: Path,
) -> None:
    workspace = tmp_path / "workspace"
    workspace.mkdir()
    staged_payloads: list[dict[str, object]] = []

    result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="write_workspace_file",
            arguments={"path": "notes.txt", "content": "jauns saturs"},
        ),
        context={
            "workspace_root": str(workspace),
            "require_workspace_approval": True,
            "task_mode": "code",
            "stage_workspace_write": lambda payload: (
                staged_payloads.append(payload)
                or {"proposal_id": "workspace-1", "status": "pending", "diff": "draft diff"}
            ),
        },
    )

    assert result["requires_approval"] is True
    assert result["saved"] is False
    assert result["saved_to_draft"] is True
    assert result["proposal_id"] == "workspace-1"
    assert result["diff"] == "draft diff"
    assert (workspace / "notes.txt").read_text(encoding="utf-8") == "jauns saturs"
    assert staged_payloads[0]["task_mode"] == "code"


def test_execute_space_agent_tool_stages_huggingface_write_when_approval_required() -> None:
    staged_payloads: list[dict[str, object]] = []

    result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="write_huggingface_repo_file",
            arguments={
                "repo_id": "MarisUK/maris.ai.agent",
                "repo_type": "space",
                "path": "README.md",
                "content": "jauns saturs",
                "commit_message": "Update README",
            },
        ),
        context={
            "require_publish_approval": True,
            "task_mode": "design",
            "stage_hf_write": lambda payload: (
                staged_payloads.append(payload)
                or {"proposal_id": "proposal-1", "status": "pending"}
            ),
        },
    )

    assert result["requires_approval"] is True
    assert result["saved"] is False
    assert result["proposal_id"] == "proposal-1"
    assert result["status"] == "pending"
    assert staged_payloads[0]["task_mode"] == "design"


def test_execute_space_agent_tool_runs_workspace_command_with_context_runner() -> None:
    result = execute_space_agent_tool(
        SpaceAgentToolCall(
            name="run_workspace_command",
            arguments={"command": "python -m pytest tests/test_space_agent.py"},
        ),
        context={
            "workspace_command_runner": lambda arguments: {
                "ok": True,
                "command_display": arguments["command"],
                "exit_code": 0,
            }
        },
    )

    assert result["ok"] is True
    assert result["command_display"] == "python -m pytest tests/test_space_agent.py"
    assert result["exit_code"] == 0


def test_execute_space_agent_tool_blocks_paths_outside_workspace(tmp_path: Path) -> None:
    workspace = tmp_path / "workspace"
    workspace.mkdir()

    with pytest.raises(ValueError):
        execute_space_agent_tool(
            SpaceAgentToolCall(name="read_workspace_file", arguments={"path": "../secret.txt"}),
            context={"workspace_root": str(workspace)},
        )


def test_generate_space_agent_reply_supports_tool_calling_roundtrip() -> None:
    request = SpaceAgentChatRequest(message="Kā man syncot Space?", tool_calling=True)

    response = generate_space_agent_reply(
        request,
        client_factory=_ToolCallingClient,
        tool_context={"training_status": {"running": False}},
    )

    assert response.used_fallback is False
    assert response.response == "Izmanto sync komandu un publicē Space profesionāli."
    assert response.task_mode == "chat"
    assert [tool_call.name for tool_call in response.tool_calls] == [
        "project_runtime",
        "sync_commands",
    ]
    assert any(event["type"] == "tool_call" for event in response.events)
    assert response.events[-1]["type"] == "final"


def test_generate_space_agent_reply_supports_multi_step_model_fix(
    monkeypatch, tmp_path: Path
) -> None:
    downloaded = tmp_path / "README.md"
    downloaded.write_text("vecs saturs", encoding="utf-8")
    staged: list[dict[str, object]] = []

    class FakeApi:
        def list_repo_files(self, **_: object):
            return ["README.md", "config.json"]

    monkeypatch.setattr("maris_core.space_agent._get_hf_api_client", lambda: FakeApi())
    monkeypatch.setattr(
        "maris_core.space_agent._download_hf_repo_file",
        lambda **_: str(downloaded),
    )

    response = generate_space_agent_reply(
        SpaceAgentChatRequest(message="Pārbaudi manu modeli un salabo to."),
        client_factory=_MultiStepToolCallingClient,
        tool_context={
            "require_publish_approval": True,
            "task_mode": "improve",
            "stage_hf_write": lambda payload: (
                staged.append(payload) or {"proposal_id": "approval-1", "status": "pending"}
            ),
        },
    )

    assert response.response == "Pārbaudīju modeli, salaboju README un saglabāju izmaiņas."
    assert [tool_call.name for tool_call in response.tool_calls] == [
        "list_huggingface_repo_files",
        "read_huggingface_repo_file",
        "write_huggingface_repo_file",
    ]
    assert staged[0]["repo_id"] == "MarisUK/maris-ai-master"
    assert staged[0]["path"] == "README.md"
    assert staged[0]["commit_message"] == "Fix model README"
    assert staged[0]["content"] == "salabots modelis"
    assert any(
        event["type"] == "tool_result" and event["tool_name"] == "write_huggingface_repo_file"
        for event in response.events
    )
    assert response.change_previews[0]["requires_approval"] is True


class _FailingClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token

    def chat_completion(self, **_: object) -> dict[str, object]:
        raise OSError("offline")


class _StopIterationChoices:
    def __bool__(self) -> bool:
        return True

    def __getitem__(self, index: int) -> object:
        raise StopIteration(index)

    def __iter__(self):
        return iter(())


class _MalformedChoicesClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token

    def chat_completion(self, **_: object) -> dict[str, object]:
        return {"choices": _StopIterationChoices()}


class _RetryAcrossModelsClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.chat_models: list[str] = []

    def chat_completion(self, **kwargs: object) -> dict[str, object]:
        model = str(kwargs["model"])
        self.chat_models.append(model)
        if model == "broken/model":
            raise OSError("broken")
        return {
            "choices": [
                {
                    "message": {
                        "content": "Profesionāli pārbaudīju MarisUK saturu un varu turpināt ar labojumiem."
                    }
                }
            ]
        }


class _MissingChatCompletionClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token


class _RuntimeErrorChatClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token

    def chat_completion(self, **_: object) -> dict[str, object]:
        raise RuntimeError("provider rejected chat completion")


class _TextModelCompatibilityClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.captured_messages: list[list[dict[str, str]]] = []

    def chat_completion(self, **kwargs: object) -> dict[str, object]:
        self.captured_messages.append(list(kwargs["messages"]))  # type: ignore[arg-type]
        return {
            "choices": [
                {
                    "message": {
                        "content": json.dumps(
                            {
                                "mode": "final",
                                "response": "Strādāju tiešā teksta režīmā ar Maris modeli.",
                            },
                            ensure_ascii=False,
                        ),
                    }
                }
            ]
        }


class _ToolFailureRecoveryClient:
    def __init__(self, token: str | None = None) -> None:
        self.token = token
        self.calls = 0

    def chat_completion(self, **_: object) -> dict[str, object]:
        self.calls += 1
        if self.calls == 1:
            return {
                "choices": [
                    {
                        "message": {
                            "content": (
                                '{"mode":"tool","tool_calls":['
                                '{"name":"read_workspace_file","arguments":{"path":"README.md"}}'
                                "]}"
                            )
                        }
                    }
                ]
            }
        return {
            "choices": [
                {
                    "message": {
                        "content": (
                            '{"mode":"final","response":"Pārbaudīju rīka kļūdu un turpinu bez iekšējas kļūdas."}'
                        )
                    }
                }
            ]
        }


def test_space_agent_accepts_external_hf_model_ids() -> None:
    request = SpaceAgentChatRequest(message="Sakārto manu Space", model="MarisUK/maris-ai-master")

    assert request.model == "MarisUK/maris-ai-master"


def test_space_agent_rejects_malformed_model_ids() -> None:
    with pytest.raises(ValueError):
        SpaceAgentChatRequest(message="Sakārto manu Space", model="bad model")


def test_resolve_space_agent_models_uses_only_explicit_request_model(monkeypatch) -> None:
    monkeypatch.setenv("MARIS_AGENT_MODEL", "MarisUK/maris-ai-master")
    monkeypatch.setenv("MARIS_AGENT_MODELS", "meta-llama/Llama-3.3-70B-Instruct")

    models = resolve_space_agent_models("deepseek-ai/DeepSeek-V3.2")

    assert models == ("deepseek-ai/DeepSeek-V3.2",)


def test_generate_space_agent_reply_raises_when_chat_completion_returns_malformed_choices() -> None:
    request = SpaceAgentChatRequest(message="Pārbaudi manu Space konfigurāciju")

    with pytest.raises(RuntimeError, match="nevarēja pieslēgties modelim"):
        generate_space_agent_reply(request, client_factory=_MalformedChoicesClient)


def test_generate_space_agent_reply_uses_requested_model_without_hidden_retry() -> None:
    response = generate_space_agent_reply(
        SpaceAgentChatRequest(
            message="Auditē manu MarisUK Space un sagatavo profesionālu atbildi.",
            model="Qwen/Qwen3-Coder-480B-A35B-Instruct",
        ),
        client_factory=_RetryAcrossModelsClient,
    )

    assert response.model == "Qwen/Qwen3-Coder-480B-A35B-Instruct"
    assert response.used_fallback is False
    assert "Profesionāli pārbaudīju" in response.response


def test_generate_space_agent_reply_raises_when_chat_completion_is_missing() -> None:
    with pytest.raises(RuntimeError, match="nevarēja pieslēgties modelim"):
        generate_space_agent_reply(
            SpaceAgentChatRequest(message="Pārbaudi manu Space konfigurāciju"),
            client_factory=_MissingChatCompletionClient,
        )


def test_generate_space_agent_reply_raises_after_runtime_error_chat_completion() -> None:
    with pytest.raises(RuntimeError, match="provider rejected chat completion"):
        generate_space_agent_reply(
            SpaceAgentChatRequest(message="Pārbaudi manu Space konfigurāciju"),
            client_factory=_RuntimeErrorChatClient,
        )


def test_generate_space_agent_reply_surfaces_tool_errors_without_internal_failure() -> None:
    response = generate_space_agent_reply(
        SpaceAgentChatRequest(message="Nolasi README un pasaki ko redzi."),
        client_factory=_ToolFailureRecoveryClient,
    )

    assert response.response == "Pārbaudīju rīka kļūdu un turpinu bez iekšējas kļūdas."
    assert any(event["type"] == "tool_error" for event in response.events)
    assert any(tool_call.name == "read_workspace_file" for tool_call in response.tool_calls)


def test_maris_text_model_uses_simple_mode() -> None:
    client = _TextModelCompatibilityClient()
    response = generate_space_agent_reply(
        SpaceAgentChatRequest(
            message="Palīdzi man saprast Space konfigurāciju.",
            model="MarisUK/maris-ai-text",
            tool_calling=True,
        ),
        client_factory=lambda token=None: client,
    )

    assert response.response == "Strādāju tiešā teksta režīmā ar Maris modeli."
    assert response.tool_calls == []
    assert response.used_fallback is False
    assert any("vienkāršotu tiešās atbildes ceļu" in event["message"] for event in response.events)
    assert client.captured_messages
    assert "Atbildi tikai ar JSON" not in client.captured_messages[0][0]["content"]
    assert "Domā kā senior programmētājs" not in client.captured_messages[0][0]["content"]


def test_generate_space_agent_reply_raises_direct_model_error_on_failure() -> None:
    request = SpaceAgentChatRequest(message="Deploy manu agent space")

    with pytest.raises(RuntimeError, match="offline"):
        generate_space_agent_reply(
            request,
            client_factory=_FailingClient,
            tool_context={"training_status": {"progress": {"label": "Gaida startu"}}},
        )