# app.py
import os
from fastapi import FastAPI
from dotenv import load_dotenv
from huggingface_hub.inference._mcp.agent import Agent
import gradio as gr
import uvicorn
from fastapi.responses import RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional, Literal

load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
HF_MODEL = os.getenv("HF_MODEL", "google/gemma-2-2b")

app = FastAPI(title="Model Card Chatbot")
app.add_middleware(
    CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
)

agent_instance: Optional[Agent] = None
DEFAULT_PROVIDER: Literal['hf-inference'] = "hf-inference"

async def get_agent():
    global agent_instance
    if agent_instance is None and HF_TOKEN:
        print("🔧 Creating Agent...")
        try:
            agent = Agent(
                model=HF_MODEL,
                provider=DEFAULT_PROVIDER,
                api_key=HF_TOKEN,
                servers=[{
                    "type": "stdio",
                    "config": {
                        "command": "python",
                        "args": ["mcp_server.py"],
                        "cwd": ".",
                        "env": {"HF_TOKEN": HF_TOKEN}
                    }
                }]
            )
            await agent.load_tools()
            agent_instance = agent
            print("✅ Agent is ready!")
        except Exception as e:
            print(f"❌ Failed to create agent: {str(e)}")
    return agent_instance

@app.on_event("startup")
async def startup_event():
    await get_agent()

def chat_function(user_message, history, model_id):
    global agent_instance
    prompt = f"""
You're an assistant helping with Hugging Face model cards.
First, run the tool `read_model_card` on repo_id `{model_id}` to get the model card.
Then answer this user question based on the model card:
User question: {user_message}
"""
    try:
        response = ""
        outputs = agent_instance.run(prompt)
        for output in outputs:
            if hasattr(output, "content") and output.content:
                response = output.content
        if not response:
            response = "⚠️ Sorry, no answer generated."
        # Append as new OpenAI-style message
        history = history + [{"role": "user", "content": user_message}]
        history.append({"role": "assistant", "content": response})
    except Exception as e:
        history = history + [{"role": "user", "content": user_message}]
        history.append({"role": "assistant", "content": f"⚠️ Error: {str(e)}"})
    return history, ""

def create_gradio_app():
    with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), title="🤖 Model Card Chatbot") as demo:
        gr.Markdown("""
        # 🤖 **Model Card Chatbot**
        _Ask anything about a Hugging Face model's card_
        """)
        chatbot = gr.Chatbot(label="Chat", type="messages", show_copy_button=True, height=400)
        with gr.Row():
            model_id = gr.Textbox(label="Model ID", value="google/gemma-2-2b", scale=2)
            user_input = gr.Textbox(
                label="Your Question", 
                placeholder="e.g., What dataset was it trained on?", 
                scale=4
            )
            ask_button = gr.Button("🔍 Ask", scale=1)
        
        ask_button.click(
            fn=chat_function,
            inputs=[user_input, chatbot, model_id],
            outputs=[chatbot, user_input]
        )
    return demo

gradio_app = create_gradio_app()
app = gr.mount_gradio_app(app, gradio_app, path="/")

@app.get("/")
async def root():
    return RedirectResponse("/")

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)