import os
import json
from typing import Any, Dict
import openai
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ---------- OpenAI ----------
def openai_chat(prompt: str, model: str = "gpt-4o-mini") -> str:
    client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7,
    )
    return resp.choices[0].message.content.strip()


# ---------- Local transformer fallback ----------
_local_model = None
_local_tokenizer = None

def _load_local_model() -> None:
    global _local_model, _local_tokenizer
    if _local_model is None:
        model_name = "microsoft/phi-2"          # лёгкая модель, подходит для CPU
        _local_tokenizer = AutoTokenizer.from_pretrained(model_name)
        _local_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            device_map="auto",
        )

def local_generate(prompt: str, max_new_tokens: int = 100) -> str:
    _load_local_model()
    inputs = _local_tokenizer(prompt, return_tensors="pt").to(_local_model.device)
    with torch.no_grad():
        out = _local_model.generate(**inputs, max_new_tokens=max_new_tokens)
    return _local_tokenizer.decode(out[0], skip_special_tokens=True)
