"""Unit tests for handler/billing.py"""

import json
import time
import threading

from gmi_ieops.handler.billing import (
    UsageRecord,
    UsageExtractor,
    BillingTracker,
    get_billing_tracker,
    reset_billing_tracker,
    _parse_usage,
    _safe_int,
)


# ==================== UsageRecord ====================

class TestUsageRecord:

    def test_defaults(self):
        r = UsageRecord()
        assert r.prompt_tokens == 0 and r.total_tokens == 0 and r.request_count == 0

    def test_accumulate(self):
        a = UsageRecord(prompt_tokens=100, completion_tokens=50, total_tokens=150, request_count=1)
        b = UsageRecord(prompt_tokens=200, completion_tokens=80, total_tokens=280, reasoning_tokens=20, request_count=1)
        a.accumulate(b)
        assert a.prompt_tokens == 300
        assert a.completion_tokens == 130
        assert a.total_tokens == 430
        assert a.reasoning_tokens == 20
        assert a.request_count == 2

    def test_to_dict(self):
        r = UsageRecord(prompt_tokens=10, completion_tokens=5, total_tokens=15, request_count=1)
        d = r.to_dict("org1", "m1", 1234567890)
        assert d == {
            "org_id": "org1", "model_name": "m1",
            "prompt_tokens": 10, "completion_tokens": 5,
            "total_tokens": 15, "reasoning_tokens": 0,
            "request_count": 1, "timestamp": 1234567890,
        }


# ==================== _parse_usage ====================

class TestParseUsage:

    def test_normal(self):
        r = _parse_usage({"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150})
        assert r.total_tokens == 150 and r.request_count == 1

    def test_calculates_total(self):
        r = _parse_usage({"prompt_tokens": 100, "completion_tokens": 50})
        assert r.total_tokens == 150

    def test_reasoning_tokens(self):
        r = _parse_usage({"prompt_tokens": 1, "total_tokens": 1, "reasoning_tokens": 99})
        assert r.reasoning_tokens == 99

    def test_zero_returns_none(self):
        assert _parse_usage({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}) is None

    def test_empty_dict_returns_none(self):
        assert _parse_usage({}) is None

    def test_not_dict_returns_none(self):
        assert _parse_usage("string") is None
        assert _parse_usage(None) is None
        assert _parse_usage(42) is None
        assert _parse_usage([]) is None


# ==================== _safe_int ====================

class TestSafeInt:

    def test_normal_int(self):
        assert _safe_int(100) == 100

    def test_zero(self):
        assert _safe_int(0) == 0

    def test_negative_clamped(self):
        assert _safe_int(-5) == 0
        assert _safe_int(-999) == 0

    def test_string_number(self):
        assert _safe_int("100") == 100
        assert _safe_int("0") == 0

    def test_float_truncated(self):
        assert _safe_int(3.14) == 3
        assert _safe_int(99.9) == 99

    def test_invalid_returns_zero(self):
        assert _safe_int(None) == 0
        assert _safe_int("abc") == 0
        assert _safe_int([]) == 0
        assert _safe_int({}) == 0
        assert _safe_int(float("inf")) == 0
        assert _safe_int(float("nan")) == 0


# ==================== _parse_usage (edge cases) ====================

class TestParseUsageEdgeCases:

    def test_negative_tokens_clamped(self):
        r = _parse_usage({"prompt_tokens": -10, "completion_tokens": 50, "total_tokens": 40})
        assert r.prompt_tokens == 0
        assert r.completion_tokens == 50
        assert r.total_tokens == 40

    def test_all_negative_returns_none(self):
        assert _parse_usage({"prompt_tokens": -10, "completion_tokens": -5, "total_tokens": -15}) is None

    def test_string_token_values(self):
        r = _parse_usage({"prompt_tokens": "100", "completion_tokens": "50", "total_tokens": "150"})
        assert r.prompt_tokens == 100 and r.total_tokens == 150

    def test_float_token_values(self):
        r = _parse_usage({"prompt_tokens": 100.7, "completion_tokens": 50.3, "total_tokens": 151.0})
        assert r.prompt_tokens == 100 and r.completion_tokens == 50 and r.total_tokens == 151

    def test_invalid_token_values_become_zero(self):
        r = _parse_usage({"prompt_tokens": "abc", "completion_tokens": None, "total_tokens": 100})
        assert r.prompt_tokens == 0 and r.completion_tokens == 0 and r.total_tokens == 100

    def test_only_invalid_values_returns_none(self):
        assert _parse_usage({"prompt_tokens": "abc", "completion_tokens": None}) is None


# ==================== UsageExtractor ====================

class TestUsageExtractor:

    # --- extract_from_json ---

    def test_json_openai_format(self):
        body = json.dumps({
            "id": "chatcmpl-123",
            "choices": [{"message": {"content": "Hi"}}],
            "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}
        }).encode()
        r = UsageExtractor.extract_from_json(body)
        assert r.prompt_tokens == 100 and r.total_tokens == 150 and r.request_count == 1

    def test_json_no_usage(self):
        assert UsageExtractor.extract_from_json(json.dumps({"choices": []}).encode()) is None

    def test_json_empty_usage(self):
        assert UsageExtractor.extract_from_json(json.dumps({"usage": {}}).encode()) is None

    def test_json_usage_not_dict(self):
        assert UsageExtractor.extract_from_json(json.dumps({"usage": "bad"}).encode()) is None

    def test_json_invalid_bytes(self):
        assert UsageExtractor.extract_from_json(b"not json") is None
        assert UsageExtractor.extract_from_json(b"") is None
        assert UsageExtractor.extract_from_json(b"null") is None

    # --- extract_from_sse_line ---

    def test_sse_with_usage(self):
        data = {"usage": {"prompt_tokens": 50, "completion_tokens": 20, "total_tokens": 70}}
        r = UsageExtractor.extract_from_sse_line(f"data: {json.dumps(data)}")
        assert r.prompt_tokens == 50 and r.total_tokens == 70

    def test_sse_done(self):
        assert UsageExtractor.extract_from_sse_line("data: [DONE]") is None

    def test_sse_no_usage_chunk(self):
        line = 'data: {"choices": [{"delta": {"content": "token"}}]}'
        assert UsageExtractor.extract_from_sse_line(line) is None

    def test_sse_non_data_lines(self):
        assert UsageExtractor.extract_from_sse_line("event: message") is None
        assert UsageExtractor.extract_from_sse_line(": keep-alive") is None
        assert UsageExtractor.extract_from_sse_line("") is None

    def test_sse_invalid_json(self):
        assert UsageExtractor.extract_from_sse_line("data: {bad}") is None

    # --- edge cases ---

    def test_json_negative_tokens(self):
        body = json.dumps({"usage": {"prompt_tokens": -10, "completion_tokens": 50, "total_tokens": 40}}).encode()
        r = UsageExtractor.extract_from_json(body)
        assert r.prompt_tokens == 0 and r.total_tokens == 40

    def test_json_string_token_values(self):
        body = json.dumps({"usage": {"prompt_tokens": "100", "completion_tokens": "50", "total_tokens": "150"}}).encode()
        r = UsageExtractor.extract_from_json(body)
        assert r.prompt_tokens == 100 and r.total_tokens == 150

    def test_json_very_large_response_body(self):
        """Large body with valid usage should still extract correctly."""
        payload = {"data": "x" * 100_000, "usage": {"prompt_tokens": 5, "total_tokens": 5}}
        r = UsageExtractor.extract_from_json(json.dumps(payload).encode())
        assert r.total_tokens == 5

    def test_sse_negative_tokens(self):
        data = {"usage": {"prompt_tokens": -1, "completion_tokens": 10, "total_tokens": 10}}
        r = UsageExtractor.extract_from_sse_line(f"data: {json.dumps(data)}")
        assert r.prompt_tokens == 0 and r.total_tokens == 10


# ==================== BillingTracker ====================

class TestBillingTracker:

    def test_record_and_flush(self):
        t = BillingTracker()
        t.record("org1", "llama", UsageRecord(100, 50, 150, 0, 1))
        t.record("org1", "llama", UsageRecord(200, 100, 300, 0, 1))
        data = t.flush()
        assert len(data) == 1
        d = data[0]
        assert d["org_id"] == "org1" and d["model_name"] == "llama"
        assert d["prompt_tokens"] == 300 and d["total_tokens"] == 450 and d["request_count"] == 2

    def test_flush_clears(self):
        t = BillingTracker()
        t.record("o", "m", UsageRecord(10, 5, 15, 0, 1))
        assert len(t.flush()) == 1
        assert len(t.flush()) == 0

    def test_multiple_dimensions(self):
        t = BillingTracker()
        t.record("o1", "m1", UsageRecord(10, 0, 10, 0, 1))
        t.record("o2", "m1", UsageRecord(20, 0, 20, 0, 1))
        t.record("o1", "m2", UsageRecord(30, 0, 30, 0, 1))
        data = t.flush()
        assert len(data) == 3
        by_key = {(d["org_id"], d["model_name"]): d for d in data}
        assert by_key[("o1", "m1")]["prompt_tokens"] == 10
        assert by_key[("o2", "m1")]["prompt_tokens"] == 20
        assert by_key[("o1", "m2")]["prompt_tokens"] == 30

    def test_skips_zero_and_none(self):
        t = BillingTracker()
        t.record("o", "m", UsageRecord())       # zero
        t.record("o", "m", None)                 # None
        assert t.pending_dimensions == 0

    def test_empty_ids_become_unknown(self):
        t = BillingTracker()
        t.record("", "m", UsageRecord(1, 0, 1, 0, 1))
        t.record("o", "", UsageRecord(1, 0, 1, 0, 1))
        by_key = {(d["org_id"], d["model_name"]): d for d in t.flush()}
        assert ("unknown", "m") in by_key and ("o", "unknown") in by_key

    def test_counters(self):
        t = BillingTracker()
        assert t.total_records == 0 and t.total_flushes == 0 and t.pending_dimensions == 0
        t.record("o", "m", UsageRecord(1, 0, 1, 0, 1))
        t.record("o", "m", UsageRecord(1, 0, 1, 0, 1))
        assert t.total_records == 2 and t.pending_dimensions == 1
        t.flush()
        assert t.total_flushes == 1 and t.pending_dimensions == 0

    def test_timestamp(self):
        t = BillingTracker()
        t.record("o", "m", UsageRecord(1, 0, 1, 0, 1))
        before = int(time.time())
        ts = t.flush()[0]["timestamp"]
        assert before <= ts <= int(time.time())

    def test_incremental_across_heartbeats(self):
        """Simulates real heartbeat: record → flush → record → flush → idle flush"""
        t = BillingTracker()
        for _ in range(3):
            t.record("o", "m", UsageRecord(100, 0, 100, 0, 1))
        d1 = t.flush()
        assert d1[0]["total_tokens"] == 300 and d1[0]["request_count"] == 3

        for _ in range(2):
            t.record("o", "m", UsageRecord(200, 0, 200, 0, 1))
        d2 = t.flush()
        assert d2[0]["total_tokens"] == 400 and d2[0]["request_count"] == 2

        assert t.flush() == []

    def test_thread_safety(self):
        t = BillingTracker()
        errors = []

        def writer(org):
            try:
                for _ in range(500):
                    t.record(org, "m", UsageRecord(10, 5, 15, 0, 1))
            except Exception as e:
                errors.append(e)

        def flusher():
            total = 0
            for _ in range(50):
                for d in t.flush():
                    total += d["request_count"]
                time.sleep(0.001)
            return total

        writers = [threading.Thread(target=writer, args=(f"o{i}",)) for i in range(4)]
        result = [0]
        ft = threading.Thread(target=lambda: result.__setitem__(0, flusher()))
        for w in writers: w.start()
        ft.start()
        for w in writers: w.join()
        ft.join()

        remaining = sum(d["request_count"] for d in t.flush())
        assert not errors
        assert result[0] + remaining == 2000  # 4 writers × 500


# ==================== Global Singleton ====================

class TestGlobalSingleton:

    def setup_method(self):
        reset_billing_tracker()

    def test_singleton(self):
        assert get_billing_tracker() is get_billing_tracker()

    def test_reset(self):
        t1 = get_billing_tracker()
        t1.record("o", "m", UsageRecord(1, 0, 1, 0, 1))
        reset_billing_tracker()
        t2 = get_billing_tracker()
        assert t1 is not t2 and t2.pending_dimensions == 0


# ==================== Interceptor Integration ====================

class TestInterceptorBillingIntegration:
    """Tests billing extraction within the Interceptor proxy flow."""

    def _run(self, coro):
        return asyncio.get_event_loop().run_until_complete(coro)

    def setup_method(self):
        reset_billing_tracker()

    def test_non_streaming_extracts_usage(self):
        """Interceptor records usage from a non-streaming JSON response."""
        from gmi_ieops.handler.billing import get_billing_tracker, UsageExtractor

        body = json.dumps({
            "choices": [{"message": {"content": "Hello"}}],
            "usage": {"prompt_tokens": 120, "completion_tokens": 40, "total_tokens": 160}
        }).encode()

        # Simulate what interceptor does
        usage = UsageExtractor.extract_from_json(body)
        assert usage is not None
        get_billing_tracker().record(org_id="org-abc", model="llama-70b", usage=usage)

        data = get_billing_tracker().flush()
        assert len(data) == 1
        assert data[0]["org_id"] == "org-abc"
        assert data[0]["model_name"] == "llama-70b"
        assert data[0]["prompt_tokens"] == 120
        assert data[0]["total_tokens"] == 160

    def test_streaming_extracts_usage_from_last_chunk(self):
        """Interceptor picks up usage from the final SSE chunk."""
        from gmi_ieops.handler.billing import get_billing_tracker, UsageExtractor

        sse_lines = [
            'data: {"choices": [{"delta": {"content": "Hello"}}]}',
            'data: {"choices": [{"delta": {"content": " world"}}]}',
            'data: {"choices": [], "usage": {"prompt_tokens": 80, "completion_tokens": 25, "total_tokens": 105}}',
            'data: [DONE]',
        ]

        # Simulate interceptor SSE loop
        usage_record = None
        for line in sse_lines:
            extracted = UsageExtractor.extract_from_sse_line(line)
            if extracted:
                usage_record = extracted

        assert usage_record is not None
        get_billing_tracker().record(org_id="org-xyz", model="gpt-4o", usage=usage_record)

        data = get_billing_tracker().flush()
        assert len(data) == 1
        assert data[0]["prompt_tokens"] == 80
        assert data[0]["completion_tokens"] == 25

    def test_no_usage_in_response_records_nothing(self):
        """Interceptor silently skips when response has no usage field."""
        from gmi_ieops.handler.billing import get_billing_tracker, UsageExtractor

        body = json.dumps({"result": "ok"}).encode()
        usage = UsageExtractor.extract_from_json(body)
        assert usage is None
        # record with None is a no-op
        get_billing_tracker().record("org", "m", usage)
        assert get_billing_tracker().flush() == []

    def test_billing_disabled_skips_extraction(self):
        """When billing is disabled, interceptor should not extract."""
        from gmi_ieops.handler.billing import get_billing_tracker

        # Simulate the interceptor check: billing_enabled = env.billing.enabled and (org_id or model)
        billing_enabled = False  # simulating BILLING_ENABLED=false
        body = json.dumps({
            "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}
        }).encode()

        if billing_enabled:
            usage = UsageExtractor.extract_from_json(body)
            get_billing_tracker().record("org", "m", usage)

        assert get_billing_tracker().flush() == []

    def test_missing_org_id_uses_unknown(self):
        """When X-Org-Id header is missing, org_id falls back to 'unknown'."""
        from gmi_ieops.handler.billing import get_billing_tracker

        get_billing_tracker().record(
            org_id="",  # empty = header missing
            model="llama",
            usage=UsageRecord(50, 20, 70, 0, 1),
        )
        data = get_billing_tracker().flush()
        assert data[0]["org_id"] == "unknown"

    def test_negative_tokens_clamped_in_extraction(self):
        """Negative token values from buggy workers are clamped to 0."""
        body = json.dumps({
            "usage": {"prompt_tokens": -100, "completion_tokens": 50, "total_tokens": 50}
        }).encode()
        usage = UsageExtractor.extract_from_json(body)
        assert usage is not None
        assert usage.prompt_tokens == 0  # clamped
        get_billing_tracker().record("org", "m", usage)
        data = get_billing_tracker().flush()
        assert data[0]["prompt_tokens"] == 0

    def test_string_tokens_coerced_in_extraction(self):
        """String token values (from non-standard backends) are coerced to int."""
        body = json.dumps({
            "usage": {"prompt_tokens": "200", "completion_tokens": "80", "total_tokens": "280"}
        }).encode()
        usage = UsageExtractor.extract_from_json(body)
        assert usage.total_tokens == 280
        get_billing_tracker().record("org", "m", usage)
        data = get_billing_tracker().flush()
        assert data[0]["total_tokens"] == 280


# ==================== Register Integration ====================

class TestRegisterBillingIntegration:
    """Tests that heartbeat payload includes usage_report correctly."""

    def setup_method(self):
        reset_billing_tracker()

    def test_flush_returns_empty_when_no_requests(self):
        """Heartbeat should not include usage_report when empty."""
        data = get_billing_tracker().flush()
        assert data == []

    def test_flush_returns_data_between_heartbeats(self):
        """Simulate two heartbeat cycles with requests in between."""
        tracker = get_billing_tracker()

        # Period 1: some requests
        tracker.record("org1", "m1", UsageRecord(100, 50, 150, 0, 1))
        tracker.record("org1", "m1", UsageRecord(200, 80, 280, 0, 1))
        report1 = tracker.flush()
        assert len(report1) == 1
        assert report1[0]["request_count"] == 2

        # Period 2: more requests, different model
        tracker.record("org1", "m1", UsageRecord(50, 20, 70, 0, 1))
        tracker.record("org1", "m2", UsageRecord(300, 100, 400, 10, 1))
        report2 = tracker.flush()
        assert len(report2) == 2
        by_model = {d["model_name"]: d for d in report2}
        assert by_model["m1"]["total_tokens"] == 70
        assert by_model["m2"]["reasoning_tokens"] == 10

        # Period 3: idle
        assert tracker.flush() == []
