import os
import gradio as gr
from huggingface_hub import InferenceClient
from datetime import datetime
import logging
from typing import Dict, List, Optional
from pathlib import Path
import json
import dotenv

# Carrega variáveis de ambiente do arquivo .env se existir
dotenv.load_dotenv()

# Configuração do logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def load_token() -> str:
    """Carrega o token HF de várias fontes possíveis"""
    # Prioridade: 1. Variável de ambiente, 2. Arquivo .env, 3. arquivo config.json
    
    # Verifica variável de ambiente
    token = os.getenv('HF_TOKEN')
    if token:
        return token
        
    # Verifica arquivo config.json
    config_file = Path('config.json')
    if config_file.exists():
        try:
            with open(config_file, 'r') as f:
                config = json.load(f)
                return config.get('hf_token', '')
        except Exception as e:
            logger.error(f"Erro ao ler config.json: {e}")
    
    return ''

# Carrega o token
HF_API_TOKEN = load_token()

class DocumentGenerator:
    """Gerencia a geração de documentos usando HF Inference API"""
    
    def __init__(self, api_key: str):
        if not api_key:
            raise ValueError("API Key não encontrada. Configure HF_TOKEN nas variáveis de ambiente ou no arquivo config.json")
        
        self.client = InferenceClient(api_key=api_key)
        self.model = "mistralai/Mistral-7B-Instruct-v0.2"
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento usando o modelo"""
        try:
            # Prepara as mensagens para o chat
            messages = [
                {
                    "role": "system",
                    "content": """Você é um advogado criminalista brasileiro experiente.
                    Gere documentos jurídicos formais no formato do direito brasileiro."""
                },
                {
                    "role": "user",
                    "content": f"""Gere um {doc_type} completo com os seguintes dados:
                    
                    QUALIFICAÇÃO:
                    Cliente: {context.get('client_name')}
                    Processo: {context.get('process_number')}
                    Tribunal: {context.get('court')}
                    Comarca: {context.get('jurisdiction')}
                    
                    FATOS:
                    {context.get('facts')}
                    
                    FUNDAMENTOS JURÍDICOS:
                    {context.get('legal_basis')}
                    
                    Use formato jurídico adequado, com todas as partes necessárias."""
                }
            ]
            
            # Faz a chamada à API
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.3,
                top_p=0.85,
                max_tokens=2048,
            )
            
            return self._format_output(completion.choices[0].message.content)
            
        except Exception as e:
            logger.error(f"Erro na geração: {str(e)}")
            return f"Erro na geração do documento: {str(e)}"
    
    def _format_output(self, text: str) -> str:
        """Formata o texto gerado"""
        if not text:
            return "Erro: Nenhum texto gerado"
        return text.strip()

class WebInterface:
    """Interface Gradio para o gerador de documentos"""
    
    def __init__(self):
        self.generator = DocumentGenerator(api_key=HF_API_TOKEN)
        self.create_interface()
    
    def create_interface(self):
        """Cria a interface web com Gradio"""
        
        with gr.Blocks(theme=gr.themes.Soft()) as self.app:
            gr.Markdown("""
            # Criminal.ai - Gerador de Documentos Jurídicos
            Sistema de geração de peças processuais criminais
            """)
            
            with gr.Row():
                with gr.Column():
                    doc_type = gr.Dropdown(
                        choices=[
                            "Habeas Corpus",
                            "Denúncia Criminal",
                            "Alegações Finais",
                            "Resposta à Acusação",
                            "Recurso em Sentido Estrito",
                            "Apelação Criminal"
                        ],
                        label="Tipo de Documento",
                        value="Habeas Corpus"
                    )
                    
                    with gr.Group():
                        gr.Markdown("### Informações do Processo")
                        client_name = gr.Textbox(
                            label="Nome do Cliente",
                            placeholder="Nome completo"
                        )
                        process_number = gr.Textbox(
                            label="Número do Processo",
                            placeholder="NNNNNNN-NN.NNNN.N.NN.NNNN"
                        )
                        court = gr.Textbox(
                            label="Tribunal",
                            value="TRIBUNAL DE JUSTIÇA DO ESTADO"
                        )
                        jurisdiction = gr.Textbox(
                            label="Comarca"
                        )
                    
                    with gr.Group():
                        gr.Markdown("### Detalhes do Caso")
                        facts = gr.Textbox(
                            label="Fatos",
                            lines=5,
                            placeholder="Descreva os fatos..."
                        )
                        legal_basis = gr.Textbox(
                            label="Fundamentos Jurídicos",
                            lines=3
                        )
                    
                    generate_btn = gr.Button("Gerar Documento", variant="primary")
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Documento Gerado",
                        lines=30,
                        show_copy_button=True
                    )
                    status = gr.Textbox(label="Status")
            
            # Eventos
            generate_btn.click(
                fn=self._generate_document,
                inputs=[
                    doc_type, client_name, process_number,
                    court, jurisdiction, facts, legal_basis
                ],
                outputs=[output, status]
            )
    
    def _generate_document(
        self, doc_type: str, client_name: str,
        process_number: str, court: str,
        jurisdiction: str, facts: str,
        legal_basis: str
    ) -> tuple:
        """Gera o documento com os parâmetros fornecidos"""
        try:
            if not all([client_name, process_number, facts, legal_basis]):
                return "Erro: Preencha todos os campos obrigatórios", "⚠️ Campos incompletos"
            
            context = {
                "client_name": client_name,
                "process_number": process_number,
                "court": court,
                "jurisdiction": jurisdiction,
                "facts": facts,
                "legal_basis": legal_basis
            }
            
            result = self.generator.generate(doc_type, context)
            return result, "✅ Documento gerado com sucesso"
            
        except Exception as e:
            logger.error(f"Erro: {str(e)}")
            return "", f"❌ Erro: {str(e)}"
    
    def launch(self):
        """Inicia a interface web"""
        self.app.launch(share=True)

if __name__ == "__main__":
    # Verifica se o token está configurado
    if not HF_API_TOKEN:
        print("\n⚠️  Token HF não encontrado!")
        print("\nPara configurar o token:")
        print("1. Crie um arquivo .env com:")
        print("   HF_TOKEN=seu_token_aqui")
        print("\nOu")
        print("2. Crie um arquivo config.json com:")
        print('   {"hf_token": "seu_token_aqui"}')
        print("\nOu")
        print("3. Configure a variável de ambiente HF_TOKEN")
        exit(1)
    
    interface = WebInterface()
    interface.launch()