from fastapi import FastAPI
from transformers import AutoTokenizer, T5ForConditionalGeneration
import json
import os
import logging

# Configura logging para capturar mais detalhes
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Carrega o arquivo questions.json
try:
    with open("questions.json", "r", encoding="utf-8") as f:
        examples = json.load(f)
    logger.info("questions.json carregado com sucesso.")
except FileNotFoundError:
    examples = []
    logger.warning("questions.json não encontrado, usando lista vazia.")

# Função para carregar o modelo e tokenizer sob demanda
def get_model():
    if not hasattr(get_model, "model_data"):
        logger.info("Carregando modelo e tokenizer pela primeira vez...")
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                "unicamp-dl/ptt5-base-portuguese-vocab",
                legacy=False,
                clean_up_tokenization_spaces=True
            )
            logger.info("Tokenizer carregado com sucesso.")
            model = T5ForConditionalGeneration.from_pretrained(
                "unicamp-dl/ptt5-base-portuguese-vocab",
                device_map="cpu"  # Força uso da CPU
            )
            logger.info("Modelo carregado com sucesso.")
            get_model.model_data = {"tokenizer": tokenizer, "model": model}
        except Exception as e:
            logger.error(f"Erro ao carregar modelo ou tokenizer: {e}")
            get_model.model_data = None
    return get_model.model_data

def generate_question_from_prompt(theme, difficulty, example_question=None):
    model_data = get_model()
    if not model_data or not model_data["tokenizer"] or not model_data["model"]:
        return {"question": "Erro: Modelo ou tokenizer não carregado.", "options": [], "answer": "", "explanation": "Por favor, verifique os logs."}

    tokenizer = model_data["tokenizer"]
    model = model_data["model"]

    if example_question:
        example_text = (
            f"Enunciado clínico: {example_question['question'].split('Considerando')[-1].strip()} "
            f"Alternativas: {', '.join(example_question['options'])} "
            f"Gabarito: {example_question['answer']} "
            f"Explicação: {example_question['explanation']}"
        )
        prompt = (
            f"Baseado no exemplo: '{example_text}', gere uma nova questão de múltipla escolha "
            f"sobre o tema '{theme}', com dificuldade {difficulty}, no estilo da prova de residência médica da USP. "
            "Use o formato: 'Enunciado clínico: [texto detalhado com dados do paciente, exames e pergunta]. "
            "Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. Gabarito: [letra]. "
            "Explicação: [texto].'"
        )
    else:
        prompt = (
            f"Gere uma questão de múltipla escolha sobre o tema '{theme}', "
            f"com dificuldade {difficulty}, no estilo da prova de residência médica da USP. "
            "Use o formato: 'Enunciado clínico: [texto detalhado com dados do paciente, exames e pergunta]. "
            "Alternativas: A) [opção], B) [opção], C) [opção], D) [opção]. Gabarito: [letra]. "
            "Explicação: [texto].'"
        )
    try:
        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
        outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7, top_p=0.9)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        parts = response.split("Alternativas:")
        if len(parts) > 1:
            question_part = parts[0].replace("Enunciado clínico:", "").strip()
            options_part = parts[1].split("Gabarito:")[0].strip()
            answer_part = parts[1].split("Gabarito:")[1].split("Explicação:")[0].strip()
            explanation_part = parts[1].split("Explicação:")[1].strip() if "Explicação:" in parts[1] else "Explicação padrão"
            options = [opt.strip() for opt in options_part.split(",")]
            if len(options) >= 4:
                return {
                    "question": f"Enunciado clínico: {question_part}",
                    "options": [f"A) {options[0]}", f"B) {options[1]}", f"C) {options[2]}", f"D) {options[3]}"],
                    "answer": answer_part,
                    "explanation": explanation_part
                }
        return {"question": response, "options": [], "answer": "", "explanation": "Explicação padrão"}
    except Exception as e:
        return {"question": f"Erro na geração: {e}", "options": [], "answer": "", "explanation": "Tente novamente."}

@app.get("/generate")
async def generate_question(theme: str, difficulty: str):
    example = examples[0] if examples else None
    return generate_question_from_prompt(theme, difficulty, example)

@app.get("/simulado")
async def get_simulado(num_questions: int = 5):
    simulado = []
    for _ in range(num_questions):
        example = examples[_ % len(examples)] if examples else None
        question_data = generate_question_from_prompt("clinica medica", "medio", example)
        simulado.append(question_data)
    return {"simulado": simulado}