from __future__ import annotations

from app.services.table_profiling import _parse_completion_payload
from app.utils.llm_usage import extract_usage


def test_parse_completion_payload_handles_array_with_trailing_text() -> None:
    response_payload = {
        "choices": [
            {
                "message": {
                    "content": """
结果如下：
[
  {"id": "snpt_a"},
  {"id": "snpt_b"}
]
附加说明：模型可能会输出额外文本。
""".strip()
                }
            }
        ]
    }

    parsed = _parse_completion_payload(response_payload)

    assert isinstance(parsed, list)
    assert [item["id"] for item in parsed] == ["snpt_a", "snpt_b"]


def test_extract_usage_info_normalizes_numeric_fields() -> None:
    response_payload = {
        "raw": {
            "usage": {
                "prompt_tokens": 12.7,
                "completion_tokens": 3,
                "total_tokens": 15.7,
                "prompt_tokens_details": {"cached_tokens": 8.9, "other": None},
                "non_numeric": "ignored",
            }
        }
    }

    usage = extract_usage(response_payload)

    assert usage == {
        "prompt_tokens": 12,
        "completion_tokens": 3,
        "total_tokens": 15,
        "prompt_tokens_details": {"cached_tokens": 8},
    }


def test_extract_usage_handles_alias_keys() -> None:
    response_payload = {
        "raw": {
            "usageMetadata": {
                "input_tokens": 20,
                "output_tokens": 4,
            }
        }
    }

    usage = extract_usage(response_payload)

    assert usage == {
        "prompt_tokens": 20,
        "completion_tokens": 4,
        "total_tokens": 24,
    }


def test_extract_usage_returns_none_when_missing() -> None:
    assert extract_usage({"raw": {}}) is None