import os
from pathlib import Path
from typing import Callable, TypeVar

from dotenv import load_dotenv


F = TypeVar("F", bound=Callable)
# Repo root (parent of puppet_theater/), not cwd — matches app.py and tests.
load_dotenv(Path(__file__).resolve().parent.parent / ".env")

def _env_flag(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


USE_ZEROGPU = _env_flag("USE_ZEROGPU", default=False)

try:
    import spaces
except ImportError:
    spaces = None

SPACES_AVAILABLE = spaces is not None
ZEROGPU_GPU_ACTIVE = USE_ZEROGPU and SPACES_AVAILABLE
LAST_GPU_CUDA_AVAILABLE: bool | None = None
LAST_GPU_FALLBACK_REASON: str | None = (
    "USE_ZEROGPU=true but the spaces package is not installed"
    if USE_ZEROGPU and not SPACES_AVAILABLE
    else None
)
import torch

def gpu_decorator(func: F) -> F:
    if not ZEROGPU_GPU_ACTIVE or spaces is None:
        return func
    return spaces.GPU(duration=30)(func)


def torch_version() -> str:
    try:
        import torch
    except ImportError:
        return "not installed"
    return str(getattr(torch, "__version__", "unknown"))


def record_gpu_failure(reason: str) -> None:
    global LAST_GPU_FALLBACK_REASON
    LAST_GPU_FALLBACK_REASON = " ".join(reason.split())[:180]


@gpu_decorator
def generate_openbmb_text_on_zerogpu(
    model_id: str,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
) -> str:
    global LAST_GPU_CUDA_AVAILABLE, LAST_GPU_FALLBACK_REASON

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    LAST_GPU_CUDA_AVAILABLE = bool(torch.cuda.is_available())
    LAST_GPU_FALLBACK_REASON = None

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto",
    )
    model.eval()

    messages = [{"role": "user", "content": prompt}]
    try:
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            enable_thinking=False,
            return_dict=True,
            return_tensors="pt",
        )
    except TypeError:
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )

    inputs = inputs.to(model.device)
    eos_token_id = tokenizer.eos_token_id
    pad_token_id = tokenizer.pad_token_id or eos_token_id
    do_sample = temperature > 0
    generation_kwargs = {
        "max_new_tokens": max_new_tokens,
        "do_sample": do_sample,
        "pad_token_id": pad_token_id,
        "eos_token_id": eos_token_id,
    }
    if do_sample:
        generation_kwargs["temperature"] = temperature
    with torch.inference_mode():
        outputs = model.generate(**inputs, **generation_kwargs)
    new_tokens = outputs[0][inputs["input_ids"].shape[-1] :]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()