import json
import logging
import tempfile
from pathlib import Path
from textwrap import dedent
from typing import Generator, get_args
from uuid import uuid4

import pytest

from ._utils.auth import _ORCA_ROOT_ACCESS_API_KEY
from .classification_model import ClassificationMetrics, ClassificationModel
from .client import OrcaClient
from .datasource import Datasource
from .embedding_model import (
    FinetunedEmbeddingModel,
    PretrainedEmbeddingModel,
    PretrainedEmbeddingModelName,
)
from .job import Status
from .memoryset import LabeledMemoryset


def test_open_pretrained_model():
    model = PretrainedEmbeddingModel.GTE_BASE
    assert model is not None
    assert isinstance(model, PretrainedEmbeddingModel)
    assert model.name == "GTE_BASE"
    assert model.embedding_dim == 768
    assert model.max_seq_length == 8192
    assert model is PretrainedEmbeddingModel.GTE_BASE


def test_open_pretrained_model_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline")


def test_open_pretrained_model_not_found():
    with pytest.raises(LookupError):
        PretrainedEmbeddingModel._get("INVALID_MODEL")  # type: ignore


def test_all_pretrained_models():
    models = PretrainedEmbeddingModel.all()
    assert len(models) > 1
    if len(models) != len(get_args(PretrainedEmbeddingModelName)):
        logging.warning("Please regenerate the SDK client! Some pretrained model names are not exposed yet.")
    model_names = [m.name for m in models]
    assert all(m in model_names for m in get_args(PretrainedEmbeddingModelName))


def test_embed_text():
    embedding = PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline", max_seq_length=32)
    assert embedding is not None
    assert isinstance(embedding, list)
    assert len(embedding) == 768
    assert isinstance(embedding[0], float)


def test_embed_text_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline", max_seq_length=32)


def test_evaluate_pretrained_model(datasource: Datasource):
    metrics = PretrainedEmbeddingModel.GTE_BASE.evaluate(datasource=datasource, label_column="label")
    assert metrics is not None
    assert isinstance(metrics, ClassificationMetrics)
    assert metrics.accuracy > 0.5


def test_finetune_model_with_datasource(finetuned_embedding_model: FinetunedEmbeddingModel):
    assert finetuned_embedding_model is not None
    assert finetuned_embedding_model.name == "test_finetuned_model"
    assert finetuned_embedding_model.base_model == PretrainedEmbeddingModel.DISTILBERT
    assert finetuned_embedding_model.embedding_dim == 768
    assert finetuned_embedding_model.max_seq_length == 512
    assert finetuned_embedding_model._status == Status.COMPLETED
    # num_params should come from the base model
    assert finetuned_embedding_model.num_params == PretrainedEmbeddingModel.DISTILBERT.num_params
    assert finetuned_embedding_model.num_params == 66362880


def test_finetune_model_with_memoryset(readonly_memoryset: LabeledMemoryset):
    finetuned_embedding_model = PretrainedEmbeddingModel.DISTILBERT.finetune(
        "test_finetuned_model_from_memoryset", readonly_memoryset
    )
    assert finetuned_embedding_model is not None
    assert finetuned_embedding_model.name == "test_finetuned_model_from_memoryset"
    assert finetuned_embedding_model.base_model == PretrainedEmbeddingModel.DISTILBERT
    assert finetuned_embedding_model.embedding_dim == 768
    assert finetuned_embedding_model.max_seq_length == 512
    assert finetuned_embedding_model._status == Status.COMPLETED
    # num_params should come from the base model
    assert finetuned_embedding_model.num_params == PretrainedEmbeddingModel.DISTILBERT.num_params
    assert finetuned_embedding_model.num_params == 66362880


def test_finetune_model_already_exists_error(datasource: Datasource, finetuned_embedding_model):
    with pytest.raises(ValueError):
        PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource)


def test_finetune_model_already_exists_return(datasource: Datasource, finetuned_embedding_model):
    with pytest.raises(ValueError):
        PretrainedEmbeddingModel.GTE_BASE.finetune("test_finetuned_model", datasource, if_exists="open")

    new_model = PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource, if_exists="open")
    assert new_model is not None
    assert new_model.name == "test_finetuned_model"
    assert new_model.base_model == PretrainedEmbeddingModel.DISTILBERT
    assert new_model.embedding_dim == 768
    assert new_model.max_seq_length == 512
    assert new_model._status == Status.COMPLETED
    # num_params should match the base model
    assert new_model.num_params == PretrainedEmbeddingModel.DISTILBERT.num_params


def test_finetune_model_unauthenticated(unauthenticated_client, datasource: Datasource):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model_unauthenticated", datasource)


def test_use_finetuned_model_in_memoryset(
    datasource: Datasource, label_names: list[str], finetuned_embedding_model: FinetunedEmbeddingModel
):
    memoryset = LabeledMemoryset.create(
        "test_memoryset_finetuned_model",
        datasource=datasource,
        embedding_model=finetuned_embedding_model,
        label_names=label_names,
    )
    assert memoryset is not None
    assert memoryset.name == "test_memoryset_finetuned_model"
    assert memoryset.embedding_model == finetuned_embedding_model
    assert memoryset.length == datasource.length


def test_open_finetuned_model(finetuned_embedding_model: FinetunedEmbeddingModel):
    model = FinetunedEmbeddingModel.open(finetuned_embedding_model.name)
    assert isinstance(model, FinetunedEmbeddingModel)
    assert model.id == finetuned_embedding_model.id
    assert model.name == finetuned_embedding_model.name
    assert model.base_model == PretrainedEmbeddingModel.DISTILBERT
    assert model.embedding_dim == 768
    assert model.max_seq_length == 512
    assert model.num_params == finetuned_embedding_model.num_params
    assert model == finetuned_embedding_model


def test_embed_finetuned_model(finetuned_embedding_model: FinetunedEmbeddingModel):
    embedding = finetuned_embedding_model.embed("I love this airline")
    assert embedding is not None
    assert isinstance(embedding, list)
    assert len(embedding) == 768
    assert isinstance(embedding[0], float)


def test_all_finetuned_models(finetuned_embedding_model: FinetunedEmbeddingModel):
    models = FinetunedEmbeddingModel.all()
    assert len(models) > 0
    assert any(model.name == finetuned_embedding_model.name for model in models)


def test_all_finetuned_models_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            FinetunedEmbeddingModel.all()


def test_all_finetuned_models_unauthorized(unauthorized_client, finetuned_embedding_model: FinetunedEmbeddingModel):
    with unauthorized_client.use():
        assert finetuned_embedding_model not in FinetunedEmbeddingModel.all()


def test_drop_finetuned_model(datasource: Datasource):
    PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)
    assert FinetunedEmbeddingModel.open("finetuned_model_to_delete")
    FinetunedEmbeddingModel.drop("finetuned_model_to_delete")
    with pytest.raises(LookupError):
        FinetunedEmbeddingModel.open("finetuned_model_to_delete")


def test_drop_finetuned_model_with_memoryset_cascade(datasource: Datasource, label_names: list[str]):
    finetuned_embedding_model = PretrainedEmbeddingModel.DISTILBERT.finetune(
        "finetuned_model_cascade_delete", datasource
    )
    memoryset = LabeledMemoryset.create(
        "test_memoryset_for_finetuned_model_cascade",
        datasource=datasource,
        embedding_model=finetuned_embedding_model,
        label_names=label_names,
    )

    # Verify memoryset exists and uses the finetuned model
    assert LabeledMemoryset.open(memoryset.name) is not None
    assert memoryset.embedding_model == finetuned_embedding_model

    # Without cascade, deletion should fail
    with pytest.raises(RuntimeError):
        FinetunedEmbeddingModel.drop(finetuned_embedding_model.id, cascade=False)

    # Model and memoryset should still exist
    assert FinetunedEmbeddingModel.exists(finetuned_embedding_model.name)
    assert LabeledMemoryset.exists(memoryset.name)

    # With cascade, deletion should succeed
    FinetunedEmbeddingModel.drop(finetuned_embedding_model.id, cascade=True)

    # Both model and memoryset should be deleted
    assert not FinetunedEmbeddingModel.exists(finetuned_embedding_model.name)
    assert not LabeledMemoryset.exists(memoryset.name)


def test_drop_finetuned_model_unauthenticated(unauthenticated_client, datasource: Datasource):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)


def test_drop_finetuned_model_not_found():
    with pytest.raises(LookupError):
        FinetunedEmbeddingModel.drop(str(uuid4()))
    # ignores error if specified
    FinetunedEmbeddingModel.drop(str(uuid4()), if_not_exists="ignore")


def test_drop_finetuned_model_unauthorized(unauthorized_client, finetuned_embedding_model: FinetunedEmbeddingModel):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            FinetunedEmbeddingModel.drop(finetuned_embedding_model.id)


def test_supports_instructions():
    model = PretrainedEmbeddingModel.GTE_BASE
    assert not model.supports_instructions

    instruction_model = PretrainedEmbeddingModel.BGE_BASE
    assert instruction_model.supports_instructions


def test_use_explicit_instruction_prompt():
    model = PretrainedEmbeddingModel.BGE_BASE
    assert model.supports_instructions
    input = "Hello world"
    assert model.embed(input, prompt="Represent this sentence for sentiment retrieval:") != model.embed(input)


def test_pretrained_embedding_model_repr():
    assert repr(PretrainedEmbeddingModel.GTE_BASE) == (
        "PretrainedEmbeddingModel({ name: GTE_BASE, embedding_dim: 768, max_seq_length: 8192, num_params: 137M })"
    )


def test_finetuned_embedding_model_repr(finetuned_embedding_model: FinetunedEmbeddingModel):
    assert repr(finetuned_embedding_model) == dedent(
        """\
        FinetunedEmbeddingModel({
            name: test_finetuned_model,
            embedding_dim: 768,
            max_seq_length: 512,
            base_model: PretrainedEmbeddingModel.DISTILBERT
        })"""
    )


@pytest.fixture(scope="module")
def root_client(org_id) -> Generator[OrcaClient, None, None]:
    """Client with root API key access for testing privileged endpoints."""
    client = OrcaClient(api_key=_ORCA_ROOT_ACCESS_API_KEY, headers={"Org-Id": org_id})
    yield client


@pytest.fixture(scope="module")
def mock_model_dir() -> Generator[Path, None, None]:
    """
    Create a minimal mock HuggingFace model directory for upload testing.

    The config.json contains the minimum fields required for EmbeddingModel.load_config
    to extract embedding dimension and max sequence length.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        model_dir = Path(temp_dir)

        # Create a minimal HuggingFace-style config.json
        config = {
            "model_type": "bert",
            "architectures": ["BertModel"],
            "hidden_size": 256,  # This becomes embedding_dim
            "max_position_embeddings": 128,  # This becomes max_seq_length
            "vocab_size": 30522,
            "num_hidden_layers": 2,
            "num_attention_heads": 4,
            "intermediate_size": 512,
        }
        config_path = model_dir / "config.json"
        config_path.write_text(json.dumps(config, indent=2))

        yield model_dir


def test_upload_embedding_model(root_client: OrcaClient, mock_model_dir: Path):
    """Test uploading a pretrained embedding model from a local directory."""
    model_name = f"test_uploaded_model_{uuid4().hex[:8]}"
    description = "Test uploaded embedding model"

    with root_client.use():
        try:
            # Upload the model
            model = FinetunedEmbeddingModel._upload(
                model_name,
                mock_model_dir,
                description=description,
            )

            # Verify the model was created with correct metadata
            assert model.name == model_name
            assert model.embedding_dim == 256  # from config hidden_size
            assert model.max_seq_length == 128  # from config max_position_embeddings
            assert model.base_model is None  # Uploaded models have no base model
            assert model.description == description
            # num_params won't be abe able to be extracted from the model
            assert model.num_params is None or model.num_params == 0

            # Verify we can open it by name
            opened_model = FinetunedEmbeddingModel.open(model_name)
            assert opened_model.id == model.id
            assert opened_model.base_model is None
            assert opened_model.description == description

            # Verify it appears in the list
            all_models = FinetunedEmbeddingModel.all()
            assert any(m.id == model.id for m in all_models)

        finally:
            # Clean up
            FinetunedEmbeddingModel.drop(model_name, if_not_exists="ignore")


def test_upload_embedding_model_already_exists_error(root_client: OrcaClient, mock_model_dir: Path):
    """Test that uploading with an existing name raises an error by default."""
    model_name = f"test_uploaded_model_exists_{uuid4().hex[:8]}"

    with root_client.use():
        try:
            # First upload
            FinetunedEmbeddingModel._upload(model_name, mock_model_dir, description="First description")

            # Second upload should fail
            with pytest.raises(ValueError, match="already exists"):
                FinetunedEmbeddingModel._upload(model_name, mock_model_dir, description="Second description")

        finally:
            FinetunedEmbeddingModel.drop(model_name, if_not_exists="ignore")


def test_upload_embedding_model_already_exists_open(root_client: OrcaClient, mock_model_dir: Path):
    """Test that uploading with if_exists='open' returns the existing model."""
    model_name = f"test_uploaded_model_open_{uuid4().hex[:8]}"
    description = "Original description"

    with root_client.use():
        try:
            # First upload
            original = FinetunedEmbeddingModel._upload(model_name, mock_model_dir, description=description)

            # Second upload with if_exists="open" should return the existing model
            existing = FinetunedEmbeddingModel._upload(model_name, mock_model_dir, if_exists="open")
            assert existing.id == original.id
            assert existing.description == description

        finally:
            FinetunedEmbeddingModel.drop(model_name, if_not_exists="ignore")


def test_upload_embedding_model_without_description(root_client: OrcaClient, mock_model_dir: Path):
    """Test that uploading without a description works correctly."""
    model_name = f"test_uploaded_model_no_desc_{uuid4().hex[:8]}"

    with root_client.use():
        try:
            # Upload the model without description
            model = FinetunedEmbeddingModel._upload(model_name, mock_model_dir)

            # Verify description is None
            assert model.description is None

            # Verify we can open it and description is still None
            opened_model = FinetunedEmbeddingModel.open(model_name)
            assert opened_model.description is None
            # num_params should be consistent
            assert opened_model.num_params == model.num_params

        finally:
            FinetunedEmbeddingModel.drop(model_name, if_not_exists="ignore")


def test_upload_embedding_model_invalid_path(root_client: OrcaClient):
    """Test that uploading from a non-directory path raises an error."""
    with root_client.use():
        with pytest.raises(ValueError, match="must be a directory"):
            FinetunedEmbeddingModel._upload("test_model", "/nonexistent/path")


def test_upload_embedding_model_unauthorized(mock_model_dir: Path):
    """Test that uploading without root access raises an error."""
    # The regular api_key fixture doesn't have root access
    with pytest.raises(PermissionError):
        FinetunedEmbeddingModel._upload("test_model_unauthorized", mock_model_dir)


def test_upload_gte_small_and_create_classification_model(root_client: OrcaClient):
    """
    Integration test: upload GTE-Small, create memoryset and classification model, verify predictions.

    This test mirrors the functionality in research/debug/upload_embedding_model.ipynb:
    1. Download and upload a real HuggingFace embedding model (GTE-Small)
    2. Verify embeddings work
    3. Create a labeled memoryset with the uploaded model
    4. Create a classification model
    5. Make predictions and verify they work correctly
    """
    hf_hub = pytest.importorskip("huggingface_hub")

    unique_suffix = uuid4().hex[:8]
    model_name = f"test_uploaded_gte_small_{unique_suffix}"
    memoryset_name = f"test_uploaded_gte_memoryset_{unique_suffix}"
    classifier_name = f"test_uploaded_gte_classifier_{unique_suffix}"

    with tempfile.TemporaryDirectory() as temp_dir:
        # Download GTE-Small from HuggingFace (only safetensors, not ONNX)
        model_dir = Path(temp_dir)
        hf_hub.snapshot_download(
            repo_id="thenlper/gte-small",
            local_dir=model_dir,
            allow_patterns=[
                "*.safetensors",
                "*.json",
                "*.txt",
                "*.py",
                "*.model",
            ],  # Include safetensors and config/tokenizer files
            ignore_patterns=["*.onnx"],  # Exclude ONNX files
        )

        with root_client.use():
            try:
                # Upload the model
                model = FinetunedEmbeddingModel._upload(
                    model_name,
                    model_dir,
                    description="GTE-Small model for integration testing",
                )

                # Verify model properties
                assert model.name == model_name
                assert model.embedding_dim == 384  # GTE-Small has 384 dimensions
                assert model.max_seq_length == 512
                assert model.base_model is None  # Uploaded models have no base model
                assert model.description == "GTE-Small model for integration testing"
                # NOTE: currently num_params cannot be reliably determined for uploaded models
                assert model.num_params == 33360512

                # Test embedding functionality
                embedding = model.embed("This is a test sentence for embedding.")
                assert isinstance(embedding, list)
                assert len(embedding) == 384
                assert isinstance(embedding[0], float)

                # Create labeled memoryset with sample sentiment data
                sample_data = [
                    {"value": "I love this product, it's amazing!", "label": 0},
                    {"value": "Great experience, highly recommend", "label": 0},
                    {"value": "Best purchase I've ever made", "label": 0},
                    {"value": "Wonderful service and quality", "label": 0},
                    {"value": "Exceeded all my expectations", "label": 0},
                    {"value": "This is terrible, waste of money", "label": 1},
                    {"value": "Horrible experience, never again", "label": 1},
                    {"value": "Worst product I've ever bought", "label": 1},
                    {"value": "Very disappointing quality", "label": 1},
                    {"value": "Complete garbage, don't buy", "label": 1},
                ]

                memoryset = LabeledMemoryset.create(
                    name=memoryset_name,
                    embedding_model=model,
                    label_names=["positive", "negative"],
                )
                memoryset.insert(sample_data)

                assert memoryset.length == len(sample_data)
                assert memoryset.embedding_model == model
                assert memoryset.label_names == ["positive", "negative"]

                # Test search on memoryset
                search_results = memoryset.search("This product is fantastic!", count=3)
                assert len(search_results) == 3
                assert search_results[0].label == 0  # Top result should be positive

                # Create classification model
                classifier = ClassificationModel.create(
                    name=classifier_name,
                    memoryset=memoryset,
                )
                assert classifier.name == classifier_name

                # Test predictions
                positive_prediction = classifier.predict("This is absolutely wonderful!")
                assert positive_prediction.label_name == "positive"

                negative_prediction = classifier.predict("I hate this, it's the worst")
                assert negative_prediction.label_name == "negative"

                # Verify memory lookups are populated
                assert len(classifier.last_prediction.memory_lookups) > 0

            finally:
                # Cleanup - cascade deletes handle dependencies
                FinetunedEmbeddingModel.drop(model_name, if_not_exists="ignore", cascade=True)
