# app.py
import gradio as gr
import os
import httpx
from huggingface_hub import list_models
from datasets import load_dataset

HF_TOKEN = os.getenv("HF_TOKEN")

def search_models(query="legal", language="pt", limit=5):
    try:
        results = list(
            list_models(
                search=query,
                language=language or None,
                sort="downloads",
                direction=-1,
                limit=int(limit),
                token=HF_TOKEN,
            )
        )
        output = ""
        for m in results[:int(limit)]:
            output += f"**{m.modelId}**\n"
            output += f"- Task: {m.pipeline_tag}\n"
            output += f"- Downloads: {m.downloads:,}\n"
            output += f"- URL: https://huggingface.co/{m.modelId}\n\n"
        return output if output else "Nenhum modelo encontrado."
    except Exception as e:
        return f"Erro: {str(e)}"

def analyze_text(text, task="summarization"):
    try:
        models = {
            "summarization": "facebook/bart-large-cnn",
            "text-classification": "nlpaueb/legal-bert-base-uncased",
        }
        model_id = models.get(task, "facebook/bart-large-cnn")
        url = f"https://api-inference.huggingface.co/models/{model_id}"
        headers = {"Content-Type": "application/json"}
        if HF_TOKEN:
            headers["Authorization"] = f"Bearer {HF_TOKEN}"
        payload = {"inputs": text[:512]}
        with httpx.Client(timeout=60.0) as client:
            resp = client.post(url, headers=headers, json=payload)
        if resp.status_code == 200:
            return str(resp.json())
        else:
            return f"Erro HTTP {resp.status_code}: {resp.text[:200]}"
    except Exception as e:
        return f"Erro: {str(e)}"

def explore_dataset(dataset_id, n_samples=3):
    try:
        ds = load_dataset(
            dataset_id,
            split=f"train[:{n_samples}]",
            token=HF_TOKEN,
            trust_remote_code=False,
        )
        output = f"**Dataset:** {dataset_id}\n\n"
        output += f"**Colunas:** {', '.join(ds.column_names)}\n\n"
        for i, sample in enumerate(ds.to_list()[:int(n_samples)]):
            output += f"--- Amostra {i+1} ---\n"
            for key, val in sample.items():
                if isinstance(val, str) and len(val) > 300:
                    val = val[:300] + "..."
                output += f"{key}: {val}\n"
            output += "\n"
        return output
    except Exception as e:
        return f"Erro: {str(e)}"

with gr.Blocks(title="LEX - Assistente Jurídico", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ⚖️ LEX - Assistente Jurídico MCP")
    gr.Markdown("Powered by Hugging Face Hub + FastMCP")
    
    with gr.Tab("🔍 Buscar Modelos"):
        with gr.Row():
            with gr.Column():
                model_query = gr.Textbox(label="Query", value="legal")
                model_lang = gr.Textbox(label="Idioma", value="pt")
                model_limit = gr.Slider(1, 20, value=5, step=1, label="Limite")
                model_btn = gr.Button("Buscar", variant="primary")
            with gr.Column():
                model_output = gr.Textbox(label="Resultados", lines=15, interactive=False)
        model_btn.click(fn=search_models, inputs=[model_query, model_lang, model_limit], outputs=model_output)
    
    with gr.Tab("📝 Analisar Texto"):
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(label="Texto Jurídico", lines=5)
                task_type = gr.Dropdown(
                    choices=["summarization", "text-classification"],
                    value="summarization",
                    label="Tarefa"
                )
                analyze_btn = gr.Button("Analisar", variant="primary")
            with gr.Column():
                analyze_output = gr.Textbox(label="Resultado", lines=10, interactive=False)
        analyze_btn.click(fn=analyze_text, inputs=[text_input, task_type], outputs=analyze_output)
    
    with gr.Tab("📊 Explorar Dataset"):
        with gr.Row():
            with gr.Column():
                dataset_input = gr.Textbox(label="Dataset ID", value="joelniklaus/brazilian_court_decisions")
                sample_count = gr.Slider(1, 10, value=3, step=1, label="Amostras")
                dataset_btn = gr.Button("Explorar", variant="primary")
            with gr.Column():
                dataset_output = gr.Textbox(label="Dataset Info", lines=15, interactive=False)
        dataset_btn.click(fn=explore_dataset, inputs=[dataset_input, sample_count], outputs=dataset_output)

if __name__ == "__main__":
    # ✅ CORREÇÃO: Configuração correta para HF Spaces
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        share=False,
        show_error=True
    )