from __future__ import annotations

import json

from click.testing import CliRunner

from logos_adk.agent import Agent
from logos_adk.cli import cli
from logos_adk.dynamic_orchestration import DepartmentSupervisor
from logos_adk.providers.mock import MockProvider
from logos_adk.team import Team
from logos_adk.types import Response, Usage


def _department_supervisor_yaml() -> str:
    return (
        "name: ceo\n"
        "type: department_supervisor\n"
        "departments:\n"
        "  - name: seo_team\n"
        "    kind: team\n"
        "    team:\n"
        "      name: seo_team_runtime\n"
        "      type: team\n"
        "      agents:\n"
        "        - name: seo\n"
        "          model: mock\n"
    )


def _supervisor() -> DepartmentSupervisor:
    team = Team(
        "seo_team_runtime", agents=[Agent("seo", model=MockProvider(responses=["mock response"]))]
    )
    return DepartmentSupervisor("ceo").department("seo_team", team)


def test_run_command_supports_department_supervisor_config(tmp_path, monkeypatch):
    config = tmp_path / "department_supervisor.yaml"
    config.write_text(_department_supervisor_yaml())
    monkeypatch.setattr("logos_adk.cli._run_with_animation", lambda coro, message="Thinking": coro)
    monkeypatch.setattr("logos_adk.cli.load_runnable_from_config", lambda _: _supervisor())

    result = CliRunner().invoke(cli, ["run", str(config), "Need SEO help"])

    assert result.exit_code == 0
    assert "mock response" in result.output


def test_test_command_supports_department_supervisor_config(tmp_path, monkeypatch):
    config = tmp_path / "department_supervisor.yaml"
    config.write_text(_department_supervisor_yaml())
    monkeypatch.setattr("logos_adk.cli._run_with_animation", lambda coro, message="Testing": coro)
    monkeypatch.setattr("logos_adk.cli.load_runnable_from_config", lambda _: _supervisor())
    monkeypatch.setattr(
        "logos_adk.cli.describe_runnable_from_config",
        lambda _: {"kind": "department_supervisor", "name": "ceo"},
    )

    result = CliRunner().invoke(cli, ["test", str(config), "--prompt", "ping"])

    assert result.exit_code == 0
    assert "Loaded department_supervisor from config" in result.output
    assert "Test successful" in result.output


def test_run_command_writes_output_and_forwards_session_id(tmp_path, monkeypatch):
    config = tmp_path / "agent.yaml"
    config.write_text("name: bot\nmodel: mock\n")
    output = tmp_path / "result.json"
    calls: list[tuple[str, str | None]] = []

    class FakeRunnable:
        async def run(self, prompt: str, session_id: str | None = None):
            calls.append((prompt, session_id))
            return Response(
                content="serialized response",
                tool_calls=[],
                usage=Usage(input_tokens=3, output_tokens=5),
                raw={"source": "test"},
                reasoning="reasoned",
            )

    monkeypatch.setattr("logos_adk.cli._run_with_animation", lambda coro, message="Thinking": coro)
    monkeypatch.setattr("logos_adk.cli.load_runnable_from_config", lambda _: FakeRunnable())

    result = CliRunner().invoke(
        cli,
        [
            "run",
            str(config),
            "Need SEO help",
            "--session-id",
            "sess-1",
            "--output",
            str(output),
        ],
    )

    assert result.exit_code == 0
    assert calls == [("Need SEO help", "sess-1")]
    payload = json.loads(output.read_text())
    assert payload["content"] == "serialized response"
    assert payload["reasoning"] == "reasoned"
    assert payload["usage"] == {"input_tokens": 3, "output_tokens": 5}


def test_describe_command_writes_output_file(tmp_path, monkeypatch):
    config = tmp_path / "agent.yaml"
    config.write_text("name: bot\nmodel: mock\n")
    output = tmp_path / "describe.yaml"

    monkeypatch.setattr(
        "logos_adk.cli.describe_runnable_from_config",
        lambda _: {
            "name": "bot",
            "type": "agent",
            "model": "mock",
            "tools": [{"name": "search"}],
            "system_prompt": "You are helpful",
        },
    )

    result = CliRunner().invoke(cli, ["describe", str(config), "--output", str(output)])

    assert result.exit_code == 0
    assert "Runnable: bot" in result.output
    written = output.read_text()
    assert "name: bot" in written
    assert "type: agent" in written
    assert "model: mock" in written
