from __future__ import annotations

import base64
import time
from typing import ClassVar

import httpx
import pytest
from fastapi.testclient import TestClient

import dify_agent.server.app as app_module
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
from dify_agent.layers.knowledge.configs import DifyKnowledgeBaseLayerConfig
from dify_agent.layers.knowledge.layer import DifyKnowledgeBaseLayer
from dify_agent.layers.shell import DifyShellLayerConfig
from dify_agent.layers.shell.layer import DifyShellLayer
from dify_agent.layers.runtime import DifyRuntimeLayerConfig
from dify_agent.layers.runtime.layer import DifyRuntimeLayer
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend
from dify_agent.runtime.compositor_factory import DifyAgentLayerProvider
from dify_agent.server.app import create_app, create_dify_api_inner_http_client, create_plugin_daemon_http_client
from dify_agent.server.settings import ServerSettings
from dify_agent.storage.redis_run_store import RedisRunStore


def _base64url_secret(value: bytes) -> str:
    return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")


def _execution_context() -> DifyExecutionContextLayerConfig:
    return DifyExecutionContextLayerConfig(
        tenant_id="tenant-1",
        user_id="user-1",
        user_from="account",
        agent_mode="workflow_run",
        invoke_from="service-api",
    )


def _patch_app_lifecycle(monkeypatch: pytest.MonkeyPatch) -> tuple[FakeRedis, FakePluginDaemonHttpClient]:
    fake_redis = FakeRedis()
    fake_http_client = FakePluginDaemonHttpClient()
    FakeRunScheduler.created.clear()
    FakeRedisModule.fake_redis = fake_redis
    monkeypatch.setattr(app_module, "Redis", FakeRedisModule)
    monkeypatch.setattr(app_module, "RunScheduler", FakeRunScheduler)

    def fake_create_plugin_daemon_http_client(_settings: ServerSettings) -> FakePluginDaemonHttpClient:
        return fake_http_client

    monkeypatch.setattr(app_module, "create_plugin_daemon_http_client", fake_create_plugin_daemon_http_client)
    return fake_redis, fake_http_client


class FakeRedis:
    closed: bool

    def __init__(self) -> None:
        self.closed = False

    async def aclose(self) -> None:
        self.closed = True


class FakeRunScheduler:
    created: list["FakeRunScheduler"] = []

    store: object
    shutdown_grace_seconds: float
    run_timeout_seconds: float
    layer_providers: tuple[DifyAgentLayerProvider, ...]
    plugin_daemon_http_client: FakePluginDaemonHttpClient
    dify_api_http_client: FakePluginDaemonHttpClient
    shutdown_called: bool

    def __init__(
        self,
        *,
        store: object,
        plugin_daemon_http_client: FakePluginDaemonHttpClient,
        dify_api_http_client: FakePluginDaemonHttpClient,
        shutdown_grace_seconds: float,
        run_timeout_seconds: float,
        layer_providers: tuple[DifyAgentLayerProvider, ...],
    ) -> None:
        self.store = store
        self.shutdown_grace_seconds = shutdown_grace_seconds
        self.run_timeout_seconds = run_timeout_seconds
        self.layer_providers = layer_providers
        self.plugin_daemon_http_client = plugin_daemon_http_client
        self.dify_api_http_client = dify_api_http_client
        self.shutdown_called = False
        self.created.append(self)

    async def shutdown(self) -> None:
        self.shutdown_called = True


class FakePluginDaemonHttpClient:
    timeout: object | None
    limits: object | None
    trust_env: bool | None
    is_closed: bool

    def __init__(
        self,
        *,
        timeout: object | None = None,
        limits: object | None = None,
        trust_env: bool | None = None,
    ) -> None:
        self.timeout = timeout
        self.limits = limits
        self.trust_env = trust_env
        self.is_closed = False

    async def aclose(self) -> None:
        self.is_closed = True


class FakeTimeout:
    connect: float
    read: float
    write: float
    pool: float

    def __init__(self, *, connect: float, read: float, write: float, pool: float) -> None:
        self.connect = connect
        self.read = read
        self.write = write
        self.pool = pool


class FakeLimits:
    max_connections: int
    max_keepalive_connections: int
    keepalive_expiry: float

    def __init__(self, *, max_connections: int, max_keepalive_connections: int, keepalive_expiry: float) -> None:
        self.max_connections = max_connections
        self.max_keepalive_connections = max_keepalive_connections
        self.keepalive_expiry = keepalive_expiry


class FakeRedisModule:
    fake_redis: ClassVar[FakeRedis | None] = None

    @staticmethod
    def from_url(_url: str) -> FakeRedis:
        assert FakeRedisModule.fake_redis is not None
        return FakeRedisModule.fake_redis


class FakeHttpxModule:
    Timeout: ClassVar[type[FakeTimeout]] = FakeTimeout
    Limits: ClassVar[type[FakeLimits]] = FakeLimits
    AsyncClient: ClassVar[type[FakePluginDaemonHttpClient]] = FakePluginDaemonHttpClient


@pytest.mark.parametrize(
    "path",
    [
        "/runs",
        "/execution-bindings",
        "/home-snapshots/from-binding",
        "/execution-bindings/files/list",
    ],
)
def test_create_app_authenticates_control_plane_routes(
    monkeypatch: pytest.MonkeyPatch,
    path: str,
) -> None:
    _patch_app_lifecycle(monkeypatch)
    settings = ServerSettings(redis_url="redis://example.invalid/0", api_token="secret-token")

    with TestClient(create_app(settings)) as client:
        assert client.post(path, json={}).status_code == 401
        assert client.post(path, headers={"Authorization": "Bearer secret-token"}, json={}).status_code != 401


def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
    fake_redis = FakeRedis()
    fake_http_client = FakePluginDaemonHttpClient()
    fake_dify_api_http_client = FakePluginDaemonHttpClient()
    FakeRunScheduler.created.clear()
    FakeRedisModule.fake_redis = fake_redis
    monkeypatch.setattr(app_module, "Redis", FakeRedisModule)
    monkeypatch.setattr(app_module, "RunScheduler", FakeRunScheduler)

    def fake_create_plugin_daemon_http_client(_settings: ServerSettings) -> FakePluginDaemonHttpClient:
        return fake_http_client

    def fake_create_dify_api_inner_http_client(_settings: ServerSettings) -> FakePluginDaemonHttpClient:
        return fake_dify_api_http_client

    monkeypatch.setattr(app_module, "create_plugin_daemon_http_client", fake_create_plugin_daemon_http_client)
    monkeypatch.setattr(app_module, "create_dify_api_inner_http_client", fake_create_dify_api_inner_http_client)

    settings = ServerSettings(
        redis_url="redis://example.invalid/0",
        redis_prefix="test",
        shutdown_grace_seconds=5,
        run_timeout_seconds=17,
        run_retention_seconds=7,
        plugin_daemon_url="http://plugin-daemon",
        plugin_daemon_api_key="daemon-secret",
        inner_api_url="http://dify-api",
        inner_api_key="inner-secret",
        sandbox_files_base_url="http://api:5001",
        local_sandbox_endpoint="http://shellctl",
        local_sandbox_auth_token="shell-secret",
        agent_stub_api_base_url="https://agent.example.com/agent-stub",
        server_secret_key=_base64url_secret(b"1" * 32),
        outbound_http_connect_timeout=1,
        outbound_http_read_timeout=2,
        outbound_http_write_timeout=3,
        outbound_http_pool_timeout=4,
        outbound_http_max_connections=5,
        outbound_http_max_keepalive_connections=3,
        outbound_http_keepalive_expiry=6,
    )

    with TestClient(create_app(settings)):
        assert len(FakeRunScheduler.created) == 1
        scheduler = FakeRunScheduler.created[0]
        assert scheduler.shutdown_grace_seconds == 5
        assert scheduler.run_timeout_seconds == 17
        layer_providers = scheduler.layer_providers
        assert isinstance(layer_providers, tuple)
        execution_context_provider = next(
            provider for provider in layer_providers if provider.type_id == "dify.execution_context"
        )
        execution_context_layer = execution_context_provider.create_layer(
            DifyExecutionContextLayerConfig(
                tenant_id="tenant-1",
                user_from="account",
                agent_mode="workflow_run",
                invoke_from="service-api",
            )
        )
        shell_provider = next(provider for provider in layer_providers if provider.type_id == "dify.shell")
        shell_layer = shell_provider.create_layer(DifyShellLayerConfig())
        assert isinstance(execution_context_layer, DifyExecutionContextLayer)
        assert isinstance(shell_layer, DifyShellLayer)
        assert execution_context_layer.daemon_url == "http://plugin-daemon"
        assert execution_context_layer.daemon_api_key == "daemon-secret"
        assert shell_layer.agent_stub_token_factory is not None
        token = shell_layer.agent_stub_token_factory(_execution_context(), session_id="abc12ff")
        token_codec = settings.create_agent_stub_token_codec()
        assert token_codec is not None
        decoded = token_codec.decode_token(token)
        assert decoded.execution_context == _execution_context()
        assert decoded.session_id == "abc12ff"
        knowledge_provider = next(provider for provider in layer_providers if provider.type_id == "dify.knowledge_base")
        knowledge_layer = knowledge_provider.create_layer(
            DifyKnowledgeBaseLayerConfig.model_validate(
                {
                    "sets": [
                        {
                            "id": "support",
                            "name": "Support KB",
                            "datasets": [{"id": "dataset-1"}],
                            "query": {"mode": "generated_query"},
                            "retrieval": {"mode": "multiple", "top_k": 2},
                        }
                    ],
                }
            )
        )
        assert isinstance(knowledge_layer, DifyKnowledgeBaseLayer)
        assert knowledge_layer.inner_api_url == "http://dify-api"
        assert knowledge_layer.inner_api_key == "inner-secret"
        runtime_provider = next(provider for provider in layer_providers if provider.type_id == "dify.runtime")
        runtime_layer = runtime_provider.create_layer(DifyRuntimeLayerConfig(backend_binding_ref="binding-1"))
        assert isinstance(runtime_layer, DifyRuntimeLayer)
        assert isinstance(runtime_layer.backend, LocalExecutionBindingBackend)
        assert shell_layer.agent_stub_api_base_url == "https://agent.example.com/agent-stub"
        http_client = scheduler.plugin_daemon_http_client
        assert http_client is fake_http_client
        assert http_client.is_closed is False
        assert scheduler.dify_api_http_client is fake_dify_api_http_client
        assert scheduler.dify_api_http_client.is_closed is False
        store = scheduler.store
        assert isinstance(store, RedisRunStore)
        assert store.run_retention_seconds == 7
        assert any(getattr(route, "path", None) == "/agent-stub/connections" for route in create_app(settings).routes)
        assert any(
            getattr(route, "path", None) == "/agent-stub/files/upload-request" for route in create_app(settings).routes
        )
        assert any(
            getattr(route, "path", None) == "/agent-stub/files/download-request"
            for route in create_app(settings).routes
        )
        route_paths = create_app(settings).openapi()["paths"]
        assert {
            "/execution-bindings/files/list",
            "/execution-bindings/files/read",
            "/execution-bindings/files/download",
        }.issubset(route_paths)
        assert "/workspace/files/list" not in route_paths
        assert "/workspace/files/read" not in route_paths
        assert "/workspace/files/upload" not in route_paths

    assert FakeRunScheduler.created[0].shutdown_called is True
    assert FakeRunScheduler.created[0].dify_api_http_client.is_closed is True
    assert FakeRunScheduler.created[0].plugin_daemon_http_client.is_closed is True
    assert fake_redis.closed is True


def test_create_app_wires_authenticated_agent_stub_connection_route(monkeypatch: pytest.MonkeyPatch) -> None:
    fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
    settings = ServerSettings(
        redis_url="redis://example.invalid/0",
        agent_stub_api_base_url="https://agent.example.com/agent-stub",
        server_secret_key=_base64url_secret(b"1" * 32),
    )
    token_codec = settings.create_agent_stub_token_codec()
    assert token_codec is not None
    token = token_codec.encode_connection_token(_execution_context(), now=int(time.time()) - 1)

    with TestClient(create_app(settings)) as client:
        response = client.post(
            "/agent-stub/connections",
            headers={"Authorization": f"Bearer {token}"},
            json={"protocol_version": 1, "argv": ["connect"]},
        )

    assert response.status_code == 200
    assert response.json()["status"] == "connected"
    assert isinstance(response.json()["connection_id"], str)
    assert FakeRunScheduler.created[0].shutdown_called is True
    assert fake_http_client.is_closed is True
    assert fake_redis.closed is True


def test_create_app_wires_authenticated_agent_stub_file_upload_route(monkeypatch: pytest.MonkeyPatch) -> None:
    fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
    settings = ServerSettings(
        redis_url="redis://example.invalid/0",
        agent_stub_api_base_url="https://agent.example.com/agent-stub",
        server_secret_key=_base64url_secret(b"1" * 32),
        inner_api_url="https://api.example.com",
        inner_api_key="inner-secret",
        sandbox_files_base_url="https://files.example.com",
    )
    token_codec = settings.create_agent_stub_token_codec()
    assert token_codec is not None
    token = token_codec.encode_connection_token(_execution_context(), now=int(time.time()) - 1)

    original_async_client = httpx.AsyncClient

    def handler(request: httpx.Request) -> httpx.Response:
        assert str(request.url) == "https://api.example.com/inner/api/agent/files/upload-request"
        assert request.headers["X-Inner-Api-Key"] == "inner-secret"
        return httpx.Response(200, json={"upload_uri": "/files/upload/for-plugin?sign=1"})

    monkeypatch.setattr(
        "dify_agent.agent_stub.server.agent_stub_files.httpx.AsyncClient",
        lambda **kwargs: original_async_client(transport=httpx.MockTransport(handler), **kwargs),
    )

    with TestClient(create_app(settings)) as client:
        response = client.post(
            "/agent-stub/files/upload-request",
            headers={"Authorization": f"Bearer {token}"},
            json={"filename": "report.pdf", "mimetype": "application/pdf"},
        )

    assert response.status_code == 200
    assert response.json() == {"upload_url": "https://files.example.com/files/upload/for-plugin?sign=1"}
    assert FakeRunScheduler.created[0].shutdown_called is True
    assert fake_http_client.is_closed is True
    assert fake_redis.closed is True


def test_create_plugin_daemon_http_client_uses_generic_outbound_httpx_construction_args(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(app_module, "httpx", FakeHttpxModule)

    client = create_plugin_daemon_http_client(
        ServerSettings(
            outbound_http_connect_timeout=1,
            outbound_http_read_timeout=2,
            outbound_http_write_timeout=3,
            outbound_http_pool_timeout=4,
            outbound_http_max_connections=5,
            outbound_http_max_keepalive_connections=3,
            outbound_http_keepalive_expiry=6,
        )
    )

    assert isinstance(client, FakePluginDaemonHttpClient)
    assert client.timeout.connect == 1
    assert client.timeout.read == 2
    assert client.timeout.write == 3
    assert client.timeout.pool == 4
    assert isinstance(client.limits, FakeLimits)
    assert client.limits.max_connections == 5
    assert client.limits.max_keepalive_connections == 3
    assert client.limits.keepalive_expiry == 6
    assert client.trust_env is False


def test_create_dify_api_inner_http_client_uses_generic_outbound_httpx_construction_args(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(app_module, "httpx", FakeHttpxModule)

    client = create_dify_api_inner_http_client(
        ServerSettings(
            outbound_http_connect_timeout=1,
            outbound_http_read_timeout=2,
            outbound_http_write_timeout=3,
            outbound_http_pool_timeout=4,
            outbound_http_max_connections=5,
            outbound_http_max_keepalive_connections=3,
            outbound_http_keepalive_expiry=6,
        )
    )

    assert isinstance(client, FakePluginDaemonHttpClient)
    assert client.timeout.connect == 1
    assert client.timeout.read == 2
    assert client.timeout.write == 3
    assert client.timeout.pool == 4
    assert isinstance(client.limits, FakeLimits)
    assert client.limits.max_connections == 5
    assert client.limits.max_keepalive_connections == 3
    assert client.limits.keepalive_expiry == 6
    assert client.trust_env is False


def test_server_settings_use_generic_outbound_http_args_for_shared_clients() -> None:
    model_fields = ServerSettings.model_fields

    assert "inner_api_url" in model_fields
    assert "inner_api_key" in model_fields
    assert "outbound_http_connect_timeout" in model_fields
    assert "outbound_http_read_timeout" in model_fields
    assert "outbound_http_write_timeout" in model_fields
    assert "outbound_http_pool_timeout" in model_fields
    assert "outbound_http_max_connections" in model_fields
    assert "outbound_http_max_keepalive_connections" in model_fields
    assert "outbound_http_keepalive_expiry" in model_fields