import os
import torch
import gradio as gr
from transformers import pipeline, AutoTokenizer
from langchain_community.llms import HuggingFacePipeline
from langchain.agents import AgentExecutor, create_react_agent
from langchain.tools import BaseTool
from langchain_core.prompts import PromptTemplate
from langchain import hub

# 1. Configuración del Modelo DeepSeek
model_id = 'deepseek-ai/deepseek-coder-6.7b-instruct'
tokenizer = AutoTokenizer.from_pretrained(model_id)

llm_pipeline = pipeline(
    'text-generation',
    model=model_id,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    device_map='auto',
    max_new_tokens=512
)

llm = HuggingFacePipeline(pipeline=llm_pipeline)

# 2. Definición de Herramientas (Tools)
class CalculatorTool(BaseTool):
    name = 'calculator'
    description = 'Useful for math. Provide only the expression.'
    def _run(self, expression: str) -> str:
        try: 
            return str(eval(expression))
        except: 
            return 'Error.'

tools = [CalculatorTool()]

# 3. Descarga del Prompt con Plan B Robusto Integrado
try:
    # Intenta traer la plantilla oficial de la comunidad para agentes ReAct
    prompt = hub.pull('hwchase17/react')
except Exception:
    print("⚠️ No se pudo conectar al Hub de LangChain. Usando plantilla ReAct local de respaldo...")
    # Estructura estricta exigida por create_react_agent para no lanzar ValueError
    template = """Answer the following questions as best you can. You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought: {agent_scratchpad}"""

    prompt = PromptTemplate(
        template=template,
        input_variables=["input", "tools", "tool_names", "agent_scratchpad"]
    )

# 4. Inicialización del Agente y su Ejecutor
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)

# 5. Función de Predicción para la Interfaz
def predict(message, history):
    res = agent_executor.invoke({'input': message})
    return res['output']

# 6. Interfaz Gráfica con Gradio 6.x (Parámetro theme movido a launch)
with gr.Blocks() as demo:
    gr.Markdown('# 🧠 DeepSeek Agent Terminal')
    gr.Markdown('Agente inteligente basado en arquitectura ReAct para resolución de problemas lógicos y matemáticos.')
    gr.ChatInterface(fn=predict)

if __name__ == '__main__':
    # Lanzamiento nativo optimizado para Hugging Face Spaces
    demo.launch(
        server_name='0.0.0.0', 
        server_port=7860,
        theme=gr.themes.Soft(primary_hue='green'),
        show_error=True
    )