import torch
import openai
import os
from transformers import AutoTokenizer
from evo_model import EvoTransformerV22
from rag_utils import extract_text_from_file
from search_utils import web_search

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = EvoTransformerV22()
model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
model.eval()

def format_input(question, options, context, web_results):
    prompt = f"{question}\n"
    if context:
        prompt += f"\nContext:\n{context}\n"
    if web_results:
        prompt += f"\nWeb Search Results:\n" + "\n".join(web_results)
    prompt += "\nOptions:\n"
    for idx, opt in enumerate(options):
        prompt += f"{idx+1}. {opt}\n"
    return prompt.strip()

def get_evo_response(question, context, options, enable_search=True):
    web_results = web_search(question) if enable_search else []
    input_text = format_input(question, options, context, web_results)
    encoded = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=256)
    with torch.no_grad():
        logits = model(encoded["input_ids"])
    probs = torch.softmax(logits, dim=1).squeeze()
    pred_index = torch.argmax(probs).item()
    confidence = probs[pred_index].item()

    suggestion = options[pred_index] if pred_index < len(options) else "N/A"
    evo_reasoning = f"Evo suggests: **{suggestion}** (Confidence: {confidence:.2f})\n\nContext used:\n" + "\n".join(web_results)
    return suggestion, evo_reasoning

def get_gpt_response(question, context, options):
    openai.api_key = os.getenv("OPENAI_API_KEY", "")
    formatted_options = "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(options)])
    prompt = f"Question: {question}\n\nContext:\n{context}\n\nOptions:\n{formatted_options}\n\nWhich option makes the most sense and why?"

    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful reasoning assistant."},
                {"role": "user", "content": prompt}
            ]
        )
        return response['choices'][0]['message']['content']
    except Exception as e:
        return f"⚠️ GPT error: {str(e)}"