import os
import gradio as gr
from fastmcp import FastMCP
from openai import OpenAI
from memory_sync import save_to_databank, load_from_databank, get_embeddings
from ebook_pipeline import create_ebook_files

# Load Environment Variables
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
HF_TOKEN = os.environ.get("HF_TOKEN")

# Initialize OpenAI Client (using GPT-4o-mini as requested)
client = OpenAI(api_key=OPENAI_API_KEY)

# Initialize MCP Server
mcp = FastMCP("Aussie Agent Hub")

# --- MCP TOOLS ---

@mcp.tool()
def generate_ebook(title: str, author: str, chapters: list) -> str:
    """Generate EPUB and PDF files from a list of chapters (title and content)."""
    epub_path, pdf_path = create_ebook_files(title, author, chapters)
    return f"Ebook generated: {epub_path}, {pdf_path}"

@mcp.tool()
def save_knowledge(module_name: str, content: str) -> str:
    """Save knowledge content to the persistent databank."""
    success = save_to_databank(f"{module_name}.md", content)
    return "Knowledge saved successfully." if success else "Failed to save knowledge."

@mcp.tool()
def query_databank(filename: str) -> str:
    """Retrieve content from the databank."""
    content = load_from_databank(filename)
    return content if content else "File not found."

# --- AGENT LOGIC (Aussie Domain Router) ---

def aussie_router(user_input, history):
    # Retrieve system instructions from databank or fallback
    system_instr = load_from_databank("router_instructions.md") or "You are the Aussie Domain Router. Orchestrate tasks for the user."
    
    messages = [{"role": "system", "content": system_instr}]
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": user_input})
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        # tools=[...] # We would list the MCP tools here if GPT-4o-mini supported native MCP schema directly
    )
    
    return response.choices[0].message.content

# --- GRADIO UI ---

with gr.Blocks(title="Aussie MCP Agent Hub") as demo:
    gr.Markdown("# 🐨 Aussie MCP Server Agent Hub")
    
    with gr.Tab("Chat with Hub"):
        chatbot = gr.Chatbot()
        msg = gr.Textbox(placeholder="Ask your Aussie Agent anything...")
        clear = gr.Button("Clear")

        def user(user_message, history):
            return "", history + [[user_message, None]]

        def bot(history):
            user_message = history[-1][0]
            bot_message = aussie_router(user_message, history[:-1])
            history[-1][1] = bot_message
            return history

        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
            bot, chatbot, chatbot
        )
        clear.click(lambda: None, None, chatbot, queue=False)

    with gr.Tab("Databank"):
        gr.Markdown("View and manage your persistent knowledge modules.")
        # Add interface elements to list and view databank files

# Start the application with MCP support
if __name__ == "__main__":
    # Gradio 5+ with MCP SSE endpoint
    # Note: FastMCP usually runs its own server, here we integrate it with Gradio or run in parallel
    import threading
    
    def run_mcp():
        mcp.run(transport="sse", host="0.0.0.0", port=7861) # Running MCP on a separate port or integrating with Gradio path

    # t = threading.Thread(target=run_mcp)
    # t.start()
    
    demo.launch(server_name="0.0.0.0", server_port=7860)