import os
import gradio as gr
import pypdf
import io
import tempfile
from fpdf import FPDF
from google import genai

# --- Configuração da API ---
# MANTENDO EXATAMENTE A SUA LÓGICA ORIGINAL
try:
    api_key = os.getenv("GOOGLE_API_KEY") 
    if not api_key:
        raise ValueError("Chave de API do Google não encontrada. Defina o 'Secret' chamado 'GOOGLE_API_KEY' no Hugging Face Spaces.")
    
    # Voltando a usar genai.Client() como no seu código
    client = genai.Client()
    print("Cliente Genai configurado com sucesso.")

except Exception as e:
    print(f"Erro ao configurar o cliente Genai: {e}")
    client = None # A verificação 'if client is None' tratará isso

# --- Funções Auxiliares ---

def read_file_content(file_obj):
    """Lê o conteúdo de um arquivo PDF ou TXT."""
    # (Sem alterações nesta função)
    if file_obj is None:
        return None, "Erro: Arquivo não fornecido."
    
    file_path = file_obj.name
    
    try:
        if file_path.endswith('.pdf'):
            reader = pypdf.PdfReader(file_path)
            text = ""
            for page in reader.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
            return text, None
        elif file_path.endswith('.txt'):
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read(), None
        else:
            return None, "Erro: Formato de arquivo não suportado. Use .pdf ou .txt"
    except Exception as e:
        return None, f"Erro ao ler o arquivo: {str(e)}"

# --- LÓGICA DE PROMPT (Com as novas especializações) ---

def build_reclamatoria_prompt(texto_caso, texto_peca):
    """Cria um prompt específico para analisar uma RECLAMATÓRIA TRABALHISTA."""
    # (Sem alterações, esta é a nova lógica de prompt)
    return f"""
    Você é um Juiz do Trabalho experiente e um Advogado Sênior da Reclamada (defesa).
    Sua tarefa é analisar uma "Reclamatória Trabalhista" (Peça) com base nos fatos do "Caso".

    **DOCUMENTO 1: O CASO (Fatos e Problema)**
    ---
    {texto_caso}
    ---

    **DOCUMENTO 2: A PEÇA (Reclamatória Trabalhista do Reclamante)**
    ---
    {texto_peca}
    ---

    Execute as seguintes tarefas com base **exclusivamente** nos documentos fornecidos:

    **TAREFA 1: ANÁLISE DE FALHAS E SUGESTÕES (Visão do Juiz/Defesa)**
    Identifique falhas, omissões, argumentos fracos ou pontos vulneráveis na "PEÇA" (Reclamatória).
    Foqie em:
    - Pedidos sem causa de pedir correspondente.
    - Argumentos frágeis ou sem base nos fatos.
    - Possíveis prescrições (bienal, quinquenal) não observadas.
    - Pedidos que parecem inepto ou que podem ser facilmente contestados.
    
    Para cada falha, forneça uma **sugestão clara de melhoria** (como o advogado do reclamante poderia ter fortalecido aquele ponto).

    **TAREFA 2: SIMULAÇÃO DE ATAQUE ADVERSÁRIO (Teses de Contestação)**
    Gere **exatamente 3 (três)** teses de **Contestação** que a empresa (Reclamada) poderia usar para atacar os pontos fracos da "PEÇA". Cada simulação deve ser uma estratégia de defesa (ex: "Arguição de prescrição", "Impugnação de horas extras por falta de prova", "Tese de justa causa").

    Formate sua resposta **exatamente** da seguinte maneira:

    ### ANÁLISE DE FALHAS E SUGESTÕES
    **Falha 1:** [Descrição da falha na Reclamatória]
    **Sugestão 1:** [Descrição da sugestão de melhoria para o Reclamante]
    ...

    ---[SEPARADOR]---

    ### SIMULAÇÃO DE ATAQUE ADVERSÁRIO (TESES DE CONTESTAÇÃO)
    **Ataque 1 (Tese de Defesa):** [Descrição da primeira tese de contestação]
    **Ataque 2 (Tese de Defesa):** [Descrição da segunda tese de contestação]
    **Ataque 3 (Tese de Defesa):** [Descrição da terceira tese de contestação]
    """

def build_contestacao_prompt(texto_caso, texto_peca):
    """Cria um prompt específico para analisar uma CONTESTAÇÃO TRABALHISTA."""
    # (Sem alterações, esta é a nova lógica de prompt)
    return f"""
    Você é um Juiz do Trabalho experiente e um Advogado Sênior do Reclamante.
    Sua tarefa é analisar uma "Contestação" (Peça) com base nos fatos do "Caso" e na provável petição inicial.

    **DOCUMENTO 1: O CASO (Fatos e Problema - que originou a Reclamatória)**
    ---
    {texto_caso}
    ---

    **DOCUMENTO 2: A PEÇA (Contestação da Reclamada)**
    ---
    {texto_peca}
    ---

    Execute as seguintes tarefas com base **exclusivamente** nos documentos fornecidos:

    **TAREFA 1: ANÁLISE DE FALHAS E SUGESTÕES (Visão do Juiz/Reclamante)**
    Identifique falhas, omissões, argumentos fracos ou pontos vulneráveis na "PEÇA" (Contestação).
    Foque em:
    - Fatos da petição inicial que NÃO foram impugnados (que se tornarão incontroversos).
    - Preliminares fracas ou mal fundamentadas.
    - Argumentos de mérito que contrariam os fatos do "Caso".
    - Falta de impugnação específica aos pedidos (presunção de veracidade).
    
    Para cada falha, forneça uma **sugestão clara de melhoria** (como o advogado da reclamada poderia ter fortalecido aquela defesa).

    **TAREFA 2: SIMULAÇÃO DE ATAQUE ADVERSÁRIO (Teses de Réplica)**
    Gere **exatamente 3 (três)** teses de **Impugnação à Contestação (Réplica)** que o Reclamante poderia usar para refutar os argumentos da "PEÇA". Cada simulação deve ser uma estratégia de réplica (ex: "Apontar fato incontroverso", "Refutar tese de mérito", "Rebater preliminar").

    Formate sua resposta **exatamente** da seguinte maneira:

    ### ANÁLISE DE FALHAS E SUGESTÕES
    **Falha 1:** [Descrição da falha na Contestação]
    **Sugestão 1:** [Descrição da sugestão de melhoria para a Reclamada]
    ...

    ---[SEPARADOR]---

    ### SIMULAÇÃO DE ATAQUE ADVERSÁRIO (TESES DE RÉPLICA)
    **Ataque 1 (Tese de Réplica):** [Descrição da primeira tese de réplica]
    **Ataque 2 (Tese de Réplica):** [Descrição da segunda tese de réplica]
    **Ataque 3 (Tese de Réplica):** [Descrição da terceira tese de réplica]
    """

def build_generic_prompt(texto_caso, texto_peca):
    """Cria o prompt genérico original."""
    # (Sem alterações, este é o prompt original)
    return f"""
    Você é um assistente sênior de análise jurídica e estratégia processual.
    Sua tarefa é analisar dois documentos: um "Caso" e uma "Peça" (a defesa ou petição).

    **DOCUMENTO 1: O CASO (Fatos e Problema)**
    ---
    {texto_caso}
    ---

    **DOCUMENTO 2: A PEÇA (Solução/Resposta ao Caso)**
    ---
    {texto_peca}
    ---

    Execute as seguintes tarefas com base **exclusivamente** nos documentos fornecidos:

    **TAREFA 1: ANÁLISE de FALHAS E SUGESTÕES**
    Identifique falhas, omissões, argumentos fracos ou pontos vulneráveis na "PEÇA". 
    Para cada falha identificada, forneça uma **sugestão clara e construtiva de melhoria** (como corrigir ou fortalecer o argumento).

    **TAREFA 2: SIMULAÇÃO DE ATAQUE ADVERSÁRIO**
    Gere **exatamente 3 (três)** simulações de como uma equipe adversária inteligente poderia atacar os pontos fracos da "PEÇA". Cada simulação deve ser uma estratégia de contra-argumentação.

    Formate sua resposta **exatamente** da seguinte maneira:

    ### ANÁLISE DE FALHAS E SUGESTÕES
    **Falha 1:** [Descrição da falha]
    **Sugestão 1:** [Descrição da sugestão de melhoria]
    ...

    ---[SEPARADOR]---

    ### SIMULAÇÃO DE ATAQUE ADVERSÁRIO
    **Ataque 1:** [Descrição do primeiro ataque]
    **Ataque 2:** [Descrição do segundo ataque]
    **Ataque 3:** [Descrição do terceiro ataque]
    """

def build_prompt(texto_caso, texto_peca, tipo_analise):
    """Seleciona o prompt correto com base na escolha do usuário."""
    # (Sem alterações, esta é a nova lógica de seleção)
    if tipo_analise == "Reclamatória Trabalhista":
        return build_reclamatoria_prompt(texto_caso, texto_peca)
    elif tipo_analise == "Contestação Trabalhista":
        return build_contestacao_prompt(texto_caso, texto_peca)
    else: # "Análise Genérica"
        return build_generic_prompt(texto_caso, texto_peca)


def create_pdf(text_content):
    """Cria um arquivo PDF temporário com o texto da simulação."""
    # (Sem alterações nesta função)
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            pdf = FPDF()
            pdf.add_page()
            
            pdf.set_font("DejaVu", size=12)
            
            pdf.multi_cell(0, 10, text_content)
            
            pdf.output(temp_file.name)
            
            return temp_file.name
    except Exception as e:
        print(f"Erro ao criar PDF: {e}")
        return None

# --- Função Principal (Modificada para aceitar 'tipo_analise') ---

def analisar_e_simular_com_progresso(tipo_analise, caso_file, peca_file): # <-- PARÂMETRO ADICIONADO
    """Função geradora que atualiza a UI com o progresso."""
    
    # 1. Estado inicial
    yield {
        btn_analisar: gr.Button(value="Analisando...", interactive=False),
        output_analise: "Aguarde, lendo arquivos...",
        output_simulacao: "",
        output_pdf: None 
    }

    # *** CORREÇÃO AQUI ***
    # Verificando 'client' como na sua lógica original
    if client is None:
        yield {
            btn_analisar: gr.Button(value="Erro na API", interactive=False),
            output_analise: "Erro: A API do Google Genai não está configurada. Adicione a GOOGLE_API_KEY aos Secrets.",
            output_simulacao: "",
            output_pdf: None
        }
        return

    if caso_file is None or peca_file is None:
        # (Lógica idêntica)
        yield {
            btn_analisar: gr.Button(value="Analisar e Simular", interactive=True),
            output_analise: "Erro: Por favor, envie ambos os arquivos.",
            output_simulacao: "",
            output_pdf: None
        }
        return

    # 2. Ler arquivos
    texto_caso, erro_caso = read_file_content(caso_file)
    texto_peca, erro_peca = read_file_content(peca_file)
    
    if erro_caso or erro_peca:
        # (Lógica idêntica)
        yield {
            btn_analisar: gr.Button(value="Analisar e Simular", interactive=True),
            output_analise: erro_caso or erro_peca,
            output_simulacao: "",
            output_pdf: None
        }
        return

    # 3. Estado de progresso (Atualizado para mostrar o tipo de análise)
    yield {
        output_analise: f"Arquivos lidos com sucesso.\n\nAnalisando como: {tipo_analise}.\n\nChamando a API das IA's... (Isso pode levar alguns segundos)",
    }

    # 4. Criar o prompt especializado
    prompt = build_prompt(texto_caso, texto_peca, tipo_analise) 
    
    try:
        # 5. Chamar a API (*** CORREÇÃO AQUI ***)
        # Revertido para a sua lógica original e o modelo que você especificou
        response = client.models.generate_content(
            model="gemini-2.5-flash",  # <-- MANTIDO CONFORME SEU PEDIDO
            contents=prompt
        )
        full_text = response.text
        
        # 6. Processar a resposta (Lógica idêntica)
        if "---[SEPARADOR]---" in full_text.upper():
            try:
                analise, simulacao = full_text.split(next(s for s in full_text.splitlines() if "---[SEPARADOR]---" in s.upper()), 1)
            except (StopIteration, ValueError):
                 analise, simulacao = full_text.split("---[SEPARADOR]---", 1)
        else:
            analise = "Não foi possível dividir a resposta. Resposta completa:\n" + full_text
            simulacao = "Não foi possível dividir a resposta. Resposta completa:\n" + full_text
        
        # 7. Criar o PDF (Lógica idêntica)
        analise_limpa = analise.strip()
        simulacao_limpa = simulacao.strip()
        pdf_path = create_pdf(simulacao_limpa) 

        # 8. Estado final (Lógica idêntica)
        yield {
            btn_analisar: gr.Button(value="Analisar e Simular", interactive=True),
            output_analise: analise_limpa,
            output_simulacao: simulacao_limpa,
            output_pdf: pdf_path 
        }
        
    except Exception as e:
        # 9. Estado de erro (Lógica idêntica)
        yield {
            btn_analisar: gr.Button(value="Analisar e Simular", interactive=True),
            output_analise: f"Erro ao chamar a API do Genai: {str(e)}",
            output_simulacao: "",
            output_pdf: None
        }

def clear_fields():
    """Limpa todos os campos de entrada e saída."""
    # Adicionado o 'tipo_analise_input' para ser limpo
    return {
        tipo_analise_input: "Análise Genérica", 
        caso_input: None,
        peca_input: None,
        output_analise: "",
        output_simulacao: "",
        output_pdf: None 
    }

# --- Interface Gráfica (Com o novo seletor) ---

with gr.Blocks(theme=gr.themes.Soft(), title="Analisador e Simulador de Peças") as demo:
    gr.Markdown(
        """
        # Analisador e Simulador de Peças (com inteligencia artificial)
        Envie o caso e sua peça. A IA analisará falhas, sugerirá melhorias e simulará 3 ataques adversários.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Entradas")
            
            # --- NOVO COMPONENTE ---
            tipo_analise_input = gr.Radio(
                label="Qual o tipo da 'Peça' a ser analisada?",
                choices=["Análise Genérica", "Reclamatória Trabalhista", "Contestação Trabalhista"],
                value="Análise Genérica"
            )
            
            caso_input = gr.File(label="Arquivo do Caso (.pdf ou .txt)", file_types=[".pdf", ".txt"])
            peca_input = gr.File(label="Arquivo da Peça (.pdf ou .txt)", file_types=[".pdf", ".txt"])
            
            with gr.Row():
                btn_clear = gr.Button("Limpar")
                btn_analisar = gr.Button("Analisar e Simular", variant="primary")
        
        with gr.Column(scale=2):
            gr.Markdown("### 2. Análise de Falhas e Sugestões")
            output_analise = gr.Markdown(label="Análise e Sugestões")
            
            gr.Markdown("### 3. Simulações de Ataque Adversário")
            output_simulacao = gr.Markdown(label="Simulações de Ataque")
            
            output_pdf = gr.File(label="Baixar Simulação em PDF")

    # (Listas atualizadas para incluir 'tipo_analise_input')
    outputs_list = [btn_analisar, output_analise, output_simulacao, output_pdf]
    inputs_list = [tipo_analise_input, caso_input, peca_input] 
    clear_list = [tipo_analise_input, caso_input, peca_input, output_analise, output_simulacao, output_pdf]

    # (Eventos atualizados)
    
    btn_analisar.click(
        fn=analisar_e_simular_com_progresso,
        inputs=inputs_list, 
        outputs=outputs_list
    )
    
    btn_clear.click(
        fn=clear_fields,
        inputs=[],
        outputs=clear_list 
    )
    
    gr.Markdown(
        """
        ---
        **Nota:** Aguarde alguns segundos até a resposta ser gerada pela IA.
        """
    )

if __name__ == "__main__":
    demo.queue().launch()