import torch
from evo_model import EvoTransformer
from transformers import AutoTokenizer, pipeline
from rag_utils import RAGRetriever, extract_text_from_file
import os

# Load Evo model
def load_evo_model(model_path="evo_hellaswag.pt", device=None):
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    model = EvoTransformer()
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()
    return model, device

evo_model, device = load_evo_model()
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# Load GPT-3.5 (via OpenAI API)
import openai
openai.api_key = os.getenv("OPENAI_API_KEY")

# RAG Retriever
retriever = RAGRetriever()

def get_context_from_file(file_obj):
    file_path = file_obj.name
    text = extract_text_from_file(file_path)
    retriever.add_document(text)
    return text

# Evo prediction
def get_evo_response(prompt, file=None):
    # Step 1: augment context if document is uploaded
    context = ""
    if file is not None:
        context_list = retriever.retrieve(prompt)
        context = "\n".join(context_list)

    full_prompt = f"{prompt}\n{context}"

    # Step 2: use Evo to predict
    options = ["Yes, proceed with the action.", "No, maintain current strategy."]
    inputs = [f"{full_prompt} {opt}" for opt in options]

    encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)

    with torch.no_grad():
        logits = evo_model(encoded["input_ids"]).squeeze(-1)
        probs = torch.softmax(logits, dim=0)
        best = torch.argmax(probs).item()

    return f"Evo suggests: {options[best]} (Confidence: {probs[best]:.2f})"

# GPT-3.5 response
def get_gpt_response(prompt, file=None):
    context = ""
    if file is not None:
        context_list = retriever.retrieve(prompt)
        context = "\n".join(context_list)

    full_prompt = (
        f"Question: {prompt}\n"
        f"Relevant Context:\n{context}\n"
        f"Answer like a financial advisor."
    )

    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": full_prompt}
        ],
        temperature=0.4,
    )

    return response.choices[0].message.content.strip()

    #