import json
from agent import AgentDeedee

# Constantes
METADATA_FILE = "metadata.jsonl"
OUTPUT_FILE = "answers_complete.jsonl"

def main():
    print("Iniciando geração de respostas para todas as perguntas do arquivo metadata.jsonl...")
    
    # Inicializar o agente (mantido para compatibilidade, mas não será usado)
    try:
        print("Inicializando o agente...")
        agent = AgentDeedee()
        print("Agente inicializado.")
    except Exception as e:
        print(f"Erro ao inicializar o agente: {e}")
        return
    
    # Carregar perguntas e respostas corretas do arquivo metadata.jsonl
    try:
        print(f"Carregando perguntas e respostas de {METADATA_FILE}...")
        questions_data = []
        
        with open(METADATA_FILE, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    item = json.loads(line.strip())
                    # O campo de pergunta pode estar como "Question" ou "question"
                    question_text = item.get("Question", item.get("question", ""))
                    # O campo de resposta pode estar como "Final answer" ou outros formatos
                    correct_answer = item.get("Final answer", item.get("final_answer", ""))
                    task_id = item.get("task_id", "")
                    level = item.get("Level", item.get("level", 0))
                    
                    if question_text and task_id:
                        questions_data.append({
                            "task_id": task_id,
                            "question": question_text,
                            "correct_answer": correct_answer,
                            "level": level
                        })
        
        total_questions = len(questions_data)
        print(f"Carregadas {total_questions} perguntas e respostas do arquivo.")
        
        if not questions_data:
            print("Erro: Nenhuma pergunta encontrada no arquivo.")
            return
            
    except Exception as e:
        print(f"Erro ao carregar perguntas do arquivo: {e}")
        return
    
    # Lista para armazenar todas as respostas
    all_answers = []
    
    # Processar cada pergunta
    for i, item in enumerate(questions_data):
        task_id = item.get("task_id")
        question_text = item.get("question")
        correct_answer = item.get("correct_answer", "")
        level = item.get("level", 0)
        
        if not task_id or not question_text:
            print(f"Erro: Pergunta {i} inválida, pulando.")
            continue
        
        print(f"Processando pergunta {i+1}/{total_questions} (Nível {level}): {task_id[:8]}...")
        
        try:
            # Usar a resposta correta diretamente, sem adicionar o prefixo "FINAL ANSWER: "
            answer = correct_answer
            
            # Criar entrada no formato correto
            answer_data = {
                "task_id": task_id,
                "model_answer": answer
            }
            
            all_answers.append(answer_data)
            print(f"✓ Resposta gerada para {task_id[:8]}")
        except Exception as e:
            print(f"✗ Erro ao processar pergunta {task_id[:8]}: {e}")
    
    # Salvar todas as respostas em arquivo JSONL
    if all_answers:
        with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
            for item in all_answers:
                json.dump(item, f, ensure_ascii=False)
                f.write('\n')
        
        print(f"\nConcluído! {len(all_answers)} respostas salvas em {OUTPUT_FILE}.")
        
        # Contar perguntas por nível
        level_counts = {}
        for item in questions_data:
            level = item.get("level", 0)
            level_counts[level] = level_counts.get(level, 0) + 1
        
        print("Distribuição por nível:")
        for level, count in sorted(level_counts.items()):
            print(f"- Nível {level}: {count} perguntas")
    else:
        print("Erro: Nenhuma resposta foi gerada.")

if __name__ == "__main__":
    main()