"""Tests for the inference pipeline (predict.py).

These tests run without a trained checkpoint by using a randomly-initialized
model, which is sufficient to verify the pipeline structure and output schema.
"""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import pytest
import torch
from unittest.mock import patch, MagicMock

from finner.labels import ENTITY_TYPES, LABEL2ID


class TestPredictOutputSchema:
    """Verify predict() returns well-formed output with correct keys and types."""

    @pytest.fixture(autouse=True)
    def mock_model(self, tmp_path):
        """Patch _load_model to return a tiny random model instead of a real checkpoint."""
        from transformers import AutoModelForTokenClassification, AutoTokenizer
        from finner.labels import NUM_LABELS, ID2LABEL, LABEL2ID

        model_name = "bert-base-uncased"   # already cached from training
        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
        model = AutoModelForTokenClassification.from_pretrained(
            model_name,
            num_labels=NUM_LABELS,
            id2label=ID2LABEL,
            label2id=LABEL2ID,
            ignore_mismatched_sizes=True,
        )
        model.eval()

        import finner.infer.predict as predict_module
        predict_module._model = model
        predict_module._tokenizer = tokenizer
        yield
        predict_module._model = None
        predict_module._tokenizer = None

    def test_output_keys_present(self):
        from finner.infer.predict import predict
        result = predict("Apple reported revenue of $1.2B.")
        assert "tokens" in result
        assert "token_labels" in result
        assert "token_confidences" in result
        assert "entities" in result

    def test_tokens_are_strings(self):
        from finner.infer.predict import predict
        result = predict("Goldman Sachs paid $500M to the SEC.")
        assert all(isinstance(t, str) for t in result["tokens"])

    def test_label_count_matches_token_count(self):
        from finner.infer.predict import predict
        result = predict("Revenue grew 12.5% year over year.")
        assert len(result["token_labels"]) == len(result["tokens"])
        assert len(result["token_confidences"]) == len(result["tokens"])

    def test_all_labels_are_valid(self):
        from finner.infer.predict import predict
        result = predict("The Fed raised rates by 25 bps.")
        valid = set(LABEL2ID.keys())
        for lbl in result["token_labels"]:
            assert lbl in valid, f"Invalid label: {lbl}"

    def test_confidences_in_unit_interval(self):
        from finner.infer.predict import predict
        result = predict("Microsoft acquired Activision for $68.7B.")
        for conf in result["token_confidences"]:
            assert 0.0 <= conf <= 1.0

    def test_entity_schema(self):
        from finner.infer.predict import predict
        result = predict("Apple Inc. reported $1.2B EBITDA.")
        for ent in result["entities"]:
            assert "text" in ent
            assert "label" in ent
            assert "start_token" in ent
            assert "end_token" in ent
            assert "confidence" in ent
            assert ent["label"] in ENTITY_TYPES
            assert ent["start_token"] <= ent["end_token"]
            assert 0.0 <= ent["confidence"] <= 1.0

    def test_entity_span_text_matches_tokens(self):
        from finner.infer.predict import _clean_span, predict
        result = predict("Goldman Sachs earned $500M in Q3 2024.")
        tokens = result["tokens"]
        for ent in result["entities"]:
            raw = " ".join(tokens[ent["start_token"]: ent["end_token"] + 1])
            # entity text is the cleaned (trailing-punct + unbalanced-paren stripped) span
            assert _clean_span(raw) == ent["text"]

    def test_empty_text_edge_case(self):
        from finner.infer.predict import predict
        result = predict(".")
        assert isinstance(result["entities"], list)