import gradio as gr
import os
import torch

from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_huggingface import HuggingFacePipeline
from langchain.memory import ConversationBufferMemory

from transformers import AutoTokenizer, pipeline

# ===================================================================
# CONFIGURAÇÃO DE COMBATE
# ===================================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TORCH_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
LLM_MODELS = {
    "TinyLlama-1.1B-Chat": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "Phi-2": "microsoft/phi-2",
    "Mistral-7B-Instruct": "mistralai/Mistral-7B-Instruct-v0.2",
    "Zephyr-7B-Beta": "HuggingFaceH4/zephyr-7b-beta"
}

# ===================================================================
# NÚCLEO DE OPERAÇÕES ESPECIAIS
# ===================================================================
class TacticalDocumentProcessor:
    @staticmethod
    def neutralize_documents(files, chunk_size=512, chunk_overlap=64):
        """Operação de desmantelamento de documentos hostis"""
        if not files:
            raise ValueError("ALVO NÃO IDENTIFICADO")
        
        try:
            loaders = [PyPDFLoader(file.name) for file in files]
            splitter = RecursiveCharacterTextSplitter(
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
                separators=["\n\n", "\n", "\. ", " ", ""]
            )
            return [page for loader in loaders for page in loader.load_and_split(splitter)]
        except Exception as e:
            raise RuntimeError(f"FALHA NA OPERAÇÃO: {str(e)}")

class VectorStrikeSystem:
    @staticmethod
    def deploy_vector_db(splits):
        """Implante imediato de sistema de vetorização"""
        if not splits:
            raise ValueError("NENHUMA INTELECÇÃO DISPONÍVEL")
        
        return Chroma.from_documents(
            documents=splits,
            embedding=HuggingFaceEmbeddings(),
            persist_directory="./combat_db"
        )

class LLMWeaponsSystem:
    @staticmethod
    def activate_weapon(model_name, temp=0.7, max_tokens=512):
        """Ativação de armamento cognitivo"""
        try:
            tokenizer = AutoTokenizer.from_pretrained(LLM_MODELS[model_name])
            
            pipe = pipeline(
                "text-generation",
                model=LLM_MODELS[model_name],
                tokenizer=tokenizer,
                device=DEVICE,
                torch_dtype=TORCH_DTYPE,
                max_new_tokens=max_tokens,
                do_sample=True,
                top_k=50,
                temperature=temp,
                model_kwargs={"load_in_4bit": True} if "cuda" in DEVICE else {}
            )
            
            return HuggingFacePipeline(pipeline=pipe)
        except KeyError:
            raise ValueError("ARMA NÃO CATALOGADA")
        except Exception as e:
            raise RuntimeError(f"FALHA NO SISTEMA DE ARMAMENTO: {str(e)}")

# ===================================================================
# INTERFACE DE COMBATE
# ===================================================================
def deploy_combat_interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="🔥 WARBOT v2.0") as interface:
        state = gr.State({
            "db": None,
            "llm": None,
            "doc_status": False,
            "model_status": False
        })
        
        # Zona de Controle Tático
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                file_upload = gr.Files(label="CARREGAMENTO DE ALVOS", file_types=[".pdf"])
                process_btn = gr.Button("INICIAR PROCESSAMENTO", variant="stop")
                process_log = gr.Textbox(label="RELATÓRIO DE PROCESSAMENTO", interactive=False)
            
            with gr.Column(scale=1):
                model_selector = gr.Dropdown(list(LLM_MODELS.keys()), label="SELECIONE O ARMAMENTO", value="TinyLlama-1.1B-Chat")
                temp_slider = gr.Slider(0, 1, 0.7, label="NÍVEL DE AGRESSIVIDADE")
                deploy_btn = gr.Button("ATIVAR ARMAMENTO", variant="primary")
                deploy_log = gr.Textbox(label="STATUS DO ARMAMENTO", interactive=False)
        
        # Campo de Batalha Principal
        chatbot = gr.Chatbot(height=650, label="ZONA DE ENGENHARIA COGNITIVA")
        msg_input = gr.Textbox(label="INSIRA COMANDO DE ATAQUE", placeholder="Aguardando ordens...")
        clear_btn = gr.Button("LIMPAR CAMPO DE BATALHA")

        # ===== OPERAÇÕES TÁTICAS =====
        @process_btn.click(inputs=[file_upload], outputs=[state, process_log])
        def execute_processing(files):
            try:
                splits = TacticalDocumentProcessor.neutralize_documents(files)
                db = VectorStrikeSystem.deploy_vector_db(splits)
                return {
                    "db": db,
                    "llm": None,
                    "doc_status": True,
                    "model_status": False
                }, "✅ ALVOS PROCESSADOS COM SUCESSO"
            except Exception as e:
                return state.value, f"☠️ FALHA CRÍTICA: {str(e)}"

        @deploy_btn.click(inputs=[model_selector, temp_slider, state], outputs=[state, deploy_log])
        def deploy_weapon(model, temp, current_state):
            try:
                if not current_state["doc_status"]:
                    raise RuntimeError("ALVOS NÃO PROCESSADOS! EXECUTE A FASE 1")
                
                llm = LLMWeaponsSystem.activate_weapon(model, temp)
                current_state["llm"] = ConversationalRetrievalChain.from_llm(
                    llm=llm,
                    retriever=current_state["db"].as_retriever(search_kwargs={"k": 3}),
                    memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True),
                    return_source_documents=True
                )
                current_state["model_status"] = True
                return current_state, f"🚀 {model} PRONTO PARA ENGAGEMENT"
            except Exception as e:
                return current_state, f"💥 FALHA NO ARMAMENTO: {str(e)}"

        @msg_input.submit(inputs=[msg_input, chatbot, state], outputs=[msg_input, chatbot])
        def execute_engagement(command, history, state):
            if not state["model_status"]:
                return command, history + [(command, "⚠️ ARMAMENTO NÃO ATIVADO")]
            
            try:
                result = state["llm"]({"question": command, "chat_history": history})
                intel_report = "\n".join(
                    f"🔍 Pg {doc.metadata['page']+1}: {doc.page_content[:100]}..." 
                    for doc in result["source_documents"][:3]
                )
                return "", history + [
                    (command, f"🎯 RESPOSTA:\n{result['answer']}\n\n📡 INTELIGÊNCIA:\n{intel_report}")
                ]
            except Exception as e:
                return command, history + [(command, f"☢️ FALHA OPERACIONAL: {str(e)}")]

        @clear_btn.click(inputs=[], outputs=[chatbot])
        def clear_battlefield():
            return []

    return interface

# ===================================================================
# INICIALIZAÇÃO DO SISTEMA
# ===================================================================
if __name__ == "__main__":
    combat_system = deploy_combat_interface()
    combat_system.launch(
        server_name="0.0.0.0",
        server_port=7860,
        auth=("commander", "tactical123"),
        share=True
    )