from __future__ import annotations

import os
import sys
from pathlib import Path

import pytest
import torch

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"

if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from inteprelens import CHG, LensAPI, TransformerTracer  # noqa: E402
from inteprelens.models.adapters import get_adapter  # noqa: E402


def _candidate_env_paths() -> list[Path]:
    return [
        ROOT / ".env",
        ROOT.parent / ".env",
    ]


def _load_env() -> None:
    try:
        from dotenv import load_dotenv
    except ImportError:
        return

    for env_path in _candidate_env_paths():
        if env_path.exists():
            load_dotenv(dotenv_path=env_path)

    token = (
        os.getenv("hf_token")
        or os.getenv("HUGGINGFACE_HUB_TOKEN")
        or os.getenv("HF_TOKEN")
    )
    if token:
        os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", token)
        os.environ.setdefault("HF_TOKEN", token)

os.environ.setdefault("CHG_DEFAULT_MODEL", "meta-llama/Llama-3.2-1B")
_load_env()


def pytest_configure() -> None:
    os.environ.setdefault("CHG_DEFAULT_MODEL", "meta-llama/Llama-3.2-1B")
    _load_env()


@pytest.fixture(scope="session")
def llama_model_name() -> str:
    return os.getenv("CHG_DEFAULT_MODEL", "meta-llama/Llama-3.2-1B")


@pytest.fixture(scope="session")
def hf_token() -> str:
    token = (
        os.getenv("hf_token")
        or os.getenv("HUGGINGFACE_HUB_TOKEN")
        or os.getenv("HF_TOKEN")
    )
    if not token:
        pytest.skip("HF_TOKEN is required for the shared real-model test suite.")
    return token


@pytest.fixture(scope="session")
def _authenticated_hf(hf_token: str) -> str:
    try:
        from huggingface_hub import login

        login(hf_token, add_to_git_credential=False, skip_if_logged_in=True)
    except Exception:
        # Environment variables are enough for most Hugging Face flows.
        pass
    return hf_token


@pytest.fixture(scope="session")
def llama_tokenizer(llama_model_name: str, _authenticated_hf: str):
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


@pytest.fixture(scope="session")
def llama_model(llama_model_name: str, _authenticated_hf: str):
    from transformers import AutoModelForCausalLM

    if not torch.cuda.is_available():
        pytest.skip("CUDA is required for the shared real-model test suite.")

    model = AutoModelForCausalLM.from_pretrained(
        llama_model_name,
        torch_dtype=torch.bfloat16,
    )
    model.to("cuda")
    model.eval()
    for param in model.parameters():
        param.requires_grad = False
    return model


@pytest.fixture()
def shared_model(llama_model):
    llama_model.zero_grad(set_to_none=True)
    return llama_model


@pytest.fixture()
def llama_adapter(shared_model):
    return get_adapter(shared_model)


@pytest.fixture()
def llama_chg(shared_model, llama_adapter):
    return CHG(shared_model, adapter=llama_adapter)


@pytest.fixture()
def llama_api(shared_model, llama_tokenizer, llama_adapter):
    return LensAPI(shared_model, llama_tokenizer, adapter=llama_adapter, device="cuda")


@pytest.fixture()
def llama_tracer(shared_model, llama_adapter):
    return TransformerTracer(shared_model, adapter=llama_adapter)
