"""Tests for YukkaClient, Session, and connection handling."""

import base64
import io
import json
import os
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta
from unittest.mock import MagicMock, Mock, patch

import polars as pl
import pytest
from httpx import ConnectError, HTTPStatusError, Request, Response, TimeoutException

from yukka.auth import Token
from yukka.client import Session, YukkaAPI, YukkaClient
from yukka.entities import Asset
from yukka.exceptions import ExpiredToken


def create_jwt_token(
    email: str = "test@yukkalab.com",
    groups: list[str] | None = None,
    exp_offset_seconds: int = 3600,
    iat_offset_seconds: int = 0,
) -> str:
    """Create a valid JWT token for testing."""
    if groups is None:
        groups = ["users", "quant"]

    now = datetime.now(tz=UTC)
    payload = {
        "exp": (now + timedelta(seconds=exp_offset_seconds)).timestamp(),
        "iat": (now + timedelta(seconds=iat_offset_seconds)).timestamp(),
        "email": email,
        "groups": groups,
    }

    header = base64.urlsafe_b64encode(json.dumps({"alg": "HS256", "typ": "JWT"}).encode()).decode().rstrip("=")
    payload_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
    signature = "fake_signature_for_testing"

    return f"{header}.{payload_b64}.{signature}"


def create_sample_parquet_bytes() -> bytes:
    """Create sample parquet data as bytes for testing (raw API format)."""
    df = pl.DataFrame(
        {
            "publish_time": [20454, 20454, 20454, 20455, 20455, 20455],  # days since epoch
            "entity_compound_key": [b"company:bmw"] * 6,
            "sentiment": [1, 0, -1, 1, 0, -1],  # pos, neu, neg for each day
            "count()": [200, 94, 115, 396, 201, 132],
        }
    )
    buffer = io.BytesIO()
    df.write_parquet(buffer)
    return buffer.getvalue()


class MockStreamResponse:
    """Mock streaming response for testing."""

    def __init__(self, parquet_bytes: bytes, status_code: int = 200, raise_error: Exception | None = None):
        """Initialise mock with raw parquet bytes and optional error to raise."""
        self.status_code = status_code
        self._data = parquet_bytes
        self._raise_error = raise_error

    def raise_for_status(self):
        """Raise the configured error, if any."""
        if self._raise_error:
            raise self._raise_error

    def iter_bytes(self, chunk_size: int = 10_000):
        """Yield parquet data in chunks."""
        for i in range(0, len(self._data), chunk_size):
            yield self._data[i : i + chunk_size]

    def __enter__(self):
        """Enter context manager."""
        return self

    def __exit__(self, *args):
        """Exit context manager."""


class TestYukkaClient:
    """Tests for YukkaClient class."""

    @pytest.fixture
    def client(self, valid_token_str: str):
        """Create a YukkaClient with valid token."""
        client = YukkaClient.from_token(valid_token_str)
        yield client
        client.close()

    @pytest.fixture
    def mock_client(self, valid_token_str: str, mock_http_client: MagicMock) -> YukkaClient:
        """Create a YukkaClient with mocked HTTP client."""
        from yarl import URL

        token = Token.from_token(valid_token_str)
        return YukkaClient(
            _token=token,
            _client=mock_http_client,
            _data_api_url=URL("https://data-dev.api.yukkalab.com/9-3-5/2023-12-13-0/"),
        )

    def test_client_from_token(self, valid_token_str: str):
        """Client should be created from valid token."""
        client = YukkaClient.from_token(valid_token_str)

        assert isinstance(client, YukkaClient)
        client.close()

    def test_client_is_abstract_api_implementation(self, client: YukkaClient):
        """YukkaClient should implement YukkaAPI."""
        assert isinstance(client, YukkaAPI)

    def test_client_context_manager(self, valid_token_str: str):
        """Client should work as context manager."""
        with YukkaClient.from_token(valid_token_str) as client:
            assert isinstance(client, YukkaClient)
        # Client should be closed after exiting context

    def test_sentiment_success(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Sentiment should be fetched successfully via streaming parquet."""
        parquet_bytes = create_sample_parquet_bytes()

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(parquet_bytes)

        mock_http_client.stream = stream_context

        df = mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert isinstance(df, pl.DataFrame)
        assert len(df) == 2  # 2 days of data
        assert list(df.columns) == ["date", "entity", "positive", "neutral", "negative"]
        assert df["positive"].to_list() == [200, 396]
        assert df["neutral"].to_list() == [94, 201]
        assert df["negative"].to_list() == [115, 132]

    def test_sentiment_empty_entities_returns_empty_df(self, mock_client: YukkaClient):
        """Empty entity list should return empty DataFrame."""
        df = mock_client.sentiment([], "2026-01-01", "2026-01-03")

        assert isinstance(df, pl.DataFrame)
        assert len(df) == 0

    def test_sentiment_calls_correct_endpoint(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Sentiment should call correct API endpoint."""
        parquet_bytes = create_sample_parquet_bytes()
        captured_args = {}

        @contextmanager
        def stream_context(*args, **kwargs):
            captured_args["args"] = args
            captured_args["kwargs"] = kwargs
            yield MockStreamResponse(parquet_bytes)

        mock_http_client.stream = stream_context

        mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        call_url = captured_args["args"][1]
        assert "api/exports/sentiment_matches" in call_url

    def test_sentiment_includes_auth_headers(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Sentiment request should include Authorization header."""
        parquet_bytes = create_sample_parquet_bytes()
        captured_kwargs = {}

        @contextmanager
        def stream_context(*args, **kwargs):
            captured_kwargs.update(kwargs)
            yield MockStreamResponse(parquet_bytes)

        mock_http_client.stream = stream_context

        mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert "headers" in captured_kwargs
        assert "Authorization" in captured_kwargs["headers"]
        assert captured_kwargs["headers"]["Authorization"].startswith("Bearer ")

    def test_sentiment_request_body_format(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Sentiment request should have correct body format."""
        parquet_bytes = create_sample_parquet_bytes()
        captured_kwargs = {}

        @contextmanager
        def stream_context(*args, **kwargs):
            captured_kwargs.update(kwargs)
            yield MockStreamResponse(parquet_bytes)

        mock_http_client.stream = stream_context

        mock_client.sentiment(["company:bmw", "company:tesla_motors"], "2026-01-01", "2026-01-03")

        json_body = captured_kwargs["json"]
        assert "time_range" in json_body
        assert "sentiment_filter" in json_body
        assert "language_filter" in json_body
        assert json_body["sentiment_filter"]["entities"] == ["company:bmw", "company:tesla_motors"]
        assert json_body["language_filter"] == ["en", "de"]


class TestYukkaClientHTTPErrors:
    """Tests for HTTP error handling in YukkaClient."""

    @pytest.fixture
    def mock_client(self, valid_token_str: str, mock_http_client: MagicMock) -> YukkaClient:
        """Create a YukkaClient with mocked HTTP client."""
        from yarl import URL

        token = Token.from_token(valid_token_str)
        return YukkaClient(
            _token=token,
            _client=mock_http_client,
            _data_api_url=URL("https://data-dev.api.yukkalab.com/9-3-5/2023-12-13-0/"),
        )

    def test_sentiment_401_unauthorized(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """401 unauthorized should propagate as HTTPStatusError."""
        mock_request = Mock(spec=Request)
        mock_response = Mock(spec=Response)
        mock_response.status_code = 401
        error = HTTPStatusError("Unauthorized", request=mock_request, response=mock_response)

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(b"", status_code=401, raise_error=error)

        mock_http_client.stream = stream_context

        with pytest.raises(HTTPStatusError) as exc_info:
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert exc_info.value.response.status_code == 401

    def test_sentiment_403_forbidden(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """403 forbidden should propagate as HTTPStatusError."""
        mock_request = Mock(spec=Request)
        mock_response = Mock(spec=Response)
        mock_response.status_code = 403
        error = HTTPStatusError("Forbidden", request=mock_request, response=mock_response)

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(b"", status_code=403, raise_error=error)

        mock_http_client.stream = stream_context

        with pytest.raises(HTTPStatusError) as exc_info:
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert exc_info.value.response.status_code == 403

    def test_sentiment_404_not_found(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """404 not found should propagate as HTTPStatusError."""
        mock_request = Mock(spec=Request)
        mock_response = Mock(spec=Response)
        mock_response.status_code = 404
        error = HTTPStatusError("Not Found", request=mock_request, response=mock_response)

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(b"", status_code=404, raise_error=error)

        mock_http_client.stream = stream_context

        with pytest.raises(HTTPStatusError) as exc_info:
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert exc_info.value.response.status_code == 404

    def test_sentiment_500_server_error(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """500 server error should propagate as HTTPStatusError."""
        mock_request = Mock(spec=Request)
        mock_response = Mock(spec=Response)
        mock_response.status_code = 500
        error = HTTPStatusError("Internal Server Error", request=mock_request, response=mock_response)

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(b"", status_code=500, raise_error=error)

        mock_http_client.stream = stream_context

        with pytest.raises(HTTPStatusError) as exc_info:
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert exc_info.value.response.status_code == 500

    def test_sentiment_429_rate_limit(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """429 rate limit should propagate as HTTPStatusError."""
        mock_request = Mock(spec=Request)
        mock_response = Mock(spec=Response)
        mock_response.status_code = 429
        error = HTTPStatusError("Too Many Requests", request=mock_request, response=mock_response)

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(b"", status_code=429, raise_error=error)

        mock_http_client.stream = stream_context

        with pytest.raises(HTTPStatusError) as exc_info:
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        assert exc_info.value.response.status_code == 429


class TestYukkaClientConnectionErrors:
    """Tests for connection error handling in YukkaClient."""

    @pytest.fixture
    def mock_client(self, valid_token_str: str, mock_http_client: MagicMock) -> YukkaClient:
        """Create a YukkaClient with mocked HTTP client."""
        from yarl import URL

        token = Token.from_token(valid_token_str)
        return YukkaClient(
            _token=token,
            _client=mock_http_client,
            _data_api_url=URL("https://data-dev.api.yukkalab.com/9-3-5/2023-12-13-0/"),
        )

    def test_sentiment_connection_error(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Connection error should propagate."""

        @contextmanager
        def stream_context(*args, **kwargs):
            raise ConnectError("Connection refused")
            yield  # type: ignore[misc]

        mock_http_client.stream = stream_context

        with pytest.raises(ConnectError):
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

    def test_sentiment_timeout_error(self, mock_client: YukkaClient, mock_http_client: MagicMock):
        """Timeout error should propagate."""

        @contextmanager
        def stream_context(*args, **kwargs):
            raise TimeoutException("Request timed out")
            yield  # type: ignore[misc]

        mock_http_client.stream = stream_context

        with pytest.raises(TimeoutException):
            mock_client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")


class TestYukkaClientExpiredToken:
    """Tests for expired token handling in YukkaClient."""

    def test_sentiment_with_expired_token(self, expired_token_str: str):
        """Expired token should raise ExpiredToken when making request."""
        client = YukkaClient.from_token(expired_token_str)

        with pytest.raises(ExpiredToken):
            client.sentiment(["company:bmw"], "2026-01-01", "2026-01-03")

        client.close()


class TestSession:
    """Tests for Session class."""

    def test_session_from_explicit_token(self, valid_token_str: str):
        """Session should be created from explicit token."""
        with Session(token=valid_token_str) as session:
            assert isinstance(session._client, YukkaClient)
            assert session.resolver is not None

    def test_session_from_env_variable(self, valid_token_str: str):
        """Session should be created from YUKKA_TOKEN env variable."""
        with patch.dict(os.environ, {"YUKKA_TOKEN": valid_token_str}), Session() as session:
            assert isinstance(session._client, YukkaClient)

    def test_session_no_token_raises_error(self):
        """Session without token should raise ValueError."""
        with patch.dict(os.environ, {}, clear=True):
            # Ensure YUKKA_TOKEN is not set
            if "YUKKA_TOKEN" in os.environ:
                del os.environ["YUKKA_TOKEN"]

            with pytest.raises(ValueError, match="token") as exc_info:
                Session()

            assert "token" in str(exc_info.value).lower()

    def test_session_context_manager(self, valid_token_str: str):
        """Session should work as context manager."""
        with Session(token=valid_token_str) as session:
            assert isinstance(session, Session)

    def test_session_today_property(self, valid_token_str: str):
        """Session.today should return current date."""
        from datetime import date

        with Session(token=valid_token_str) as session:
            assert session.today == date.today()

    def test_session_resolver_property(self, valid_token_str: str):
        """Session.resolver should return EntityResolver."""
        from yukka.entities import EntityResolver

        with Session(token=valid_token_str) as session:
            assert isinstance(session.resolver, EntityResolver)

    def test_session_sentiment_with_single_asset(self, valid_token_str: str):
        """Session.sentiment should work with single Asset."""
        parquet_bytes = create_sample_parquet_bytes()

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(parquet_bytes)

        with patch("yukka.client.Client") as mock_client, Session(token=valid_token_str) as session:
            mock_client.return_value.stream = stream_context
            asset = Asset.from_yukka_id("company:bmw")
            df = session.sentiment(asset, "2026-01-01", "2026-01-03")

            assert isinstance(df, pl.DataFrame)

    def test_session_sentiment_with_asset_list(self, valid_token_str: str):
        """Session.sentiment should work with list of Assets."""
        parquet_bytes = create_sample_parquet_bytes()

        @contextmanager
        def stream_context(*args, **kwargs):
            yield MockStreamResponse(parquet_bytes)

        with patch("yukka.client.Client") as mock_client, Session(token=valid_token_str) as session:
            mock_client.return_value.stream = stream_context
            assets = Asset.from_yukka_id(["company:bmw", "company:apple"])
            df = session.sentiment(assets, "2026-01-01", "2026-01-03")

            assert isinstance(df, pl.DataFrame)

    def test_session_sentiment_defaults_date_to_today(self, valid_token_str: str):
        """Session.sentiment should default date_to to today."""
        parquet_bytes = create_sample_parquet_bytes()
        stream_called = False

        @contextmanager
        def stream_context(*args, **kwargs):
            nonlocal stream_called
            stream_called = True
            yield MockStreamResponse(parquet_bytes)

        with patch("yukka.client.Client") as mock_client, Session(token=valid_token_str) as session:
            mock_client.return_value.stream = stream_context
            asset = Asset.from_yukka_id("company:bmw")
            # date_to should default to today
            session.sentiment(asset, "2026-01-01")

            assert stream_called


class TestSessionAuthenticationErrors:
    """Tests for authentication error handling in Session."""

    def test_session_with_expired_token(self, expired_token_str: str):
        """Session with expired token should raise on API call."""
        with Session(token=expired_token_str) as session:
            asset = Asset.from_yukka_id("company:bmw")

            with pytest.raises(ExpiredToken):
                session.sentiment(asset, "2026-01-01", "2026-01-03")

    def test_session_explicit_token_takes_precedence(self, valid_token_str: str):
        """Explicit token should take precedence over env variable."""
        different_token = create_jwt_token(email="different@yukkalab.com")

        with patch.dict(os.environ, {"YUKKA_TOKEN": different_token}), Session(token=valid_token_str) as session:
            # Should use explicit token, not env var
            assert session._client._token._payload.email == "test@yukkalab.com"


class TestYukkaAPIAbstract:
    """Tests for the YukkaAPI abstract base class."""

    def test_abstract_sentiment_body_returns_none(self):
        """Abstract sentiment method body (pass) should return None when called via super()."""

        class ConcreteAPI(YukkaAPI):
            def sentiment(self, entities, date_from, date_to):
                return super().sentiment(entities, date_from, date_to)

        api = ConcreteAPI()
        result = api.sentiment([], "2026-01-01", "2026-01-07")
        assert result is None
