# Copyright 2026 Datastrato, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for ClientSession."""

import asyncio
import json
import unittest
from typing import Any

from adp_sdk.client import ClientSession
from adp_sdk.client.session import basic_auth
from adp_sdk.shared import (
    ADPError,
    InvalidParamsError,
    MethodNotFoundError,
    ResourceNotFoundError,
)
from adp_sdk.transports import Transport
from adp_sdk.types import (
    LATEST_PROTOCOL_VERSION,
    DescribeResult,
    DiscoverResult,
    Implementation,
)
from adp_sdk.types.common import PredicateGroup
from adp_sdk.types.intents import QueryIntent

# -- Helpers -----------------------------------------------------------------


class _MockTransport(Transport):
    """A transport that records written messages and returns canned responses."""

    def __init__(self, responses: list[dict[str, Any]] | None = None) -> None:
        self.responses: list[dict[str, Any]] = responses or []
        self.written: list[str] = []
        self.started = False
        self.closed = False
        self._response_index = 0

    async def start(self) -> None:
        self.started = True

    async def read_message(self) -> str | None:
        if self._response_index >= len(self.responses):
            return None
        resp = self.responses[self._response_index]
        self._response_index += 1
        return json.dumps(resp)

    async def write_message(self, message: str) -> None:
        self.written.append(message)

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


def _make_init_response(request_id: int = 1) -> dict[str, Any]:
    """Create a standard adp.initialize response."""
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "protocolVersion": LATEST_PROTOCOL_VERSION,
            "capabilities": {},
            "serverInfo": {"name": "test-server", "version": "1.0.0"},
        },
    }


def _make_ping_response(request_id: int = 2) -> dict[str, Any]:
    return {"jsonrpc": "2.0", "id": request_id, "result": {}}


def _make_discover_response(request_id: int = 2) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "resources": [
                {
                    "resourceId": "com.example:users",
                    "version": 1,
                    "intentClasses": ["QUERY", "LOOKUP"],
                    "description": "User data",
                }
            ],
        },
    }


def _make_describe_response(request_id: int = 2) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "resourceId": "com.example:users",
            "version": 1,
            "intentClass": "QUERY",
            "usageContract": {
                "fields": [{"fieldId": "id", "type": "INTEGER"}],
                "capabilities": {},
            },
        },
    }


def _make_validate_response(request_id: int = 2) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {"valid": True},
    }


def _make_execute_response(request_id: int = 2) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "results": [{"id": 1, "name": "Alice"}],
        },
    }


def _make_error_response(
    request_id: int,
    code: int = -32001,
    message: str = "Resource not found",
) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": request_id,
        "error": {"code": code, "message": message},
    }


def _make_query_intent(resource_id: str = "com.example:users") -> QueryIntent:
    return QueryIntent(
        intentClass="QUERY",
        resourceId=resource_id,
        predicates=PredicateGroup(op="AND", predicates=[]),
    )


# ============================================================================


class TestClientSessionLifecycle(unittest.IsolatedAsyncioTestCase):
    """Tests for session lifecycle management."""

    async def test_initialize(self) -> None:
        transport = _MockTransport([_make_init_response()])
        session = ClientSession(transport)
        await transport.start()

        result = await session.initialize()

        self.assertTrue(session.initialized)
        self.assertIsNotNone(session.server_info)
        self.assertEqual(session.server_info.name, "test-server")  # type: ignore[union-attr]
        self.assertIsNotNone(session.server_capabilities)
        self.assertEqual(result.protocol_version, LATEST_PROTOCOL_VERSION)

    async def test_initialize_sends_correct_request(self) -> None:
        transport = _MockTransport([_make_init_response()])
        session = ClientSession(transport)
        await transport.start()

        await session.initialize()

        self.assertEqual(len(transport.written), 1)
        data = json.loads(transport.written[0])
        self.assertEqual(data["method"], "adp.initialize")
        self.assertEqual(data["params"]["protocolVersion"], LATEST_PROTOCOL_VERSION)
        self.assertIn("clientInfo", data["params"])

    async def test_initialize_twice_raises(self) -> None:
        transport = _MockTransport([_make_init_response()])
        session = ClientSession(transport)
        await transport.start()
        await session.initialize()

        with self.assertRaises(RuntimeError):
            await session.initialize()

    async def test_initialize_unsupported_protocol_version(self) -> None:
        response = {
            "jsonrpc": "2.0",
            "id": 1,
            "result": {
                "protocolVersion": "1970-01-01",
                "capabilities": {},
                "serverInfo": {"name": "old-server", "version": "0.1.0"},
            },
        }
        transport = _MockTransport([response])
        session = ClientSession(transport)
        await transport.start()

        with self.assertRaises(RuntimeError, msg="Unsupported protocol version"):
            await session.initialize()
        self.assertFalse(session.initialized)

    async def test_context_manager(self) -> None:
        transport = _MockTransport([_make_init_response()])
        async with ClientSession(transport) as session:
            self.assertTrue(session.initialized)
            self.assertTrue(transport.started)
        self.assertTrue(transport.closed)

    async def test_custom_client_info(self) -> None:
        info = Implementation(name="my-agent", version="2.0")
        transport = _MockTransport([_make_init_response()])
        session = ClientSession(transport, client_info=info)
        await transport.start()
        await session.initialize()

        data = json.loads(transport.written[0])
        self.assertEqual(data["params"]["clientInfo"]["name"], "my-agent")
        self.assertEqual(data["params"]["clientInfo"]["version"], "2.0")

    async def test_not_initialized_raises(self) -> None:
        transport = _MockTransport()
        session = ClientSession(transport)
        await transport.start()

        with self.assertRaises(RuntimeError):
            await session.ping()

    async def test_properties_before_init(self) -> None:
        transport = _MockTransport()
        session = ClientSession(transport)
        self.assertFalse(session.initialized)
        self.assertIsNone(session.server_info)
        self.assertIsNone(session.server_capabilities)


# ============================================================================


class TestClientSessionPing(unittest.IsolatedAsyncioTestCase):
    """Tests for the ping operation."""

    async def test_ping(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_ping_response()])
        async with ClientSession(transport) as session:
            result = await session.ping()
            self.assertIsNotNone(result)

    async def test_ping_sends_correct_request(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_ping_response()])
        async with ClientSession(transport) as session:
            await session.ping()
            data = json.loads(transport.written[1])
            self.assertEqual(data["method"], "adp.ping")
            self.assertEqual(data["id"], 2)


# ============================================================================


class TestClientSessionDiscover(unittest.IsolatedAsyncioTestCase):
    """Tests for the discover operation."""

    async def test_discover(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_discover_response()])
        async with ClientSession(transport) as session:
            result = await session.discover()
            self.assertEqual(len(result.resources), 1)
            self.assertEqual(result.resources[0].resource_id, "com.example:users")

    async def test_discover_with_filter(self) -> None:
        from adp_sdk.types import DiscoverFilter

        transport = _MockTransport([_make_init_response(), _make_discover_response()])
        async with ClientSession(transport) as session:
            await session.discover(
                filter=DiscoverFilter(domainPrefix="com.example"),
            )
            data = json.loads(transport.written[1])
            self.assertEqual(data["params"]["filter"]["domainPrefix"], "com.example")

    async def test_discover_with_cursor(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_discover_response()])
        async with ClientSession(transport) as session:
            await session.discover(cursor="page2")
            data = json.loads(transport.written[1])
            self.assertEqual(data["params"]["cursor"], "page2")


# ============================================================================


class TestClientSessionDescribe(unittest.IsolatedAsyncioTestCase):
    """Tests for the describe operation."""

    async def test_describe(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_describe_response()])
        async with ClientSession(transport) as session:
            result = await session.describe("com.example:users", "QUERY")
            self.assertEqual(result.resource_id, "com.example:users")
            self.assertEqual(result.intent_class, "QUERY")
            self.assertIsNotNone(result.usage_contract)

    async def test_describe_with_version(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_describe_response()])
        async with ClientSession(transport) as session:
            await session.describe("com.example:users", "QUERY", version=2)
            data = json.loads(transport.written[1])
            self.assertEqual(data["params"]["version"], 2)


# ============================================================================


class TestClientSessionValidate(unittest.IsolatedAsyncioTestCase):
    """Tests for the validate operation."""

    async def test_validate(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_validate_response()])
        async with ClientSession(transport) as session:
            intent = _make_query_intent()
            result = await session.validate(intent)
            self.assertTrue(result.valid)

    async def test_validate_sends_correct_request(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_validate_response()])
        async with ClientSession(transport) as session:
            intent = _make_query_intent()
            await session.validate(intent)
            data = json.loads(transport.written[1])
            self.assertEqual(data["method"], "adp.validate")
            self.assertEqual(data["params"]["intent"]["resourceId"], "com.example:users")
            self.assertEqual(data["params"]["intent"]["intentClass"], "QUERY")


# ============================================================================


class TestClientSessionExecute(unittest.IsolatedAsyncioTestCase):
    """Tests for the execute operation."""

    async def test_execute(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_execute_response()])
        async with ClientSession(transport) as session:
            intent = _make_query_intent()
            result = await session.execute(intent)
            self.assertEqual(len(result.results), 1)
            self.assertEqual(result.results[0]["name"], "Alice")

    async def test_execute_with_cursor(self) -> None:
        transport = _MockTransport([_make_init_response(), _make_execute_response()])
        async with ClientSession(transport) as session:
            intent = _make_query_intent()
            await session.execute(intent, cursor="next")
            data = json.loads(transport.written[1])
            self.assertEqual(data["params"]["cursor"], "next")


# ============================================================================


class TestClientSessionErrors(unittest.IsolatedAsyncioTestCase):
    """Tests for error handling."""

    async def test_server_error_raises_adp_error(self) -> None:
        transport = _MockTransport(
            [
                _make_init_response(),
                _make_error_response(2, -32001, "Resource not found"),
            ]
        )
        async with ClientSession(transport) as session:
            with self.assertRaises(ResourceNotFoundError):
                await session.discover()

    async def test_method_not_found_error(self) -> None:
        transport = _MockTransport(
            [
                _make_init_response(),
                _make_error_response(2, -32601, "Method not found"),
            ]
        )
        async with ClientSession(transport) as session:
            with self.assertRaises(MethodNotFoundError):
                await session.ping()

    async def test_invalid_params_error(self) -> None:
        transport = _MockTransport(
            [
                _make_init_response(),
                _make_error_response(2, -32602, "Invalid params"),
            ]
        )
        async with ClientSession(transport) as session:
            with self.assertRaises(InvalidParamsError):
                await session.discover()

    async def test_unknown_error_code(self) -> None:
        transport = _MockTransport(
            [
                _make_init_response(),
                _make_error_response(2, -99999, "Something weird"),
            ]
        )
        async with ClientSession(transport) as session:
            with self.assertRaises(ADPError):
                await session.discover()

    async def test_connection_lost(self) -> None:
        transport = _MockTransport([_make_init_response()])
        async with ClientSession(transport) as session:
            with self.assertRaises(ConnectionError):
                await session.ping()

    async def test_init_error(self) -> None:
        transport = _MockTransport(
            [
                _make_error_response(1, -32600, "Bad init request"),
            ]
        )
        session = ClientSession(transport)
        await transport.start()
        with self.assertRaises(ADPError):
            await session.initialize()
        self.assertFalse(session.initialized)


# ============================================================================


class TestClientSessionRequestIds(unittest.IsolatedAsyncioTestCase):
    """Tests for request ID management."""

    async def test_incremental_ids(self) -> None:
        transport = _MockTransport(
            [
                _make_init_response(1),
                _make_ping_response(2),
                _make_ping_response(3),
            ]
        )
        async with ClientSession(transport) as session:
            await session.ping()
            await session.ping()

            ids = [json.loads(w)["id"] for w in transport.written]
            self.assertEqual(ids, [1, 2, 3])

    async def test_malformed_response_no_result(self) -> None:
        """Response with neither result nor error should raise ConnectionError."""
        transport = _MockTransport(
            [
                _make_init_response(),
                {"jsonrpc": "2.0", "id": 2},
            ]
        )
        async with ClientSession(transport) as session:
            with self.assertRaises(ConnectionError):
                await session.ping()

    async def test_session_not_initialized_after_close(self) -> None:
        """After close, session should be marked as not initialized."""
        transport = _MockTransport([_make_init_response()])
        session = ClientSession(transport)
        await transport.start()
        await session.initialize()
        self.assertTrue(session.initialized)
        await session.close()
        self.assertFalse(session.initialized)
        with self.assertRaises(RuntimeError):
            await session.ping()


# ============================================================================


class TestClientSessionAuthorization(unittest.IsolatedAsyncioTestCase):
    """Tests for authorization support."""

    async def test_authorization_injected_in_meta(self) -> None:
        """Authorization is injected into _meta of every request."""
        transport = _MockTransport([_make_init_response(), _make_discover_response()])
        async with ClientSession(transport, authorization="Basic dXNlcjpwYXNz") as session:
            await session.discover()

            for msg in transport.written:
                data = json.loads(msg)
                self.assertEqual(data["params"]["_meta"]["authorization"], "Basic dXNlcjpwYXNz")

    async def test_no_authorization_by_default(self) -> None:
        """Without authorization, _meta should not contain authorization."""
        transport = _MockTransport([_make_init_response(), _make_ping_response()])
        async with ClientSession(transport) as session:
            await session.ping()

            data = json.loads(transport.written[1])
            meta = data["params"].get("_meta", {})
            self.assertNotIn("authorization", meta)


# ============================================================================


class TestBasicAuth(unittest.TestCase):
    """Tests for basic_auth helper function."""

    def test_basic_auth_with_password(self) -> None:
        result = basic_auth("alice", "secret")
        # base64("alice:secret") == "YWxpY2U6c2VjcmV0"
        self.assertEqual(result, "Basic YWxpY2U6c2VjcmV0")

    def test_basic_auth_without_password(self) -> None:
        result = basic_auth("alice")
        # base64("alice:") == "YWxpY2U6"
        self.assertEqual(result, "Basic YWxpY2U6")

    def test_basic_auth_round_trip(self) -> None:
        """Authorization value produced by basic_auth is accepted by ClientSession."""
        import base64

        auth = basic_auth("user", "pass")
        _, _, payload = auth.partition(" ")
        decoded = base64.b64decode(payload).decode()
        self.assertEqual(decoded, "user:pass")

    def test_basic_auth_empty_username_raises(self) -> None:
        with self.assertRaises(ValueError, msg="username must not be empty"):
            basic_auth("")

    def test_basic_auth_username_with_colon_raises(self) -> None:
        with self.assertRaises(ValueError, msg="username must not contain a colon character"):
            basic_auth("user:name", "pass")


# ============================================================
# Concurrent access safety
# ============================================================


class _DelayedMockTransport(Transport):
    """Mock transport that introduces async delay to simulate real I/O."""

    def __init__(self, responses: list[str], delay: float = 0.05) -> None:
        self.responses = list(responses)
        self.written: list[str] = []
        self._delay = delay
        self._index = 0
        self._started = False

    async def start(self) -> None:
        self._started = True

    async def close(self) -> None:
        self._started = False

    async def write_message(self, message: str) -> None:
        self.written.append(message)

    async def read_message(self) -> str | None:
        await asyncio.sleep(self._delay)
        if self._index >= len(self.responses):
            return None
        resp = self.responses[self._index]
        self._index += 1
        return resp


class TestClientSessionConcurrency(unittest.IsolatedAsyncioTestCase):
    """Verify that concurrent protocol calls are serialized by _send_lock."""

    async def test_concurrent_calls_do_not_swap_responses(self) -> None:
        """Two concurrent requests must each receive their own response."""
        discover_response = json.dumps(
            {
                "jsonrpc": "2.0",
                "id": 2,
                "result": {
                    "resources": [{"resourceId": "t:x", "version": 1, "intentClasses": ["QUERY"]}],
                },
            }
        )
        describe_response = json.dumps(
            {
                "jsonrpc": "2.0",
                "id": 3,
                "result": {
                    "resourceId": "t:x",
                    "version": 1,
                    "intentClass": "QUERY",
                    "usageContract": {
                        "fields": [{"fieldId": "id", "type": "INTEGER"}],
                        "capabilities": {},
                    },
                },
            }
        )
        transport = _DelayedMockTransport(
            responses=[json.dumps(_make_init_response()), discover_response, describe_response],
            delay=0.05,
        )
        async with ClientSession(transport) as session:
            discover_result, describe_result = await asyncio.gather(
                session.discover(),
                session.describe("t:x", "QUERY"),
            )
            self.assertIsInstance(discover_result, DiscoverResult)
            self.assertIsInstance(describe_result, DescribeResult)
            self.assertEqual(len(discover_result.resources), 1)
            self.assertEqual(describe_result.resource_id, "t:x")

    async def test_concurrent_calls_no_runtime_error(self) -> None:
        """Concurrent calls must not raise RuntimeError from transport."""
        responses: list[str] = [json.dumps(_make_init_response())]
        for i in range(2, 12):
            responses.append(
                json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": i,
                        "result": {"resources": []},
                    }
                )
            )
        transport = _DelayedMockTransport(responses=responses, delay=0.01)
        async with ClientSession(transport) as session:
            tasks = [session.discover() for _ in range(10)]
            results = await asyncio.gather(*tasks)
            for result in results:
                self.assertIsInstance(result, DiscoverResult)
