import os
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = None
tokenizer = None

def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")

device = get_device()

def load_model(weights_filename: str = "trained_weights.pt"):
    global model, tokenizer, device
    device = get_device()
    
    print("[ENGINE] Initializing tokenizer and GPT-2 model...", flush=True)
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token
    
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    
    root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
    weights_path = os.path.join(root_dir, weights_filename)
    
    if os.path.exists(weights_path):
        print(f"[ENGINE] Loading custom weights from {weights_path}", flush=True)
        state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict)
    else:
        print(f"[ENGINE] Warning: {weights_filename} not found at {weights_path}. Using base GPT-2.", flush=True)
        
    model.to(device)
    model.eval()

def generate_sql(prompt_string: str, n_candidates: int = 5) -> list[str]:
    global model, tokenizer
    
    # Lazy initialization: auto-load if not already in memory
    if tokenizer is None or model is None:
        load_model()

    inputs = tokenizer(prompt_string, return_tensors="pt").to(device)
    prompt_len = inputs.input_ids.shape[1]

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=48,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            do_sample=True,
            temperature=0.4,
            num_return_sequences=n_candidates
        )

    candidates = []
    for i in range(n_candidates):
        gen_tokens = output_ids[i][prompt_len:]
        decoded = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
        candidates.append(decoded)

    return candidates