"""WebSocket /v1/stream: auth, subscribe, delivery of published events, resume, heartbeat, connection limit."""
from __future__ import annotations

import pytest
from starlette.websockets import WebSocketDisconnect



@pytest.fixture
def keys(make_user):
    """Two real high_usage API keys (the rate-limit middleware validates keys before the stream handler)."""
    _, k1, _, _ = make_user(tier="high_usage")
    _, k2, _, _ = make_user(tier="high_usage")
    return k1, k2


@pytest.fixture(autouse=True)
def _reset(app):
    from stream import broker
    broker.reset_for_tests()
    yield


def test_keyless_is_refused_with_json_reason(client):
    with client.websocket_connect("/v1/stream") as ws:
        msg = ws.receive_json()
        assert msg["type"] == "error" and msg["code"] == "AUTH_REQUIRED" and msg["docs"].endswith("#auth_required")
        with pytest.raises(WebSocketDisconnect) as e:
            ws.receive_json()
        assert e.value.code == 4001


def test_bad_key_is_refused(client):
    try:
        with client.websocket_connect("/v1/stream", headers={"Authorization": "Bearer not-a-key"}) as ws:
            msg = ws.receive_json()
            assert msg["code"] in ("AUTH_REQUIRED", "INVALID_API_KEY")
    except WebSocketDisconnect as e:   # refused by the rate-limit middleware before the handler
        assert e.code in (4401, 4001, 1008)


def _key_id(raw: str) -> int:
    from sqlalchemy import select

    from accounts.models import ApiKey
    from accounts.security import hash_key
    from core.db import session
    with session() as s:
        return int(s.scalar(select(ApiKey.id).where(ApiKey.key_hash == hash_key(raw))))


def test_principal_is_the_key_id_not_a_hash(keys):
    """The stream must bill the same `key:<id>` principal as the HTTP limiter (a hashed pseudo-principal fed an
    orphan counter and resolved to the free tier)."""
    import hashlib

    from core.config import settings
    from stream.auth import authenticate, principal_from_scope
    kid = _key_id(keys[0])
    assert authenticate(keys[0]) == f"key:{kid}"
    fake = "key:" + hashlib.sha256((settings.key_hash_salt + keys[0]).encode()).hexdigest()[:16]
    assert authenticate(keys[0]) != fake
    # well-formed but unknown keys are refused, whatever their shape
    assert authenticate("hfmd_live_" + "a" * 32) is None and authenticate("not-a-key") is None and authenticate(None) is None
    assert principal_from_scope({"state": {"principal": "key:7", "principal_kind": "key", "tier": "free"}}) == ("key:7", "key", "free")
    assert principal_from_scope({}) == (None, None, None)


def test_well_formed_unknown_key_is_refused(client):
    with pytest.raises(WebSocketDisconnect) as e:
        with client.websocket_connect("/v1/stream?api_key=hfmd_live_" + "Z" * 32) as ws:
            ws.receive_json()
    assert e.value.code in (4401, 4001)


def test_handler_refuses_without_middleware_context(keys):
    """`_resolve_principal`: middleware state wins; no state while the limiter is enabled → refuse, never guess."""
    import types

    from core.config import settings
    from stream.routes import _resolve_principal
    ws = types.SimpleNamespace(scope={"state": {"principal": "key:42", "principal_kind": "key", "tier": "high_usage"}},
                               headers={"authorization": f"Bearer {keys[0]}"})
    assert _resolve_principal(ws, None) == ("key:42", "high_usage")
    ws.scope = {"state": {"principal": "ip:abc", "principal_kind": "keyless", "tier": "keyless"}}
    assert _resolve_principal(ws, keys[0]) == (None, "keyless")
    ws.scope = {}
    expected = (None, None) if settings.ratelimit_enabled else (f"key:{_key_id(keys[0])}", None)
    assert _resolve_principal(ws, keys[0]) == expected


def test_subscribe_receive_and_accounting(keys, client, fundamentals_data):
    from stream import accounting, broker
    from stream.auth import authenticate
    principal = authenticate(keys[0])
    assert principal == f"key:{_key_id(keys[0])}"
    before = accounting.rows_charged(principal)
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        hello = ws.receive_json()
        assert hello["type"] == "hello" and hello["heartbeat_seconds"] == 20
        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": ["AAPL"], "forms": ["10-Q", "10-K"]})
        sub = ws.receive_json()
        assert sub["type"] == "subscribed" and sub["tickers"] == ["AAPL"] and sub["forms"] == ["10-K", "10-Q"]
        # events published by the ingestion (Redis) — one matching, two filtered out
        broker.publish({"type": "filing", "ticker": "MSFT", "form": "10-Q", "accn": "x1"})
        broker.publish({"type": "filing", "ticker": "AAPL", "form": "8-K", "accn": "x2"})
        seq = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-Q", "accn": "x3", "filed_date": "2024-05-03",
                              "summary": {"revenue": 90_753_000_000}})
        ev = ws.receive_json()
        assert ev["type"] == "filing" and ev["accn"] == "x3" and ev["seq"] == seq and ev["summary"]["revenue"] == 90_753_000_000
        ws.send_json({"action": "ping"})
        pong = ws.receive_json()
        assert pong["type"] == "pong" and pong["seq"] == seq
    assert accounting.rows_charged(principal) == before + 1


def test_real_filing_event_from_ingest(keys, client, fundamentals_data):
    from datetime import date

    from fundamentals import ingest
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        ws.receive_json()
        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": "all"})
        ws.receive_json()
        ev = ingest.publish_filing_event(320193, "AAPL", "10-Q", date(2024, 5, 3), date(2024, 3, 30), "0000320193-24-000069",
                                         "aapl-20240330.htm")
        assert ev is not None
        got = ws.receive_json()
        assert got["ticker"] == "AAPL" and got["form"] == "10-Q" and got["filed_date"] == "2024-05-03"
        assert got["url"].endswith("/000032019324000069/aapl-20240330.htm")
        s = got["summary"]
        assert s["revenue"] == 90_753_000_000 and s["net_income"] == 23_636_000_000 and s["eps_diluted"] == 1.53
        assert s["total_assets"] == 337_411_000_000 and s["operating_cash_flow"] == 22_690_000_000
        assert s["yoy"]["revenue"] == pytest.approx(90_753 / 94_836 - 1, rel=1e-4)


def test_resume_token_replays_buffer(keys, client):
    from stream import broker
    s1 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r1"})
    s2 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r2"})
    s3 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r3"})
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        assert ws.receive_json()["seq"] == s3
        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": ["AAPL"], "resume_token": s1})
        assert ws.receive_json()["resume_from"] == s1
        assert [ws.receive_json()["accn"] for _ in range(2)] == ["r2", "r3"]
        assert s2 < s3
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:   # without resume: only new events
        ws.receive_json()
        ws.send_json({"action": "subscribe"})
        ws.receive_json()
        broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r4"})
        assert ws.receive_json()["accn"] == "r4"


def test_protocol_errors(keys, client):
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        ws.receive_json()
        ws.send_text("not json")
        assert ws.receive_json()["code"] == "VALIDATION_ERROR"
        ws.send_json({"action": "subscribe", "channel": "trades"})
        assert ws.receive_json()["code"] == "NOT_FOUND"
        ws.send_json({"action": "subscribe", "tickers": []})
        assert ws.receive_json()["code"] == "VALIDATION_ERROR"
        ws.send_json({"action": "dance"})
        assert ws.receive_json()["code"] == "VALIDATION_ERROR"
        ws.send_json({"action": "subscribe", "resume_token": "abc"})
        assert ws.receive_json()["code"] == "VALIDATION_ERROR"


def test_heartbeat(keys, client, monkeypatch):
    from stream import routes
    monkeypatch.setattr(routes, "HEARTBEAT_SECONDS", 0.3)
    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:
        ws.receive_json()
        beat = ws.receive_json()
        assert beat["type"] == "heartbeat" and "ts" in beat and "seq" in beat


def test_connection_limit_per_key(keys, client):
    from contextlib import ExitStack

    from stream import routes
    with ExitStack() as stack:
        for _ in range(routes.MAX_CONNECTIONS_PER_KEY):
            ws = stack.enter_context(client.websocket_connect(f"/v1/stream?api_key={keys[0]}"))
            assert ws.receive_json()["type"] == "hello"
        extra = stack.enter_context(client.websocket_connect(f"/v1/stream?api_key={keys[0]}"))
        msg = extra.receive_json()
        assert msg["type"] == "error" and msg["code"] == "STREAM_CONNECTION_LIMIT" and msg["limit"] == 5
        with pytest.raises(WebSocketDisconnect) as e:
            extra.receive_json()
        assert e.value.code == 4029
    # slots are released on disconnect
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        assert ws.receive_json()["type"] == "hello"


def test_phantom_connections_are_pruned(keys, client):
    """A worker killed mid-stream never releases its slot: members whose heartbeat is older than 90 s must not
    count (the old INCR/DECR counter stayed inflated for an hour and produced unjustified 4029s)."""
    import time

    from stream import broker, routes
    principal = f"key:{_key_id(keys[1])}"
    r = broker.get_redis()
    key = routes.conn_key(principal)
    r.delete(key)
    stale = time.time() - routes.CONN_STALE_SECONDS - 5
    r.zadd(key, {f"dead-{i}": stale for i in range(routes.MAX_CONNECTIONS_PER_KEY)})       # 5 phantom slots
    r.zadd(key, {"alive": time.time()})                                                     # 1 real one elsewhere
    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:
        assert ws.receive_json()["type"] == "hello"                                         # phantoms pruned
        members = r.zrange(key, 0, -1)
        assert len(members) == 2 and "alive" in members and not any(m.startswith("dead-") for m in members)
    assert r.zrange(key, 0, -1) == ["alive"]                                                # own slot released
    r.delete(key)


def test_heartbeat_refreshes_connection_score(keys, client, monkeypatch):
    import time

    from stream import broker, routes
    monkeypatch.setattr(routes, "HEARTBEAT_SECONDS", 0.2)
    monkeypatch.setattr(routes, "CONN_HEARTBEAT_SECONDS", 0.0)
    principal = f"key:{_key_id(keys[1])}"
    r = broker.get_redis()
    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:
        ws.receive_json()
        conn_id = r.zrange(routes.conn_key(principal), 0, -1)[0]
        s0 = r.zscore(routes.conn_key(principal), conn_id)
        assert ws.receive_json()["type"] == "heartbeat"
        assert ws.receive_json()["type"] == "heartbeat"
        assert r.zscore(routes.conn_key(principal), conn_id) >= s0 and time.time() - s0 < 5


def test_delivery_is_blocking_read_not_polling(keys, client):
    """Events published after subscribe arrive through XREAD BLOCK; events before a resume token are skipped;
    unsubscribe stops delivery (no reader), resubscribe with a token replays."""
    from stream import broker
    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:
        ws.receive_json()
        ws.send_json({"action": "subscribe", "tickers": ["NVDA"]})
        ws.receive_json()
        s1 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b1"})
        assert ws.receive_json()["accn"] == "b1"
        ws.send_json({"action": "unsubscribe"})
        assert ws.receive_json()["type"] == "unsubscribed"
        s2 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b2"})
        ws.send_json({"action": "ping"})
        assert ws.receive_json()["type"] == "pong"                     # nothing delivered while unsubscribed
        ws.send_json({"action": "subscribe", "tickers": ["NVDA"], "resume_token": s1})
        assert ws.receive_json()["resume_from"] == s1
        assert ws.receive_json()["accn"] == "b2" and s2 > s1
        s3 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b3"})
        got = ws.receive_json()
        assert got["accn"] == "b3" and got["seq"] == s3


def test_accounting_hook_interface(app):
    from stream import accounting
    calls = []
    accounting.set_charger(lambda principal, rows: calls.append((principal, rows)))
    try:
        accounting.charge("key:abc", 1)
        assert calls == [("key:abc", 1)] and accounting.get_charger() is not None
    finally:
        accounting.set_charger(None)


def test_stream_info(client_hu):
    r = client_hu.get("/v1/stream/info")
    assert r.status_code == 200
    d = r.json()["data"]
    assert d["channels"] == ["filings"] and d["close_codes"]["4001"] == "AUTH_REQUIRED" and d["url"].startswith("wss://")