app.py

import gradio as gr
import requests
import json
import os
import copy
import re # For more robust parsing

# --- Configuration ---
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    print("ALERTA: GEMINI_API_KEY não encontrada como Secret no Hugging Face Space. O app não funcionará.")
    # GEMINI_API_KEY = "YOUR_LOCAL_KEY_HERE" # Descomente para teste local

MODEL_NAME = "gemini-1.5-flash-latest"
API_URL_BASE = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL_NAME}:generateContent"

# --- Fases da AIJ ---
FASE_INICIAL_CONFIG = "CONFIGURACAO"
FASE_ABERTURA_CONCILIACAO = "ABERTURA_E_CONCILIACAO"
FASE_TESTEMUNHAS_AUTOR_INICIO = "TESTEMUNHAS_AUTOR_INICIO"
FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU = "TESTEMUNHAS_AUTOR_CONTRADITA_REU"
FASE_TESTEMUNHAS_REU_INICIO = "TESTEMUNHAS_REU_INICIO"
FASE_TESTEMUNHAS_REU_PERGUNTAS_REU = "TESTEMUNHAS_REU_PERGUNTAS_REU"
FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR = "TESTEMUNHAS_REU_INQUIRICAO_AUTOR"
FASE_ALEGACOES_FINAIS_AUTOR = "ALEGACOES_FINAIS_AUTOR"
FASE_ALEGACOES_FINAIS_REU = "ALEGACOES_FINAIS_REU"
FASE_ANALISE_FINAL_PROFESSOR = "ANALISE_FINAL_PROFESSOR"
FASE_CONCLUIDA = "AIJ_CONCLUIDA"

# --- Predefined Case Data ---
CASO_TITULO_PREDEFINIDO = "O Caso da Reportagem Revisitada"
CASO_PARTES_PREDEFINIDO = "Autor: Carlos Alberto Mendes\nRéu: Jornal 'A Verdade Diária'"
CASO_DESCRICAO_PREDEFINIDO = """
Uma pessoa (Carlos Alberto Mendes) processou um jornal (A Verdade Diária) alegando violação do direito ao esquecimento.
O jornal publicou uma matéria relembrando um crime ocorrido há 15 anos, no qual o autor foi
acusado de fraude, mas posteriormente absolvido por falta de provas. A matéria original foi publicada na época dos fatos
e a nova matéria, parte de uma série "Crimes que Marcaram a Cidade", revisitou o caso, mencionando a absolvição.

O autor alega que a republicação da história, mesmo com a menção da absolvição, causou-lhe
constrangimentos sociais, dificuldades em obter novo emprego e reacendeu traumas passados,
afetando sua vida familiar. Ele requer a remoção da matéria do site do jornal e uma indenização
por danos morais.

O jornal (réu) defende a liberdade de imprensa e o direito à informação, argumentando que
a matéria tem interesse público histórico, é factualmente correta (inclusive quanto à absolvição)
e faz parte do acervo jornalístico que documenta a história da comunidade.
"""
CASO_FATOS_COMPLETO_PREDEFINIDO = f"Título: {CASO_TITULO_PREDEFINIDO}\nPartes:\n{CASO_PARTES_PREDEFINIDO}\n\n---INÍCIO DOS FATOS---\n{CASO_DESCRICAO_PREDEFINIDO.strip()}"

ARGUMENTOS_AUTOR_PREDEFINIDOS = [
    "Argumento: A matéria publicada reabriu feridas do passado e causou danos psicológicos e sociais ao autor. Prova: Laudo psicológico atestando abalo emocional e depoimentos sobre constrangimento social.",
    "Argumento: O autor já foi absolvido judicialmente e tem direito a reconstruir sua vida sem o estigma do passado. Prova: Cópia da sentença absolutória.",
    "Argumento: O tempo transcorrido (15 anos) torna o caso desatualizado e sem relevância pública atual que justifique a exposição. Prova: Irrelevância pública atual, foco no dano individual.",
    "Argumento: A publicação violou o direito à privacidade e à honra do autor. Prova: Doutrina e jurisprudência sobre direitos da personalidade."
]
PROVAS_AUTOR_SUPPOSTAS_TEXTO_PREDEFINIDO = "\n".join(ARGUMENTOS_AUTOR_PREDEFINIDOS)

CONTRA_ARGUMENTOS_REU_PREDEFINIDOS_LIST = [
    "Argumento: A liberdade de imprensa é um direito fundamental e a matéria é de interesse público histórico. Prova: Art. 5º, IX e XIV da CF/88; Doutrina sobre liberdade de imprensa.",
    "Argumento: A matéria baseia-se inteiramente em informações verídicas e de acesso público (processo judicial). Prova: Cópia da matéria original e da recente, certidão do processo.",
    "Argumento: O STF (RE 1.010.606) não reconheceu o direito ao esquecimento de forma absoluta, ponderando-o com a liberdade de informação. Prova: Inteiro teor do acórdão do RE 1.010.606.",
    "Argumento: A matéria destacou a absolvição do autor, não havendo intenção de difamar. Prova: Conteúdo da própria matéria."
]
PROVAS_REU_TEXTO_PREDEFINIDO = "\n".join(CONTRA_ARGUMENTOS_REU_PREDEFINIDOS_LIST)

TESTEMUNHAS_REU_PREDEFINIDAS = [
    {"nome": "João Silva", "pontos_chave": "Editor-chefe do jornal. Confirmará o processo editorial, o interesse público da série de reportagens e a ausência de intenção de prejudicar o autor."},
    {"nome": "Dra. Maria Santos", "pontos_chave": "Historiadora local. Atestará a relevância histórica do caso para a comunidade e a importância de preservar a memória de tais eventos."},
    {"nome": "Roberto Almeida", "pontos_chave": "Jornalista responsável pela matéria. Detalhará a pesquisa, a verificação dos fatos e a menção explícita à absolvição do autor."}
]
TESTEMUNHAS_REU_TEXTO_PREDEFINIDO = "\n".join([f"{t['nome']}: {t['pontos_chave']}" for t in TESTEMUNHAS_REU_PREDEFINIDAS])

TESTEMUNHAS_AUTOR_SUPPOSTAS_PREDEFINIDAS_TEXTO = """
Dr. Felipe Antunes: Psicólogo do autor. Relatará o impacto emocional da nova reportagem na vida do autor.
Sra. Clara Dias: Vizinha do autor. Comentará sobre as mudanças no tratamento social que o autor recebeu após a reportagem.
"""

# ... (LEGISLACAO_PRECEDENTES_PREDEFINIDO as before) ...
LEGISLACAO_PRECEDENTES_PREDEFINIDO = {
    "favoraveis_ao_reu": [
        "Art. 5º, IV, IX e XIV da Constituição Federal - garantia da liberdade de expressão e acesso à informação",
        "ADPF 130 (STF) - declarou a Lei de Imprensa incompatível com a ordem constitucional atual",
        "RE 1.010.606 (STF) - não reconheceu o direito ao esquecimento na esfera civil de forma irrestrita",
        "Pacto de San José da Costa Rica - proteção à liberdade de expressão",
    ],
    "favoraveis_ao_autor": [
        "Art. 5º, X da Constituição Federal - inviolabilidade da intimidade, vida privada e honra",
        "Art. 21 do Código Civil - proteção à vida privada",
        "Enunciado 531 da VI Jornada de Direito Civil - reconhece o direito ao esquecimento",
        "LGPD (Lei Geral de Proteção de Dados) - embora não diretamente, reforça a proteção a dados pessoais."
    ]
}

def estado_inicial_simulacao(): # As before
    return {
        "fase_aij": FASE_INICIAL_CONFIG, "caso_fatos": "", "provas_reu_texto": "",
        "provas_autor_supostas_texto": "", "testemunhas_reu_defs": [],
        "testemunhas_autor_supostas_defs": [], "testemunha_autor_atual_idx": -1,
        "testemunha_reu_atual_idx": -1, "testemunha_em_depoimento": None,
        "transcricao_chat": [], "transcricao_logica": [],
        "feedback_para_usuario": "Aguardando início da simulação.",
        "instrucao_para_usuario": "Configure os detalhes do caso na aba 'Configuração' ou carregue o caso padrão."
    }

def chamar_gemini_api(prompt_text, temperature=0.7, max_output_tokens=800): # As before
    if not GEMINI_API_KEY:
        return "ERRO_FATAL: API Key do Gemini não configurada (verifique os Secrets do Space)."
    # ... (rest of API call logic)
    api_url_com_chave = f"{API_URL_BASE}?key={GEMINI_API_KEY}"
    headers = {'Content-Type': 'application/json'}
    data = {
        "contents": [{"parts": [{"text": prompt_text}]}],
        "generationConfig": {"temperature": temperature, "maxOutputTokens": max_output_tokens},
        "safetySettings": [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"}
        ]
    }
    try:
        response = requests.post(api_url_com_chave, headers=headers, data=json.dumps(data), timeout=180)
        response.raise_for_status()
        response_json = response.json()
        if response_json.get("candidates") and response_json["candidates"][0].get("content"):
            return response_json["candidates"][0]["content"]["parts"][0]["text"].strip()
        elif response_json.get("promptFeedback") and response_json["promptFeedback"].get("blockReason"):
            return f"ERRO_API: Conteúdo bloqueado. Razão: {response_json['promptFeedback']['blockReason']}"
        return "ERRO_API: Resposta inesperada da API."
    except requests.exceptions.Timeout:
        return "ERRO_CONEXAO: Timeout ao chamar a API Gemini. Tente novamente."
    except requests.exceptions.RequestException as e:
        error_msg = f"ERRO_CONEXAO: {e}"
        if e.response is not None:
            try: error_msg += f" - Detalhes: {e.response.json().get('error', {}).get('message', e.response.text)}"
            except json.JSONDecodeError: error_msg += f" - Resposta não JSON: {e.response.text}"
        print(error_msg)
        return error_msg
    except Exception as e:
        print(f"Erro desconhecido na API: {e}")
        return f"ERRO_DESCONHECIDO: {e}"


def formatar_historico_para_prompt(transcricao_logica, max_turns=15): # As before
    historico_formatado = "\nContexto da Audiência (últimas interações):\n"
    for ator, fala in transcricao_logica[-max_turns:]:
        fala_escaped = fala.replace('"', '\"')
        historico_formatado += f"- {ator}: {fala_escaped}\n"
    return historico_formatado

def gerar_caso_ficticio_com_ia():
    prompt_geracao_caso = """
**Instrução de Persona:** Você é um roteirista jurídico criativo e experiente, especializado em criar casos simulados realistas para treinamento de advogados.
**Tarefa:** Crie um caso fictício detalhado para uma simulação de Audiência de Instrução e Julgamento (AIJ) sobre o tema "Direito ao Esquecimento vs. Liberdade de Imprensa".
**Elementos Obrigatórios a Serem Gerados:**
1.  **Título do Caso Fictício:** (Ex: "O Caso da Notícia Ressurgida")
2.  **Partes:** (Autor: Nome Completo; Réu: Nome do Veículo de Imprensa)
3.  **Resumo Detalhado dos Fatos (Mínimo 250 palavras):** (Detalhes da publicação original, lapso temporal, detalhes da republicação/nova matéria, motivos do autor para se sentir lesado, impacto alegado em sua vida).
4.  **Argumentos e "Provas" Hipotéticas da Parte AUTORA (mínimo 3 itens, cada um em uma nova linha):** (Formato: `Argumento: [Descrição concisa do argumento]. Prova: [Descrição da prova hipotética que sustenta o argumento].`)
5.  **Argumentos e "Provas" Hipotéticas da Parte RÉ (Jornal - mínimo 3 itens, cada um em uma nova linha):** (Formato: `Argumento: [Descrição concisa do argumento]. Prova: [Descrição da prova hipotética que sustenta o argumento].`)
6.  **Testemunhas Hipotéticas da Parte AUTORA (Pelo menos 1, idealmente 2, cada uma em uma nova linha):** (Formato: `Nome Completo da Testemunha: [Breve descrição dos pontos chave que ela irá depor]. Ex: Maria Silva: Testemunhou o constrangimento social do autor após a reportagem.`)
7.  **Testemunhas Hipotéticas da Parte RÉ (Pelo menos 1, idealmente 2, cada uma em uma nova linha):** (Formato: `Nome Completo da Testemunha: [Breve descrição dos pontos chave que ela irá depor]. Ex: João Pereira: Editor responsável, falará sobre o interesse público e o processo editorial.`)
**Formato da Saída:** Use os seguintes marcadores exatos para cada seção (iniciando em nova linha):
`## Título do Caso Fictício`
`## Partes`
`## Resumo Detalhado dos Fatos`
`## Argumentos e "Provas" Hipotéticas da Parte AUTORA`
`## Argumentos e "Provas" Hipotéticas da Parte RÉ (Jornal)`
`## Testemunhas Hipotéticas da Parte AUTORA`
`## Testemunhas Hipotéticas da Parte RÉ`
O conteúdo de cada seção deve vir APENAS APÓS o respectivo marcador e antes do próximo marcador (ou fim do texto). Para Argumentos/Provas e Testemunhas, liste cada item individual em uma NOVA LINHA, seguindo o formato especificado. NÃO adicione marcadores de lista como '-' ou '*' a menos que explicitamente pedido no formato do item.
"""
    print("Chamando API para gerar caso...")
    texto_caso_completo = chamar_gemini_api(prompt_geracao_caso, temperature=0.8, max_output_tokens=3000) # Increased tokens
    if "ERRO_" in texto_caso_completo:
        print(f"Erro da API ao gerar caso: {texto_caso_completo}")
        return {"fatos": f"Falha ao gerar caso com IA: {texto_caso_completo}", "provas_autor": "", "provas_reu": "", "testemunhas_autor": "", "testemunhas_reu": ""}
    print("Caso gerado pela API (primeiros 500 chars):", texto_caso_completo[:500])
    parsed_data = parse_caso_gerado(texto_caso_completo)
    print("Dados parseados:", parsed_data)
    return parsed_data


def parse_caso_gerado(texto_completo_caso):
    print("Iniciando parse_caso_gerado...")
    dados_caso = {
        "fatos": "Falha parsing: Fatos não encontrados",
        "provas_autor": "Falha parsing: Provas do Autor não encontradas",
        "provas_reu": "Falha parsing: Provas do Réu não encontradas",
        "testemunhas_autor": "Falha parsing: Testemunhas do Autor não encontradas",
        "testemunhas_reu": "Falha parsing: Testemunhas do Réu não encontradas"
    }
    secoes_marcadores = {
        "titulo": "## Título do Caso Fictício",
        "partes": "## Partes",
        "fatos_resumo": "## Resumo Detalhado dos Fatos",
        "provas_autor": "## Argumentos e \"Provas\" Hipotéticas da Parte AUTORA",
        "provas_reu": "## Argumentos e \"Provas\" Hipotéticas da Parte RÉ (Jornal)",
        "testemunhas_autor_raw": "## Testemunhas Hipotéticas da Parte AUTORA", # Raw text for witnesses
        "testemunhas_reu_raw": "## Testemunhas Hipotéticas da Parte RÉ"    # Raw text for witnesses
    }

    # Helper to extract content between markers
    def extract_section(text, start_marker, berikutnya_markers):
        try:
            start_idx = text.lower().find(start_marker.lower())
            if start_idx == -1:
                print(f"Marcador inicial '{start_marker}' não encontrado.")
                return None
            
            content_start_idx = text.find("\n", start_idx + len(start_marker))
            if content_start_idx == -1: # If no newline, content starts right after marker
                 content_start_idx = start_idx + len(start_marker)
            else:
                content_start_idx +=1 # Move past the newline

            end_idx = len(text)
            for next_marker in berikutnya_markers:
                next_marker_idx = text.lower().find(next_marker.lower(), content_start_idx)
                if next_marker_idx != -1:
                    end_idx = min(end_idx, next_marker_idx)
            
            return text[content_start_idx:end_idx].strip()
        except Exception as e:
            print(f"Erro ao extrair seção para '{start_marker}': {e}")
            return None

    all_markers = list(secoes_marcadores.values())
    
    titulo_texto = extract_section(texto_completo_caso, secoes_marcadores["titulo"], [secoes_marcadores["partes"]])
    partes_texto = extract_section(texto_completo_caso, secoes_marcadores["partes"], [secoes_marcadores["fatos_resumo"]])
    fatos_resumo_texto = extract_section(texto_completo_caso, secoes_marcadores["fatos_resumo"], [secoes_marcadores["provas_autor"]])

    if titulo_texto and partes_texto and fatos_resumo_texto:
        dados_caso["fatos"] = f"Título: {titulo_texto}\nPartes:\n{partes_texto}\n\n---INÍCIO DOS FATOS---\n{fatos_resumo_texto}"
    elif fatos_resumo_texto: # Fallback if title/partes parsing fails
        dados_caso["fatos"] = fatos_resumo_texto


    dados_caso["provas_autor"] = extract_section(texto_completo_caso, secoes_marcadores["provas_autor"], [secoes_marcadores["provas_reu"]]) or dados_caso["provas_autor"]
    dados_caso["provas_reu"] = extract_section(texto_completo_caso, secoes_marcadores["provas_reu"], [secoes_marcadores["testemunhas_autor_raw"]]) or dados_caso["provas_reu"]
    
    # Store raw text for witnesses first, then parse them into structured list if needed by other functions
    dados_caso["testemunhas_autor"] = extract_section(texto_completo_caso, secoes_marcadores["testemunhas_autor_raw"], [secoes_marcadores["testemunhas_reu_raw"]]) or dados_caso["testemunhas_autor"]
    dados_caso["testemunhas_reu"] = extract_section(texto_completo_caso, secoes_marcadores["testemunhas_reu_raw"], []) or dados_caso["testemunhas_reu"] # No next marker for the last one

    # Check if any section remained as "Falha parsing" and log if so
    for key, value in dados_caso.items():
        if isinstance(value, str) and ("Falha parsing" in value or "não encontrado" in value):
            print(f"Alerta de Parsing - {key}: {value}")
            # Attempt to grab any text if section was missed due to marker order assumption
            if key == "testemunhas_reu" and "Falha parsing" in value: # Last item, try to grab till end
                try:
                    start_idx = texto_completo_caso.lower().find(secoes_marcadores["testemunhas_reu_raw"].lower())
                    if start_idx != -1:
                        content_start_idx = texto_completo_caso.find("\n", start_idx + len(secoes_marcadores["testemunhas_reu_raw"])) + 1
                        dados_caso[key] = texto_completo_caso[content_start_idx:].strip()
                        print(f"Reparsed {key} with fallback.")
                except:
                    pass


    print("Parse_caso_gerado concluído.")
    return dados_caso


def acao_gerar_caso_ficticio_ui(): # As before
    gr.Info("Gerando caso fictício com IA... Isso pode levar um momento.")
    dados_gerados = gerar_caso_ficticio_com_ia()
    return [dados_gerados["fatos"], dados_gerados["provas_reu"], dados_gerados["provas_autor"], dados_gerados["testemunhas_reu"], dados_gerados["testemunhas_autor"], "Caso fictício gerado. Revise e salve."]

def acao_carregar_caso_padrao_ui(): # As before
    gr.Info("Carregando caso padrão...")
    return [
        CASO_FATOS_COMPLETO_PREDEFINIDO,
        PROVAS_REU_TEXTO_PREDEFINIDO,
        PROVAS_AUTOR_SUPPOSTAS_TEXTO_PREDEFINIDO,
        TESTEMUNHAS_REU_TEXTO_PREDEFINIDO,
        TESTEMUNHAS_AUTOR_SUPPOSTAS_PREDEFINIDAS_TEXTO,
        "Caso padrão carregado. Revise e salve."
    ]

# --- Funções de Geração de Respostas da IA (Corrigidas, como antes) ---
def gerar_resposta_juiz_ia(estado_simulacao, acao_especifica=""):
    fase = estado_simulacao["fase_aij"]
    historico = formatar_historico_para_prompt(estado_simulacao["transcricao_logica"])
    prompt_base = f"**Persona:** Juiz(a) experiente em AIJ (Direito Esquecimento vs Liberdade Imprensa).\n**Fatos:** {estado_simulacao['caso_fatos'] or 'Não detalhado.'}\n{historico}**Fase Atual:** {fase}\n"
    tarefa_especifica_str = ""
    # ... (lógica de tarefa_especifica_str como antes) ...
    if fase == FASE_ABERTURA_CONCILIACAO: tarefa_especifica_str = "**Tarefa:** Declare audiência aberta, pergunte sobre conciliação. Se já houve manifestação e sem acordo, anuncie prosseguimento para testemunhas do autor."
    elif fase == FASE_TESTEMUNHAS_AUTOR_INICIO:
        testemunhas_autor_nomes = ', '.join([t['nome'] for t in estado_simulacao.get('testemunhas_autor_supostas_defs', [])]) or 'Nenhuma' # Added .get for safety
        tarefa_especifica_str = f"**Tarefa:** Anuncie oitiva testemunhas do Autor. Peça ao Adv. Autor (IA) para qualificar e inquirir a primeira. (Supostas test. Autor: {testemunhas_autor_nomes})"
    elif fase == FASE_TESTEMUNHAS_REU_INICIO:
        testemunhas_reu_nomes = ', '.join([t['nome'] for t in estado_simulacao.get('testemunhas_reu_defs', [])]) or 'Nenhuma' # Added .get for safety
        tarefa_especifica_str = f"**Tarefa:** Anuncie oitiva testemunhas do Réu. Peça ao Adv. Réu (Usuário) para indicar a primeira. (Test. Réu: {testemunhas_reu_nomes})"
    elif fase == FASE_ALEGACOES_FINAIS_AUTOR: tarefa_especifica_str = "**Tarefa:** Anuncie encerramento instrução. Conceda palavra ao Adv. Autor (IA) para alegações finais."
    elif fase == FASE_ALEGACOES_FINAIS_REU: tarefa_especifica_str = "**Tarefa:** Após alegações do Autor, conceda palavra ao Adv. Réu (Usuário) para suas alegações."
    elif fase == FASE_ANALISE_FINAL_PROFESSOR: tarefa_especifica_str = "**Tarefa:** Declare audiência encerrada para julgamento (simulado)."
    else: tarefa_especifica_str = f"**Tarefa:** Conduza a audiência. Ação específica: {acao_especifica or 'Nenhuma'}. Mantenha ordem e fluxo."
    prompt_completo = prompt_base + tarefa_especifica_str + "\n**Sua Fala (Juiz):**"
    return chamar_gemini_api(prompt_completo, temperature=0.6, max_output_tokens=350) # Slightly increased tokens

def gerar_resposta_advogado_autor_ia(estado_simulacao, fala_adv_reu=""): # As before (corrected)
    fase = estado_simulacao["fase_aij"]
    historico = formatar_historico_para_prompt(estado_simulacao["transcricao_logica"])
    prompt_base = f"**Persona:** Advogado(a) do AUTOR em AIJ (Direito Esquecimento vs Liberdade Imprensa). Cliente busca esquecimento/reparação.\n**Réu:** Jornal.\n**Fatos:** {estado_simulacao['caso_fatos'] or 'Não detalhado.'}\n**Provas/Argumentos Autor:** {estado_simulacao['provas_autor_supostas_texto'] or 'Passagem tempo, desinteresse público, dano vida privada.'}\n{historico}**Fase Atual:** {fase}\n"
    tarefa_especifica_str = ""
    fala_adv_reu_escaped = fala_adv_reu.replace('"', '\"')
    # ... (lógica de tarefa_especifica_str como antes) ...
    if fase == FASE_ABERTURA_CONCILIACAO:
        tarefa_especifica_str = f"**Tarefa:** Juiz perguntou sobre conciliação. Apresente proposta ou declare desinteresse. Se Réu já falou (\"{fala_adv_reu_escaped}\"), responda."
    elif fase == FASE_TESTEMUNHAS_AUTOR_INICIO:
        if estado_simulacao.get("testemunhas_autor_supostas_defs"): # .get for safety
            proxima_testemunha_idx = estado_simulacao.get("testemunha_autor_atual_idx", -1) + 1
            if proxima_testemunha_idx < len(estado_simulacao["testemunhas_autor_supostas_defs"]):
                testemunha_nome = estado_simulacao["testemunhas_autor_supostas_defs"][proxima_testemunha_idx]["nome"]
                tarefa_especifica_str = f"**Tarefa:** Juiz pediu para chamar testemunha. Chame {testemunha_nome}. Faça 1-2 perguntas iniciais (Ex: 'Sra. {testemunha_nome}, confirme conhecimento sobre X?')."
            else: return "Advogado do Autor (IA): Não tenho mais testemunhas a apresentar, Excelência."
        else: return "Advogado do Autor (IA): Não arrolei testemunhas, Excelência."
    elif fase == FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR:
        testemunha_reu_obj = estado_simulacao.get("testemunha_em_depoimento")
        if testemunha_reu_obj and testemunha_reu_obj["parte"] == "REU":
            pontos_chave_testemunha = testemunha_reu_obj.get('pontos_chave', 'Não informado')
            tarefa_especifica_str = f"**Tarefa:** Testemunha do Réu ({testemunha_reu_obj['nome']}) depôs. Formule 1-2 perguntas incisivas. (Pontos chave dela: {pontos_chave_testemunha})"
        else: return "Advogado do Autor (IA): (Erro interno - sem testemunha do réu para inquirir)"
    elif fase == FASE_ALEGACOES_FINAIS_AUTOR: tarefa_especifica_str = "**Tarefa:** Juiz concedeu palavra para alegações finais. Apresente resumo persuasivo e reitere pedidos."
    else:
        tarefa_especifica_str = f"**Contexto:** Adv. Réu (Usuário) disse: \"{fala_adv_reu_escaped}\"\n**Tarefa:** Responda estrategicamente."
    prompt_completo = prompt_base + tarefa_especifica_str + "\n**Sua Resposta (Advogado Autor):**"
    return chamar_gemini_api(prompt_completo, temperature=0.75, max_output_tokens=500) # Slightly increased

def gerar_resposta_testemunha_autor_ia(estado_simulacao, pergunta_feita): # As before (corrected)
    testemunha_obj = estado_simulacao.get("testemunha_em_depoimento")
    if not testemunha_obj or testemunha_obj["parte"] != "AUTOR": return "Testemunha do Autor (IA): (Erro interno - Nenhuma testemunha do autor em depoimento)"
    historico = formatar_historico_para_prompt(estado_simulacao["transcricao_logica"])
    pergunta_feita_escaped = pergunta_feita.replace('"', '\"')
    prompt_completo = f"**Persona:** Testemunha {testemunha_obj['nome']} (AUTOR).\n**Caso:** Direito Esquecimento vs Liberdade Imprensa.\n**Fatos:** {estado_simulacao['caso_fatos'] or 'Não detalhado.'}\n**Seus Pontos Chave:** {testemunha_obj['pontos_chave'] or 'Conhecimento geral sobre autor e impactos.'}\n{historico}**Pergunta Feita:** \"{pergunta_feita_escaped}\"\n**Tarefa:** Responda claro, conciso, baseado nos seus pontos. Se Adv. Réu, seja cauteloso(a).\n**Sua Resposta (Testemunha {testemunha_obj['nome']}):**"
    return chamar_gemini_api(prompt_completo, temperature=0.6, max_output_tokens=300) # Slightly increased

def analisar_atuacao_defesa_ia(estado_simulacao, ultima_fala_reu=""): # As before (corrected)
    historico_completo_dialogo = formatar_historico_para_prompt(estado_simulacao["transcricao_logica"], max_turns=len(estado_simulacao["transcricao_logica"]))
    ultima_fala_reu_escaped = ultima_fala_reu.replace('"', '\"')
    prompt_analise = f"**Persona:** Professor(a) de Prática Jurídica experiente.\n**Tarefa:** Feedback construtivo para aluno (Adv. RÉU - Jornal) em AIJ (Direito Esquecimento vs Liberdade Imprensa).\n**Fatos:** {estado_simulacao['caso_fatos'] or 'Não detalhado.'}\n**Provas/Argumentos Réu:** {estado_simulacao['provas_reu_texto'] or 'Não detalhado.'}\n**Transcrição AIJ:** {historico_completo_dialogo}\n**Última Intervenção Réu:** \"{ultima_fala_reu_escaped}\"\n**Feedback Detalhado (Pontos Fortes, Melhorias, Sugestões, Fundamentação, Estratégia):**"
    return chamar_gemini_api(prompt_analise, temperature=0.65, max_output_tokens=1500) # Increased

def avaliar_argumentos_defesa_ia(fatos_caso, provas_reu_texto_arg, tipo_argumento, argumento_usuario): # As before (corrected)
    # ... (prompt building logic as before)
    if not argumento_usuario.strip():
        return "Por favor, insira seus argumentos para avaliação."
    fatos_formatado = fatos_caso if fatos_caso else CASO_DESCRICAO_PREDEFINIDO
    provas_reu_formatado = provas_reu_texto_arg if provas_reu_texto_arg else PROVAS_REU_TEXTO_PREDEFINIDO
    contexto_caso_parts = [
        "Você está avaliando argumentos no contexto de um caso jurídico no Brasil sobre direito ao esquecimento vs. liberdade de imprensa.\n",
        "O usuário representa a defesa (jornal).\n",
        f"**Fatos do Caso:** {fatos_formatado}\n",
        f"**Principais Argumentos/Provas da Defesa (Jornal):** {provas_reu_formatado}\n"
    ]
    contexto_caso = "".join(contexto_caso_parts)
    argumento_usuario_escaped = argumento_usuario.replace('"', '\"') 
    tipo_argumento_escaped = tipo_argumento.replace('"', '\"')
    prompt_parts = [
        contexto_caso,
        "\n\n**Persona:** Você é um Juiz experiente em um tribunal brasileiro...\n", 
        f"**Tarefa:** Avalie o seguinte argumento apresentado pelo advogado de defesa (jornal) para a peça/momento processual: \"{tipo_argumento_escaped}\".\n",
        "Forneça feedback detalhado e construtivo em português brasileiro, incluindo:\n",
        "1.  **Pontos Fortes:** Destaque aspectos positivos (ex: fundamentação legal sólida, coerência lógica, clareza, persuasão, uso de fatos).\n",
        "2.  **Pontos a Melhorar:** Identifique fraquezas (ex: falta de aprofundamento, generalizações, omissão de contrapontos importantes, linguagem inadequada, falta de nexo com os fatos/provas do réu).\n",
        "3.  **Sugestões Específicas:** Ofereça conselhos práticos para aprimoramento (ex: citar precedentes específicos como o RE 1.010.606/STF, refutar explicitamente possíveis argumentos do autor, conectar melhor o argumento às provas do réu, sugerir redação alternativa).\n",
        "4.  **Nota Geral (0 a 5):** Atribua uma nota considerando:\n",
        "    *   Fundamentação Legal/Jurisprudencial (Peso 3)\n",
        "    *   Coerência e Lógica (Peso 2)\n",
        "    *   Persuasão e Clareza (Peso 2)\n",
        "    *   Adequação ao tipo de peça/momento (Peso 1)\n",
        "    *   Conexão com os fatos/provas do caso (Peso 2)\n\n",
        "Formato da Resposta: Utilize linguagem formal e jurídica. Estruture a resposta rigorosamente com os títulos:\n",
        "`### Pontos Fortes`\n",
        "`### Pontos a Melhorar`\n",
        "`### Sugestões Específicas`\n",
        "`### Nota Geral`\n\n",
        f"**Argumento do Advogado de Defesa (Jornal):**\n\"{argumento_usuario_escaped}\"\n\n",
        "**Sua Avaliação:**\n"
    ]
    prompt = "".join(prompt_parts)
    return chamar_gemini_api(prompt, temperature=0.6, max_output_tokens=1000)


def gerar_resposta_testemunha_defesa_para_treino_ia(fatos_caso, nome_testemunha, pontos_chave_testemunha, historico_perguntas_respostas_tuples, pergunta_atual_adv_defesa): # As before (corrected)
    # ... (prompt building logic as before)
    fatos_formatado = fatos_caso if fatos_caso else CASO_DESCRICAO_PREDEFINIDO
    contexto_caso_parts = [
        "Você está simulando o depoimento de uma testemunha de DEFESA (Jornal)...\n", 
        f"**Fatos do Caso:** {fatos_formatado}\n"
    ]
    contexto_caso = "".join(contexto_caso_parts)
    historico_formatado_str = "\nDepoimento até agora:\n"
    if historico_perguntas_respostas_tuples:
        for p, r in historico_perguntas_respostas_tuples: 
            p_escaped = p.replace('"', '\"')
            r_escaped = r.replace('"', '\"')
            historico_formatado_str += f"- Advogado Defesa: {p_escaped}\n- {nome_testemunha}: {r_escaped}\n"
    else:
        historico_formatado_str = "Este é o início do seu depoimento."
    pergunta_atual_escaped = pergunta_atual_adv_defesa.replace('"', '\"')
    prompt_parts = [
        contexto_caso,
        f"\n\n**Persona:** Você é {nome_testemunha}, testemunha arrolada pela DEFESA (Jornal).\n",
        f"**Seus Pontos Chave a Defender/Explicar:** {pontos_chave_testemunha}\n",
        "Você deve responder às perguntas do advogado de defesa de forma consistente com esses pontos e com a tese da defesa (liberdade de imprensa, interesse público, veracidade, ausência de dolo/culpa, destaque à absolvição do autor no caso original).\n",
        "Use linguagem formal, mas acessível para uma testemunha. Seja objetivo e direto. Evite admitir qualquer falha ou negligência por parte do jornal.\n",
        historico_formatado_str,
        f"\n\n**Pergunta do Advogado de Defesa (Jornal):** \"{pergunta_atual_escaped}\"\n\n",
        f"**Sua Resposta ({nome_testemunha}):**\n"
    ]
    prompt = "".join(prompt_parts)
    return chamar_gemini_api(prompt, temperature=0.5, max_output_tokens=300)

# --- Funções de Lógica da Simulação e UI (Adaptadas para type="messages") ---
def adicionar_fala_transcricao(estado_simulacao, ator, fala): # As before (corrected for type="messages")
    estado_simulacao["transcricao_logica"].append((ator, fala)) 
    if ator == "Advogado do Réu (Você)":
        estado_simulacao["transcricao_chat"].append({"role": "user", "content": fala})
    elif ator == "Sistema": 
        estado_simulacao["transcricao_chat"].append({"role": "assistant", "content": f"*{fala}*"})
    else: 
        estado_simulacao["transcricao_chat"].append({"role": "assistant", "content": f"**{ator}:**\n{fala}"})
    return estado_simulacao

def mostrar_botoes_corretos(fase_atual): # As before
    # ... (logic as before)
    vis = {
        "btn_responder_conciliacao": fase_atual == FASE_ABERTURA_CONCILIACAO,
        "btn_chamar_testemunha_reu": fase_atual == FASE_TESTEMUNHAS_REU_INICIO,
        "dropdown_testemunhas_reu": fase_atual == FASE_TESTEMUNHAS_REU_INICIO,
        "btn_perguntar_testemunha_autor": fase_atual == FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU,
        "btn_adv_reu_pergunta_sua_testemunha": fase_atual == FASE_TESTEMUNHAS_REU_PERGUNTAS_REU,
        "btn_adv_reu_responde_pela_sua_testemunha": fase_atual == FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR,
        "btn_apresentar_alegacoes_finais_reu": fase_atual == FASE_ALEGACOES_FINAIS_REU,
        "btn_juiz_avancar_fase": fase_atual not in [FASE_INICIAL_CONFIG, FASE_CONCLUIDA, FASE_ANALISE_FINAL_PROFESSOR] and fase_atual != FASE_ALEGACOES_FINAIS_REU
    }
    all_button_keys = [
        "btn_responder_conciliacao", "btn_chamar_testemunha_reu", "dropdown_testemunhas_reu",
        "btn_perguntar_testemunha_autor", "btn_adv_reu_pergunta_sua_testemunha",
        "btn_adv_reu_responde_pela_sua_testemunha", "btn_apresentar_alegacoes_finais_reu",
        "btn_juiz_avancar_fase"
    ]
    return [gr.update(visible=vis.get(k, False)) for k in all_button_keys]

def configurar_e_iniciar_simulacao(estado_s, fatos, p_reu, p_autor, t_reu_str, t_autor_str): # As before
    # ... (logic as before)
    novo_estado = estado_inicial_simulacao()
    if not fatos or not p_reu or not p_autor:
        gr.Warning("Por favor, preencha todos os campos de configuração do caso ou carregue o caso padrão.")
        initial_botoes_vis = mostrar_botoes_corretos(novo_estado["fase_aij"])
        return [novo_estado, novo_estado["transcricao_chat"], "Erro: Configuração incompleta.", "Preencha a configuração.", gr.Dropdown(choices=[], value=None, interactive=True)] + initial_botoes_vis

    novo_estado["caso_fatos"], novo_estado["provas_reu_texto"], novo_estado["provas_autor_supostas_texto"] = fatos, p_reu, p_autor
    def parse_testemunhas_local(t_str_local): 
        defs = []
        if t_str_local and t_str_local.strip():
            # More robust parsing for witnesses: Nome: Pontos Chave
            # Allow for variations, but prioritize the colon.
            for linha in t_str_local.strip().split('\n'):
                linha = linha.strip()
                if not linha: continue
                match = re.match(r"([^:]+):\s*(.*)", linha)
                if match:
                    nome, pontos = match.groups()
                    defs.append({"nome": nome.strip(), "pontos_chave": pontos.strip() or "Não especificado."})
                else: # If no colon, assume it's just a name or a malformed line
                    defs.append({"nome": linha, "pontos_chave": "Pontos chave não especificados claramente."})
        return defs

    novo_estado["testemunhas_reu_defs"] = parse_testemunhas_local(t_reu_str)
    novo_estado["testemunhas_autor_supostas_defs"] = parse_testemunhas_local(t_autor_str)

    novo_estado["fase_aij"] = FASE_ABERTURA_CONCILIACAO
    novo_estado = adicionar_fala_transcricao(novo_estado, "Sistema", "Configuração salva. Iniciando simulação...")
    gr.Info("Iniciando simulação... Gerando falas iniciais da IA.")
    fala_juiz = gerar_resposta_juiz_ia(novo_estado)
    novo_estado = adicionar_fala_transcricao(novo_estado, "Juiz (IA)", fala_juiz)
    fala_adv_autor = gerar_resposta_advogado_autor_ia(novo_estado)
    novo_estado = adicionar_fala_transcricao(novo_estado, "Advogado do Autor (IA)", fala_adv_autor)
    novo_estado["instrucao_para_usuario"] = "Responda à proposta de conciliação do Autor ou declare sua posição."
    novo_estado["feedback_para_usuario"] = analisar_atuacao_defesa_ia(novo_estado)
    nomes_testemunhas_reu = [t["nome"] for t in novo_estado["testemunhas_reu_defs"]] if novo_estado["testemunhas_reu_defs"] else []
    botoes_vis = mostrar_botoes_corretos(novo_estado["fase_aij"])
    gr.Info("Simulação iniciada.")
    return [novo_estado, novo_estado["transcricao_chat"], novo_estado["feedback_para_usuario"], novo_estado["instrucao_para_usuario"], gr.Dropdown(choices=nomes_testemunhas_reu, value=None, interactive=True)] + botoes_vis


def lidar_com_fala_usuario(estado_s, fala_usuario_txt): # As before
    # ... (entire logic from the previous corrected version)
    if not fala_usuario_txt and estado_s["fase_aij"] not in [FASE_TESTEMUNHAS_REU_INICIO]:
        is_text_needed = True
        if estado_s["fase_aij"] in [FASE_ABERTURA_CONCILIACAO, FASE_ALEGACOES_FINAIS_REU]:
            pass # Text is generally expected here unless user clicks "Juiz Avançar" (handled separately)
        elif estado_s["fase_aij"] in [FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU, FASE_TESTEMUNHAS_REU_PERGUNTAS_REU, FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR]:
            # Allow empty if user intends to click "Juiz Avançar"
            is_text_needed = False if not fala_usuario_txt else True # Only need text if user *types*
        
        if is_text_needed and not fala_usuario_txt:
            estado_s["instrucao_para_usuario"] = "Por favor, digite sua fala ou selecione uma ação que não exija texto."
            return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.Textbox(value="")] + mostrar_botoes_corretos(estado_s["fase_aij"])

    estado_s = copy.deepcopy(estado_s)
    ator_usuario = "Advogado do Réu (Você)"
    # ... rest of ator_usuario logic ...
    if estado_s["fase_aij"] == FASE_TESTEMUNHAS_REU_PERGUNTAS_REU and estado_s.get("testemunha_em_depoimento"):
        ator_usuario = f"Advogado do Réu (Você) perguntando a {estado_s['testemunha_em_depoimento']['nome']}"
    elif estado_s["fase_aij"] == FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR and estado_s.get("testemunha_em_depoimento"):
        ator_usuario = f"Testemunha {estado_s['testemunha_em_depoimento']['nome']} (Réu) respondendo"

    if fala_usuario_txt:
        estado_s = adicionar_fala_transcricao(estado_s, ator_usuario, fala_usuario_txt)
        gr.Info(f"Sua fala registrada. Aguardando respostas da IA...")
    
    # --- Lógica de Transição de Fases e Respostas da IA ---
    if estado_s["fase_aij"] == FASE_ABERTURA_CONCILIACAO:
        fala_adv_autor = gerar_resposta_advogado_autor_ia(estado_s, fala_adv_reu=fala_usuario_txt)
        estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor)
        fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Partes se manifestaram sobre conciliação. Decida o prosseguimento.")
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
        if any(k in fala_juiz.lower() for k in ["testemunha", "instrução", "prossegu", "oitiva"]):
            if estado_s.get("testemunhas_autor_supostas_defs"): # Use .get for safety
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_INICIO
                estado_s["instrucao_para_usuario"] = "Aguarde Juiz/Adv.Autor chamarem testemunha."
                fala_adv_autor_chamando = gerar_resposta_advogado_autor_ia(estado_s)
                estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor_chamando)
                if not ("Não tenho mais testemunhas" in fala_adv_autor_chamando or "Não arrolei testemunhas" in fala_adv_autor_chamando):
                    estado_s["testemunha_autor_atual_idx"] = estado_s.get("testemunha_autor_atual_idx", -1) + 1
                    idx = estado_s["testemunha_autor_atual_idx"]
                    if idx < len(estado_s["testemunhas_autor_supostas_defs"]):
                        t_info = estado_s["testemunhas_autor_supostas_defs"][idx]
                        estado_s["testemunha_em_depoimento"] = {"nome": t_info["nome"], "pontos_chave": t_info["pontos_chave"], "parte": "AUTOR"}
                        fala_test_autor = gerar_resposta_testemunha_autor_ia(estado_s, pergunta_feita=fala_adv_autor_chamando)
                        estado_s = adicionar_fala_transcricao(estado_s, f"Testemunha {t_info['nome']} (Autor - IA)", fala_test_autor)
                        estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU
                        estado_s["instrucao_para_usuario"] = f"Contradite/pergunte à testemunha {t_info['nome']} (Autor)."
            elif estado_s.get("testemunhas_reu_defs"): # Use .get for safety
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
                estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
            else: 
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
                estado_s["instrucao_para_usuario"] = "Aguarde Juiz iniciar Alegações Finais."
                fala_juiz_aleg = gerar_resposta_juiz_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_aleg)
                fala_aleg_autor = gerar_resposta_advogado_autor_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_aleg_autor)
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU 
                estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."
        else: 
            estado_s["instrucao_para_usuario"] = "Responda ao Juiz ou Adv. Autor sobre conciliação."

    elif estado_s["fase_aij"] == FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU:
        if not fala_usuario_txt:
             estado_s["instrucao_para_usuario"] = "Faça uma pergunta à testemunha do Autor ou clique em 'Juiz (IA): Avançar Fase' para encerrar."
        else:
            fala_test_autor = gerar_resposta_testemunha_autor_ia(estado_s, pergunta_feita=fala_usuario_txt)
            estado_s = adicionar_fala_transcricao(estado_s, f"Testemunha {estado_s['testemunha_em_depoimento']['nome']} (Autor - IA)", fala_test_autor)
            estado_s["instrucao_para_usuario"] = f"Faça outra pergunta à testemunha {estado_s['testemunha_em_depoimento']['nome']} (Autor) ou clique 'Juiz Avançar'."

    elif estado_s["fase_aij"] == FASE_TESTEMUNHAS_REU_PERGUNTAS_REU:
        if not fala_usuario_txt : 
            if estado_s.get("testemunha_em_depoimento"):
                estado_s["instrucao_para_usuario"] = f"Por favor, faça uma pergunta à sua testemunha {estado_s['testemunha_em_depoimento']['nome']}."
            else: # Should not happen if logic is correct for calling witness first
                estado_s["instrucao_para_usuario"] = "Erro: Nenhuma testemunha do réu em depoimento para perguntar."
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO # Go back to selection
        else:
            estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR
            fala_adv_autor_ia = gerar_resposta_advogado_autor_ia(estado_s) 
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor_ia)
            estado_s["instrucao_para_usuario"] = f"Como testemunha {estado_s['testemunha_em_depoimento']['nome']}, responda à pergunta do Adv. Autor."

    elif estado_s["fase_aij"] == FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR: 
        if not fala_usuario_txt:
            if estado_s.get("testemunha_em_depoimento"):
                estado_s["instrucao_para_usuario"] = f"Como testemunha {estado_s['testemunha_em_depoimento']['nome']}, por favor, responda à pergunta do Adv. Autor."
            else:
                estado_s["instrucao_para_usuario"] = "Erro: Nenhuma testemunha do réu em depoimento para responder."
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO # Go back
        else:
            fala_adv_autor_ia = gerar_resposta_advogado_autor_ia(estado_s, fala_adv_reu=fala_usuario_txt)
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor_ia)
            if ("pergunta" in fala_adv_autor_ia.lower() or "?" in fala_adv_autor_ia) and not any(k in fala_adv_autor_ia.lower() for k in ["não tenho mais perguntas", "encerro", "nada mais"]):
                estado_s["instrucao_para_usuario"] = f"Como testemunha {estado_s['testemunha_em_depoimento']['nome']}, responda à nova pergunta do Adv. Autor."
            else: 
                fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica=f"Adv. Autor encerrou inquirição da testemunha {estado_s.get('testemunha_em_depoimento',{}).get('nome','do Réu')}. Adv. Réu pode fazer mais perguntas ou dispensar.")
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_PERGUNTAS_REU 
                estado_s["instrucao_para_usuario"] = f"Faça mais perguntas à sua testemunha {estado_s.get('testemunha_em_depoimento',{}).get('nome','do Réu')}, ou clique 'Juiz Avançar' para dispensá-la."

    elif estado_s["fase_aij"] == FASE_ALEGACOES_FINAIS_REU:
        if not fala_usuario_txt:
            estado_s["instrucao_para_usuario"] = "Por favor, apresente suas Alegações Finais."
        else:
            estado_s["fase_aij"] = FASE_ANALISE_FINAL_PROFESSOR
            fala_juiz = gerar_resposta_juiz_ia(estado_s)
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
            estado_s["instrucao_para_usuario"] = "Audiência encerrada para sentença (simulada). Veja o feedback final do Professor IA."
    
    estado_s["feedback_para_usuario"] = analisar_atuacao_defesa_ia(estado_s, ultima_fala_reu=fala_usuario_txt if ator_usuario == "Advogado do Réu (Você)" else "")
    gr.Info("Respostas da IA geradas.")
    return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.Textbox(value="")] + mostrar_botoes_corretos(estado_s["fase_aij"])


def acao_botao_juiz_avancar(estado_s):
    # ... (logic of acao_botao_juiz_avancar as before, but ensure it returns 13 items)
    estado_s = copy.deepcopy(estado_s)
    fase_atual = estado_s["fase_aij"]
    gr.Info(f"Avançando fase a partir de: {fase_atual}. Aguarde IA...")
    # ... (all the phase transition logic from the previous correct version) ...

    # Ensure this structure is consistent across all branches of acao_botao_juiz_avancar
    # The original logic from your previous correct version should be here.
    # For brevity, I'm showing the end part.
    if fase_atual == FASE_ABERTURA_CONCILIACAO:
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", "Considerando frustrada a tentativa de conciliação ou por avanço solicitado, passaremos à instrução.")
        if estado_s.get("testemunhas_autor_supostas_defs"):
            estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_INICIO
            fala_juiz = gerar_resposta_juiz_ia(estado_s)
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
            fala_adv_autor = gerar_resposta_advogado_autor_ia(estado_s)
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor)
            if not ("Não tenho mais testemunhas" in fala_adv_autor or "Não arrolei testemunhas" in fala_adv_autor):
                estado_s["testemunha_autor_atual_idx"] = estado_s.get("testemunha_autor_atual_idx",-1) + 1 # Safe increment
                idx = estado_s["testemunha_autor_atual_idx"]
                if idx < len(estado_s["testemunhas_autor_supostas_defs"]):
                    t_info = estado_s["testemunhas_autor_supostas_defs"][idx]
                    estado_s["testemunha_em_depoimento"] = {"nome": t_info["nome"], "pontos_chave": t_info["pontos_chave"], "parte": "AUTOR"}
                    fala_test_autor = gerar_resposta_testemunha_autor_ia(estado_s, pergunta_feita=fala_adv_autor)
                    estado_s = adicionar_fala_transcricao(estado_s, f"Testemunha {t_info['nome']} (Autor - IA)", fala_test_autor)
                    estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU
                    estado_s["instrucao_para_usuario"] = f"Contradite ou faça perguntas à testemunha {t_info['nome']} (Autor)."
            else: 
                 if estado_s.get("testemunhas_reu_defs"): 
                     estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
                     estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
                     fala_juiz_trans = gerar_resposta_juiz_ia(estado_s)
                     estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_trans)

                 else: estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
        elif estado_s.get("testemunhas_reu_defs"):
            estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
            fala_juiz = gerar_resposta_juiz_ia(estado_s)
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
            estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
        else: 
            estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
            fala_juiz = gerar_resposta_juiz_ia(estado_s) 
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
            fala_adv_autor = gerar_resposta_advogado_autor_ia(estado_s) 
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor)
            estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
            estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."

    elif fase_atual == FASE_TESTEMUNHAS_AUTOR_INICIO:
        fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Verificar se há mais testemunhas do Autor ou encerrar esta fase.")
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
        fala_adv_autor = gerar_resposta_advogado_autor_ia(estado_s) 
        estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor)
        if "Não tenho mais testemunhas" in fala_adv_autor or "Não arrolei testemunhas" in fala_adv_autor or estado_s.get("testemunha_autor_atual_idx", -1) +1 >= len(estado_s.get("testemunhas_autor_supostas_defs",[])):
            estado_s["testemunha_em_depoimento"] = None
            if estado_s.get("testemunhas_reu_defs"):
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
                fala_juiz_trans = gerar_resposta_juiz_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_trans)
                estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
            else: 
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
                fala_juiz_trans = gerar_resposta_juiz_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_trans)
                fala_aleg_autor = gerar_resposta_advogado_autor_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_aleg_autor)
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
                estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."
        else: 
            estado_s["testemunha_autor_atual_idx"] = estado_s.get("testemunha_autor_atual_idx", -1) + 1
            idx = estado_s["testemunha_autor_atual_idx"]
            if idx < len(estado_s.get("testemunhas_autor_supostas_defs",[])):
                t_info = estado_s["testemunhas_autor_supostas_defs"][idx]
                estado_s["testemunha_em_depoimento"] = {"nome": t_info["nome"], "pontos_chave": t_info["pontos_chave"], "parte": "AUTOR"}
                fala_test_autor = gerar_resposta_testemunha_autor_ia(estado_s, pergunta_feita=fala_adv_autor) 
                estado_s = adicionar_fala_transcricao(estado_s, f"Testemunha {t_info['nome']} (Autor - IA)", fala_test_autor)
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU
                estado_s["instrucao_para_usuario"] = f"Contradite ou faça perguntas à testemunha {t_info['nome']} (Autor)."

    elif fase_atual == FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU:
        estado_s["testemunha_em_depoimento"] = None 
        if estado_s.get("testemunha_autor_atual_idx",-1) + 1 < len(estado_s.get("testemunhas_autor_supostas_defs",[])):
            estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_INICIO 
            fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Advogado do Réu encerrou perguntas à testemunha do Autor. Verificar próxima testemunha do Autor.")
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
            fala_adv_autor = gerar_resposta_advogado_autor_ia(estado_s) 
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor)
            if not ("Não tenho mais testemunhas" in fala_adv_autor or "Não arrolei testemunhas" in fala_adv_autor):
                estado_s["testemunha_autor_atual_idx"] = estado_s.get("testemunha_autor_atual_idx",-1) + 1
                idx_new = estado_s["testemunha_autor_atual_idx"]
                if idx_new < len(estado_s.get("testemunhas_autor_supostas_defs",[])):
                    t_info_new = estado_s["testemunhas_autor_supostas_defs"][idx_new]
                    estado_s["testemunha_em_depoimento"] = {"nome": t_info_new["nome"], "pontos_chave": t_info_new["pontos_chave"], "parte": "AUTOR"}
                    fala_test_autor_new = gerar_resposta_testemunha_autor_ia(estado_s, pergunta_feita=fala_adv_autor)
                    estado_s = adicionar_fala_transcricao(estado_s, f"Testemunha {t_info_new['nome']} (Autor - IA)", fala_test_autor_new)
                    estado_s["fase_aij"] = FASE_TESTEMUNHAS_AUTOR_CONTRADITA_REU
                    estado_s["instrucao_para_usuario"] = f"Contradite ou faça perguntas à testemunha {t_info_new['nome']} (Autor)."
            else: 
                 if estado_s.get("testemunhas_reu_defs"):
                     estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
                     fala_juiz_trans = gerar_resposta_juiz_ia(estado_s)
                     estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_trans)
                     estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
                 else:
                     estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
                     fala_juiz_trans = gerar_resposta_juiz_ia(estado_s)
                     estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_trans)
                     fala_aleg_autor = gerar_resposta_advogado_autor_ia(estado_s)
                     estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_aleg_autor)
                     estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
                     estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."
        else: 
            if estado_s.get("testemunhas_reu_defs"):
                estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO
                fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Encerrada oitiva das testemunhas do Autor. Iniciar oitiva das testemunhas do Réu.")
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
                estado_s["instrucao_para_usuario"] = "Selecione sua testemunha para chamar."
            else: 
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
                fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Instrução encerrada. Passar para alegações finais do Autor.")
                estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
                fala_aleg_autor = gerar_resposta_advogado_autor_ia(estado_s) 
                estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_aleg_autor)
                estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
                estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."

    elif fase_atual in [FASE_TESTEMUNHAS_REU_PERGUNTAS_REU, FASE_TESTEMUNHAS_REU_INQUIRICAO_AUTOR]:
        nome_test_dispensada = "a testemunha do Réu"
        if estado_s.get("testemunha_em_depoimento") and estado_s["testemunha_em_depoimento"]["parte"] == "REU":
            nome_test_dispensada = estado_s["testemunha_em_depoimento"]["nome"]
        estado_s["testemunha_em_depoimento"] = None
        
        fala_juiz_dispensa = gerar_resposta_juiz_ia(estado_s, acao_especifica=f"Dispensar {nome_test_dispensada}. Verificar se Adv. Réu tem mais testemunhas ou se encerrou sua instrução probatória.")
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_dispensa)

        more_reu_witnesses = estado_s.get("testemunha_reu_atual_idx",-1) + 1 < len(estado_s.get("testemunhas_reu_defs",[]))
        juiz_wants_to_continue_with_reu_witnesses = not any(k in fala_juiz_dispensa.lower() for k in ["alegações", "encerrada instrução", "sentença"])

        if more_reu_witnesses and juiz_wants_to_continue_with_reu_witnesses:
            estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_INICIO 
            estado_s["instrucao_para_usuario"] = "Selecione sua próxima testemunha para chamar ou clique 'Juiz Avançar' novamente para encerrar oitiva das suas testemunhas."
        else: 
            estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_AUTOR
            fala_juiz_aleg = gerar_resposta_juiz_ia(estado_s, acao_especifica="Encerrada instrução probatória. Passar para alegações finais do Autor.")
            estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz_aleg)
            fala_adv_autor_aleg = gerar_resposta_advogado_autor_ia(estado_s)
            estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Autor (IA)", fala_adv_autor_aleg)
            estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
            estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."
            
    elif fase_atual == FASE_ALEGACOES_FINAIS_AUTOR:
        estado_s["fase_aij"] = FASE_ALEGACOES_FINAIS_REU
        fala_juiz = gerar_resposta_juiz_ia(estado_s)
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
        estado_s["instrucao_para_usuario"] = "Apresente suas Alegações Finais."
    else:
        fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica="Conduza para a próxima etapa lógica da audiência.")
        estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
        estado_s["instrucao_para_usuario"] = "Siga as instruções do Juiz ou avance a fase."


    estado_s["feedback_para_usuario"] = analisar_atuacao_defesa_ia(estado_s)
    gr.Info("Fase avançada. Interações da IA geradas.")
    botoes_vis = mostrar_botoes_corretos(estado_s["fase_aij"])
    # Corrected: Add placeholder for fala_usuario_input
    return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.update(value="")] + botoes_vis


def chamar_testemunha_reu_selecionada(estado_s, nome_testemunha_selecionada): # As before
    # ... (logic as before, but also ensure it returns 13 items)
    estado_s = copy.deepcopy(estado_s)
    if not nome_testemunha_selecionada:
        estado_s["instrucao_para_usuario"] = "Nenhuma testemunha selecionada. Por favor, selecione uma testemunha da lista."
        gr.Warning("Selecione uma testemunha da lista.")
        return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.update(value="")] + mostrar_botoes_corretos(estado_s["fase_aij"])

    idx_testemunha = next((i for i, t in enumerate(estado_s["testemunhas_reu_defs"]) if t["nome"] == nome_testemunha_selecionada), -1)
    if idx_testemunha == -1:
        estado_s["instrucao_para_usuario"] = f"Testemunha '{nome_testemunha_selecionada}' não encontrada na configuração."
        gr.Error(f"Testemunha '{nome_testemunha_selecionada}' não encontrada.")
        return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.update(value="")] + mostrar_botoes_corretos(estado_s["fase_aij"])

    estado_s["testemunha_reu_atual_idx"] = idx_testemunha
    testemunha_obj = estado_s["testemunhas_reu_defs"][idx_testemunha]
    estado_s["testemunha_em_depoimento"] = {"nome": testemunha_obj["nome"], "pontos_chave": testemunha_obj["pontos_chave"], "parte": "REU"}

    gr.Info(f"Chamando testemunha {testemunha_obj['nome']}. Aguarde IA...")
    estado_s = adicionar_fala_transcricao(estado_s, "Advogado do Réu (Você)", f"Chamo para depor minha testemunha, Sr(a). {testemunha_obj['nome']}.")
    fala_juiz = gerar_resposta_juiz_ia(estado_s, acao_especifica=f"Advogado do Réu chamou a testemunha {testemunha_obj['nome']}. Qualifique-a e passe a palavra ao Advogado do Réu para inquirição inicial.")
    estado_s = adicionar_fala_transcricao(estado_s, "Juiz (IA)", fala_juiz)
    estado_s["fase_aij"] = FASE_TESTEMUNHAS_REU_PERGUNTAS_REU
    estado_s["instrucao_para_usuario"] = f"Faça sua primeira pergunta à sua testemunha, {testemunha_obj['nome']}. Após sua pergunta, o Advogado do Autor (IA) poderá inquirir (você responderá como sua testemunha)."
    estado_s["feedback_para_usuario"] = analisar_atuacao_defesa_ia(estado_s)
    gr.Info("Testemunha chamada. Interações da IA geradas.")
    botoes_vis = mostrar_botoes_corretos(estado_s["fase_aij"])
    return [estado_s, estado_s["transcricao_chat"], estado_s["feedback_para_usuario"], estado_s["instrucao_para_usuario"], gr.update(value="")] + botoes_vis

_historico_treino_testemunha_atual_messages = [] # Adapted for type="messages"
def gerar_dialogo_testemunha_defesa_para_treino_ui(fatos_caso_atual, testemunhas_reu_defs_atual_str, testemunha_selecionada_nome, pergunta_adv_defesa_treino):
    # ... (logic adapted as shown in the previous fix for this function for type="messages")
    global _historico_treino_testemunha_atual_messages
    if not GEMINI_API_KEY:
        _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": "ERRO: API Key não configurada."} )
        return _historico_treino_testemunha_atual_messages, ""

    if not testemunha_selecionada_nome:
        _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": "Por favor, selecione uma testemunha."} )
        return _historico_treino_testemunha_atual_messages, ""
    if not pergunta_adv_defesa_treino:
        _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": "Por favor, digite sua pergunta para a testemunha."} )
        return _historico_treino_testemunha_atual_messages, ""

    testemunhas_reu_defs_usar = []
    # ... (parsing testemunhas_reu_defs_usar from testemunhas_reu_defs_atual_str)
    if testemunhas_reu_defs_atual_str:
        try:
            linhas = testemunhas_reu_defs_atual_str.strip().split('\n')
            for linha in linhas:
                if linha.strip():
                    partes = linha.split(':', 1)
                    if len(partes) == 2:
                        testemunhas_reu_defs_usar.append({"nome": partes[0].strip(), "pontos_chave": partes[1].strip()})
                    elif len(partes) == 1:
                         testemunhas_reu_defs_usar.append({"nome": partes[0].strip(), "pontos_chave": "Não especificado."})
        except Exception as e:
            print(f"Erro ao parsear testemunhas_reu_defs_atual_str: {e}")
            testemunhas_reu_defs_usar = TESTEMUNHAS_REU_PREDEFINIDAS
    else:
        testemunhas_reu_defs_usar = TESTEMUNHAS_REU_PREDEFINIDAS


    testemunha_obj = next((t for t in testemunhas_reu_defs_usar if t["nome"] == testemunha_selecionada_nome), None)
    if not testemunha_obj: 
        testemunha_obj = next((t for t in TESTEMUNHAS_REU_PREDEFINIDAS if t["nome"] == testemunha_selecionada_nome), None)
        if not testemunha_obj:
            _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": f"Testemunha '{testemunha_selecionada_nome}' não encontrada."} )
            return _historico_treino_testemunha_atual_messages, ""

    reset_history = True
    if _historico_treino_testemunha_atual_messages:
        first_message = _historico_treino_testemunha_atual_messages[0]
        if first_message.get("role") == "assistant" and first_message.get("content", "").startswith(f"Iniciando inquirição de **{testemunha_obj['nome']}**"):
            reset_history = False
            
    if reset_history:
        _historico_treino_testemunha_atual_messages = []
        _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": f"Iniciando inquirição de **{testemunha_obj['nome']}**. Pontos chave: {testemunha_obj['pontos_chave']}"} )

    _historico_treino_testemunha_atual_messages.append({"role": "user", "content": pergunta_adv_defesa_treino})
    gr.Info(f"Gerando resposta da testemunha {testemunha_obj['nome']}... Aguarde.")
    
    llm_prompt_history_tuples = []
    temp_q = None
    for msg in _historico_treino_testemunha_atual_messages:
        if msg["role"] == "user":
            temp_q = msg["content"]
        elif msg["role"] == "assistant" and temp_q is not None:
            if not msg["content"].startswith("Iniciando inquirição"):
                 llm_prompt_history_tuples.append((temp_q, msg["content"]))
            temp_q = None

    resposta_testemunha = gerar_resposta_testemunha_defesa_para_treino_ia(
        fatos_caso_atual if fatos_caso_atual else CASO_FATOS_COMPLETO_PREDEFINIDO,
        testemunha_obj['nome'],
        testemunha_obj['pontos_chave'],
        llm_prompt_history_tuples, 
        pergunta_adv_defesa_treino
    )
    _historico_treino_testemunha_atual_messages.append( {"role": "assistant", "content": resposta_testemunha} )
    gr.Info("Resposta da testemunha gerada.")
    return _historico_treino_testemunha_atual_messages, ""


def limpar_chat_treino_tab5(): # As before
    global _historico_treino_testemunha_atual_messages
    _historico_treino_testemunha_atual_messages = []
    return [], "" 

# --- Interface Gráfica (Gradio UI) ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="orange"), title="Simulador AIJ Interativo") as demo:
    gr.Markdown("# ⚖️ Simulador Interativo de Audiência de Instrução e Julgamento (AIJ)")
    gr.Markdown("Tema: **Direito ao Esquecimento vs. Liberdade de Imprensa**. Você assume o papel de **Advogado(a) do Réu (Jornal)**.")
    if not GEMINI_API_KEY:
        gr.Warning("API Key do Gemini (GEMINI_API_KEY) não encontrada! O aplicativo pode não funcionar corretamente. Configure-a como Secret neste Space ou localmente como variável de ambiente.")

    sim_state = gr.State(value=estado_inicial_simulacao())

    with gr.Tabs() as abas_principais:
        with gr.TabItem("1. Configuração do Caso", id=0):
            gr.Markdown("## Detalhes do Caso para Simulação")
            with gr.Row():
                btn_gerar_caso_ia = gr.Button("✨ Gerar Novo Caso com IA", icon="🧠", scale=1)
                btn_carregar_caso_padrao = gr.Button("📄 Carregar Caso Padrão (Esquecimento)", icon="📜", scale=1)
            with gr.Row():
                with gr.Column(scale=2):
                    fatos_input = gr.Textbox(lines=7, label="Resumo dos Fatos, Partes e Título:", placeholder="Descreva o caso ou gere/carregue.")
                    provas_reu_input = gr.Textbox(lines=4, label="Argumentos e Provas do Réu (Seu Cliente - Jornal):", placeholder="Ex: Arg: Liberdade de imprensa. Prova: Art. 5 CF.")
                    provas_autor_input = gr.Textbox(lines=4, label="Argumentos e Provas (Supostas) do Autor:", placeholder="Ex: Arg: Dano à imagem. Prova: Laudo X.")
                with gr.Column(scale=1):
                    testemunhas_reu_input = gr.Textbox(lines=5, label="Suas Testemunhas (Réu) - Formato: Nome: Pontos Chave a Depor", placeholder="João Silva: Viu X e Y.\nMaria Costa: Sabe Z.")
                    testemunhas_autor_supostas_input = gr.Textbox(lines=5, label="Supostas Testemunhas do Autor - Formato: Nome: Pontos Chave", placeholder="Carlos Lima: Sofreu Z.\nAna B.: Testemunhará sobre Q.")
            btn_config_iniciar = gr.Button("💾 Salvar Configurações e Iniciar Simulação Interativa", variant="primary")
            status_config_output = gr.Label(value="Aguardando configuração do caso...")

            outputs_gerar_caso = [fatos_input, provas_reu_input, provas_autor_input, testemunhas_reu_input, testemunhas_autor_supostas_input, status_config_output]
            btn_gerar_caso_ia.click(acao_gerar_caso_ficticio_ui, inputs=None, outputs=outputs_gerar_caso)
            btn_carregar_caso_padrao.click(acao_carregar_caso_padrao_ui, inputs=None, outputs=outputs_gerar_caso)

        with gr.TabItem("2. Audiência Interativa", id=1):
            gr.Markdown("## Andamento da Audiência Simulada")
            instrucao_usuario_output = gr.Markdown("Instrução: Configure o caso na aba 'Configuração' e clique em 'Salvar e Iniciar' para começar.")
            chat_aij_output = gr.Chatbot(
                label="Transcrição da Audiência", 
                height=500, 
                show_copy_button=True, 
                avatar_images=(None, "https://img.icons8.com/fluency/96/robot-2.png"),
                type="messages" # Corrected
            )
            fala_usuario_input = gr.Textbox(show_label=False, placeholder="Digite sua fala aqui ou use os botões de ação abaixo...", scale=4, lines=3)
            
            btn_responder_conciliacao = gr.Button("Responder à Conciliação", visible=False)
            dropdown_testemunhas_reu_select = gr.Dropdown(label="Sua Testemunha a Chamar:", choices=[], visible=False, interactive=True)
            btn_chamar_testemunha_reu_action = gr.Button("Chamar Testemunha Selecionada", visible=False)
            btn_perguntar_testemunha_autor_action = gr.Button("Perguntar à Test. do Autor", visible=False)
            btn_adv_reu_pergunta_sua_testemunha_action = gr.Button("Perguntar à Minha Testemunha", visible=False)
            btn_adv_reu_responde_pela_sua_testemunha_action = gr.Button("Responder pela Minha Testemunha (à IA)", visible=False)
            btn_apresentar_alegacoes_finais_reu_action = gr.Button("Apresentar Alegações Finais", visible=False)
            btn_juiz_avancar_fase_action = gr.Button("⚖️ Juiz (IA): Concluir Etapa / Avançar Fase", visible=False, variant="secondary")

            with gr.Row():
                 btn_enviar_fala_usuario_display = gr.Button("🎤 Enviar Fala / Ação", variant="stop", scale=2) # Renamed for clarity in layout
                 # Removed gr.Space
                 btn_juiz_avancar_fase_action_display = btn_juiz_avancar_fase_action 

            with gr.Accordion("Ações Contextuais do Advogado do Réu (Você)", open=True):
                with gr.Row():
                    btn_responder_conciliacao_display = btn_responder_conciliacao
                    dropdown_testemunhas_reu_select_display = dropdown_testemunhas_reu_select
                    btn_chamar_testemunha_reu_action_display = btn_chamar_testemunha_reu_action
                with gr.Row():
                    btn_perguntar_testemunha_autor_action_display = btn_perguntar_testemunha_autor_action
                    btn_adv_reu_pergunta_sua_testemunha_action_display = btn_adv_reu_pergunta_sua_testemunha_action
                    btn_adv_reu_responde_pela_sua_testemunha_action_display = btn_adv_reu_responde_pela_sua_testemunha_action
                    btn_apresentar_alegacoes_finais_reu_action_display = btn_apresentar_alegacoes_finais_reu_action
            
            lista_botoes_ui_objects = [ # The actual UI component objects
                btn_responder_conciliacao, btn_chamar_testemunha_reu_action, dropdown_testemunhas_reu_select,
                btn_perguntar_testemunha_autor_action, btn_adv_reu_pergunta_sua_testemunha_action,
                btn_adv_reu_responde_pela_sua_testemunha_action, btn_apresentar_alegacoes_finais_reu_action,
                btn_juiz_avancar_fase_action
            ]

        with gr.TabItem("3. Feedback da Atuação (Professor IA)", id=2):
            gr.Markdown("## 💡 Análise da Sua Atuação na Audiência Interativa")
            feedback_output = gr.Textbox(lines=20, label="Feedback Detalhado:", interactive=False, show_copy_button=True)

        with gr.TabItem("4. Treinamento de Argumentação (Defesa)", id=3):
            gr.Markdown("## 🗣️ Pratique e Avalie Seus Argumentos (Advogado de Defesa)")
            fatos_caso_treino_arg_display_tab4 = gr.Textbox(label="Fatos do Caso (da Configuração):", value=CASO_FATOS_COMPLETO_PREDEFINIDO, lines=5, interactive=False)
            provas_reu_treino_arg_display_tab4 = gr.Textbox(label="Provas do Réu (da Configuração):", value=PROVAS_REU_TEXTO_PREDEFINIDO, lines=5, interactive=False)
            tipo_argumento_treino_tab4 = gr.Radio(
                ["Contestação (teses principais)", "Perguntas para Testemunha do Autor", "Perguntas para Testemunha do Réu", "Alegações Finais (memoriais)"],
                label="Selecione o Tipo de Argumento/Peça para Treinar:", value="Alegações Finais (memoriais)"
            )
            argumento_input_treino_tab4 = gr.Textbox(lines=10, label="Digite seu argumento aqui:", placeholder="Ex: 'Excelência, a matéria jornalística é um retrato fiel dos fatos públicos...'")
            btn_avaliar_argumento_treino_tab4 = gr.Button("🤖 Avaliar Argumento com IA", variant="primary")
            avaliacao_output_treino_tab4 = gr.Markdown(label="Feedback da IA sobre o Argumento:")

            fatos_input.change(lambda x: x, inputs=fatos_input, outputs=fatos_caso_treino_arg_display_tab4)
            provas_reu_input.change(lambda x: x, inputs=provas_reu_input, outputs=provas_reu_treino_arg_display_tab4)

            btn_avaliar_argumento_treino_tab4.click(
                avaliar_argumentos_defesa_ia,
                inputs=[fatos_caso_treino_arg_display_tab4, provas_reu_treino_arg_display_tab4, tipo_argumento_treino_tab4, argumento_input_treino_tab4],
                outputs=avaliacao_output_treino_tab4
            )

        with gr.TabItem("5. Preparar Testemunhas (Defesa)", id=4):
            gr.Markdown("## 🎙️ Simule a Inquirição das Suas Testemunhas (Defesa)")
            fatos_caso_treino_test_display_tab5 = gr.Textbox(label="Fatos do Caso (da Configuração):", value=CASO_FATOS_COMPLETO_PREDEFINIDO, lines=5, interactive=False)
            testemunhas_reu_config_treino_display_tab5 = gr.Textbox(label="Definição das Testemunhas do Réu (da Configuração):", value=TESTEMUNHAS_REU_TEXTO_PREDEFINIDO, lines=3, interactive=False)

            fatos_input.change(lambda x: x, inputs=fatos_input, outputs=fatos_caso_treino_test_display_tab5)
            testemunhas_reu_input.change(lambda x: x, inputs=testemunhas_reu_input, outputs=testemunhas_reu_config_treino_display_tab5)

            dropdown_testemunhas_reu_treino_tab5 = gr.Dropdown(label="Selecione Sua Testemunha para Inquirir:", choices=[t['nome'] for t in TESTEMUNHAS_REU_PREDEFINIDAS], interactive=True)
            
            def update_dropdown_treino_choices_tab5(t_reu_str_val_tab5):
                nomes = []
                if t_reu_str_val_tab5 and t_reu_str_val_tab5.strip():
                    for linha in t_reu_str_val_tab5.strip().split('\n'):
                        if linha.strip():
                            partes = linha.split(':', 1)
                            nomes.append(partes[0].strip())
                return gr.Dropdown(choices=nomes if nomes else [t['nome'] for t in TESTEMUNHAS_REU_PREDEFINIDAS], value=nomes[0] if nomes else None, interactive=True)
            testemunhas_reu_input.change(update_dropdown_treino_choices_tab5, inputs=testemunhas_reu_input, outputs=dropdown_testemunhas_reu_treino_tab5)

            chat_treino_testemunha_tab5 = gr.Chatbot(
                label="Simulação de Depoimento (Treino)", 
                height=300, 
                show_copy_button=True, 
                avatar_images=(None, "https://img.icons8.com/officel/80/person-male.png"),
                type="messages" # Corrected
            )
            pergunta_testemunha_treino_input_tab5 = gr.Textbox(label="Sua Pergunta para a Testemunha:", placeholder="Ex: 'Senhor Silva, poderia descrever o procedimento editorial?'")
            btn_perguntar_testemunha_treino_tab5 = gr.Button("💬 Perguntar à Testemunha (IA)", variant="primary")
            btn_limpar_chat_treino_testemunha_tab5 = gr.Button("🗑️ Limpar Chat de Treino")
            
            btn_limpar_chat_treino_testemunha_tab5.click(limpar_chat_treino_tab5, outputs=[chat_treino_testemunha_tab5, pergunta_testemunha_treino_input_tab5])
            btn_perguntar_testemunha_treino_tab5.click(
                gerar_dialogo_testemunha_defesa_para_treino_ui,
                inputs=[fatos_caso_treino_test_display_tab5, testemunhas_reu_config_treino_display_tab5, dropdown_testemunhas_reu_treino_tab5, pergunta_testemunha_treino_input_tab5],
                outputs=[chat_treino_testemunha_tab5, pergunta_testemunha_treino_input_tab5]
            )

        with gr.TabItem("ℹ️ Ajuda / Sobre o Caso Padrão", id=5):
            # ... (Content of Help Tab as before) ...
            gr.Markdown("## Sobre o Caso Padrão: Direito ao Esquecimento vs. Liberdade de Imprensa")
            gr.Markdown(f"**Título:** {CASO_TITULO_PREDEFINIDO}")
            md_partes = CASO_PARTES_PREDEFINIDO.replace("\n", "\n\n")
            gr.Markdown(f"**Partes:**\n{md_partes}")
            gr.Markdown(f"**Resumo do Caso:**\n{CASO_DESCRICAO_PREDEFINIDO}") 
            with gr.Accordion("Argumentos Principais do Autor (Caso Padrão)", open=False):
                gr.Markdown("\n".join([f"- {arg}" for arg in ARGUMENTOS_AUTOR_PREDEFINIDOS]))
            with gr.Accordion("Argumentos Principais do Réu/Jornal (Caso Padrão)", open=False):
                gr.Markdown("\n".join([f"- {arg}" for arg in CONTRA_ARGUMENTOS_REU_PREDEFINIDOS_LIST]))
            with gr.Accordion("Testemunhas do Réu/Jornal (Caso Padrão)", open=False):
                gr.Markdown("\n".join([f"- **{t['nome']}**: {t['pontos_chave']}" for t in TESTEMUNHAS_REU_PREDEFINIDAS]))

            gr.Markdown("### Legislação e Precedentes Relevantes (Geral)")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("**Potencialmente Favoráveis ao Réu (Jornal):**")
                    gr.Markdown("\n".join([f"- {lei}" for lei in LEGISLACAO_PRECEDENTES_PREDEFINIDO["favoraveis_ao_reu"]]))
                with gr.Column():
                    gr.Markdown("**Potencialmente Favoráveis ao Autor:**")
                    gr.Markdown("\n".join([f"- {lei}" for lei in LEGISLACAO_PRECEDENTES_PREDEFINIDO["favoraveis_ao_autor"]]))
            gr.Markdown("""
            ## Como utilizar este simulador: 
            ... (rest of help text as before) ...
            """)


    # --- Event Handlers ---
    # Outputs for config_iniciar: 5 base + 8 buttons = 13
    outputs_config_list = [
        sim_state, chat_aij_output, feedback_output, instrucao_usuario_output, 
        dropdown_testemunhas_reu_select 
    ] + lista_botoes_ui_objects # Use the actual UI objects for outputs

    btn_config_iniciar.click(
        configurar_e_iniciar_simulacao,
        inputs=[sim_state, fatos_input, provas_reu_input, provas_autor_input, testemunhas_reu_input, testemunhas_autor_supostas_input],
        outputs=outputs_config_list
    ).then(lambda: gr.Textbox(value=""), outputs=[fala_usuario_input]
    ).then(lambda: gr.Tabs(selected=1), inputs=None, outputs=[abas_principais])

    # Outputs for fala_usuario and general actions: 5 base + 8 buttons = 13
    outputs_fala_usuario_and_actions_list = [
        sim_state, chat_aij_output, feedback_output, instrucao_usuario_output, 
        fala_usuario_input # This output is to clear the textbox after submit/click
    ] + lista_botoes_ui_objects

    # Use the renamed btn_enviar_fala_usuario_display if it's the one in the layout
    btn_enviar_fala_usuario_display.click(lidar_com_fala_usuario, inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list)
    fala_usuario_input.submit(lidar_com_fala_usuario, inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list)

    btn_juiz_avancar_fase_action.click(acao_botao_juiz_avancar, inputs=[sim_state], outputs=outputs_fala_usuario_and_actions_list) # .then to clear input removed as main func now returns update for it
    btn_chamar_testemunha_reu_action.click(chamar_testemunha_reu_selecionada, inputs=[sim_state, dropdown_testemunhas_reu_select], outputs=outputs_fala_usuario_and_actions_list) # .then removed

    # ... (Button click handlers as before, using outputs_fala_usuario_and_actions_list)
    def create_button_click_handler_wrapper(info_message):
        def actual_handler_for_button(current_sim_state, current_fala_usuario_input_text):
            gr.Info(info_message)
            return lidar_com_fala_usuario(current_sim_state, current_fala_usuario_input_text)
        return actual_handler_for_button

    btn_responder_conciliacao.click(
        create_button_click_handler_wrapper("Conciliação: Digite sua resposta no campo de fala e clique 'Enviar Fala / Ação'."),
        inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list
    )
    btn_perguntar_testemunha_autor_action.click(
        create_button_click_handler_wrapper("Perg. Test. Autor: Digite sua pergunta no campo de fala e clique 'Enviar Fala / Ação'."),
        inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list
    )
    btn_adv_reu_pergunta_sua_testemunha_action.click(
        create_button_click_handler_wrapper("Perg. Minha Test.: Digite sua pergunta no campo de fala e clique 'Enviar Fala / Ação'."),
        inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list
    )
    btn_adv_reu_responde_pela_sua_testemunha_action.click(
        create_button_click_handler_wrapper("Resp. p/ Minha Test.: Digite a resposta da sua testemunha no campo de fala e clique 'Enviar Fala / Ação'."),
        inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list
    )
    btn_apresentar_alegacoes_finais_reu_action.click(
        create_button_click_handler_wrapper("Alegações Finais: Digite suas alegações no campo de fala e clique 'Enviar Fala / Ação'."),
        inputs=[sim_state, fala_usuario_input], outputs=outputs_fala_usuario_and_actions_list
    )

    # Outputs for demo.load: 5 base + 8 buttons = 13
    outputs_on_load_list = [
        sim_state, chat_aij_output, feedback_output, instrucao_usuario_output, 
        dropdown_testemunhas_reu_select 
    ] + lista_botoes_ui_objects

    def setup_initial_ui_state_load_wrapper(s_val_load):
        initial_fase = s_val_load["fase_aij"]
        dropdown_choices_load = [t["nome"] for t in s_val_load.get("testemunhas_reu_defs", [])] if s_val_load.get("testemunhas_reu_defs") else []
        updates_for_buttons_load = mostrar_botoes_corretos(initial_fase)
        
        # Order: sim_state, chat, feedback, instrucao, dropdown, then the 8 buttons
        final_return_list = [
            s_val_load,                             # state
            s_val_load["transcricao_chat"],         # chatbot
            s_val_load["feedback_para_usuario"],    # textbox (feedback)
            s_val_load["instrucao_para_usuario"],   # markdown (instrucao)
            gr.Dropdown(choices=dropdown_choices_load, value=None, interactive=True) # dropdown
        ]
        final_return_list.extend(updates_for_buttons_load) # the 8 button updates
        return tuple(final_return_list)

    demo.load(
        setup_initial_ui_state_load_wrapper,
        inputs=[sim_state],
        outputs=outputs_on_load_list
    )

if __name__ == "__main__":
    if not GEMINI_API_KEY:
        print("#####################################################################")
        print("ATENÇÃO: GEMINI_API_KEY não configurada...")
        print("#####################################################################")
    demo.queue(max_size=20).launch(debug=True)