import os
import gradio as gr
from huggingface_hub import InferenceClient
import json

HF_TOKEN = os.getenv("HF_TOKEN")

client = InferenceClient(
    model="google/gemma-2-9b-it",
    token=HF_TOKEN
)

# Seu AGENTE_JSON completo
AGENTE_JSON = """
{
  "nome_agente": "IndenizaAI",
  "objetivo": "Calcular indenizações trabalhistas automaticamente usando TODOS os dados fornecidos pelo usuário, sem solicitar informações adicionais quando já estiverem presentes.",
  "regras_gerais": [
    "Nunca pedir informações extras se os dados já incluem idade, renda, data do acidente, sequelas, custos médicos, tempo de afastamento e dados contratuais.",
    "Sempre realizar TODOS os cálculos com base SOMENTE nos dados enviados pelo usuário.",
    "Se algum dado realmente estiver faltando, pedir apenas o dado específico e uma única vez.",
    "Nunca repetir perguntas já respondidas pelo usuário.",
    "Nunca devolver a responsabilidade do cálculo ao usuário."
  ],
  "calculos_obrigatorios": [
    "Cálculo de pensão mensal (sem projetar valor vitalício total)",
    "Cálculo de danos morais",
    "Cálculo de danos estéticos",
    "Cálculo de danos materiais",
    "Cálculo de lucros cessantes",
    "Tabela final com todos os valores calculados",
    "Total geral da indenização"
  ],
  "comportamento_resposta": {
    "estilo": "técnico, claro e organizado",
    "estrutura": [
      "👉 Resumo inicial do caso",
      "👉 Tabela detalhada de cada cálculo",
      "👉 Total geral",
      "👉 Fundamentação jurídica resumida"
    ],
    "proibicoes": [
      "Não pedir dados que já estão no texto do usuário",
      "Não afirmar que faltam informações quando elas já foram fornecidas",
      "Não delegar cálculos ao usuário"
    ]
  }
}
"""

# Campos que queremos armazenar do usuário
CAMPOS_NECESSARIOS = [
    "nome_trabalhador", "idade", "cargo", "tempo_servico",
    "renda_mensal", "contrato_trabalho",
    "data_acidente", "tipo_acidente", "descricao_acidente", "tempo_afastamento",
    "sequelas", "custos_medicos", "perda_capacidade",
    "gastos_tratamento", "perda_rendimentos"
]

# Armazenamento em memória (pode ser substituído por DB)
usuarios_dados = {}

def montar_mensagem_usuario(msg, dados_completos):
    return (
        "Você é o IndenizaAI.\n"
        "Siga integralmente todas as instruções do JSON abaixo:\n"
        "### CONFIGURAÇÃO DO AGENTE ###\n"
        + AGENTE_JSON +
        "\n### DADOS COLETADOS DO USUÁRIO ###\n"
        + json.dumps(dados_completos, indent=2) +
        "\n### MENSAGEM DO USUÁRIO ###\n"
        + msg +
        "\n### INSTRUÇÃO CRÍTICA ###\n"
        "- Se o usuário fornecer dados suficientes, execute TODOS os cálculos automaticamente.\n"
        "- Se faltar algum dado, peça apenas o que estiver faltando.\n"
        "- Sempre incluir cálculos de pensão vitalícia, danos morais, estéticos e materiais.\n"
        "- Sempre entregar uma tabela final e o total geral da indenização.\n"
        "Responda como **IndenizaAI**."
    )

def verificar_dados_completos(dados):
    faltando = [campo for campo in CAMPOS_NECESSARIOS if campo not in dados or not dados[campo]]
    return faltando

def responder(mensagem, historico, usuario_id="default"):
    # Recupera dados armazenados do usuário
    dados_usuario = usuarios_dados.get(usuario_id, {})

    # Tenta atualizar dados a partir de JSON enviado pelo usuário
    try:
        dados_parciais = json.loads(mensagem)
        dados_usuario.update(dados_parciais)
    except:
        pass  # mensagem normal, sem dados estruturados

    usuarios_dados[usuario_id] = dados_usuario

    faltando = verificar_dados_completos(dados_usuario)

    if faltando:
        resposta = f"⚠️ Ainda faltam os seguintes dados para realizar o cálculo: {', '.join(faltando)}"
    else:
        prompt = montar_mensagem_usuario(mensagem, dados_usuario)
        messages = [
            {"role": "system", "content": "Você é o IndenizaAI, especialista em cálculos de indenizações trabalhistas."},
            {"role": "user", "content": prompt}
        ]
        result = client.chat_completion(
            messages=messages,
            max_tokens=1200,
            temperature=0.1,
            top_p=0.9
        )
        resposta = result.choices[0].message["content"]

    historico.append({"role": "user", "content": mensagem})
    historico.append({"role": "assistant", "content": resposta})
    return historico, ""

with gr.Blocks() as demo:
    gr.Markdown("## ⚖️ IndenizaAI — Assistente Jurídico Especializado")
    chat = gr.Chatbot(type="messages")
    inp = gr.Textbox(label="Digite sua mensagem")
    inp.submit(responder, [inp, chat], [chat, inp])

demo.launch()