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
import stripe
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")
STRIPE_API_KEY = os.environ.get("STRIPE_API_KEY")

# Initialize Clients
client = OpenAI(api_key=OPENAI_API_KEY)
if STRIPE_API_KEY:
    stripe.api_key = STRIPE_API_KEY

# ... (MCP Server Init)

# --- MCP TOOLS ---

@mcp.tool()
def create_stripe_checkout_session(price_id: str, success_url: str, cancel_url: str) -> str:
    """Create a Stripe Checkout Session for a given Price ID."""
    try:
        session = stripe.checkout.Session.create(
            payment_method_types=['card'],
            line_items=[{'price': price_id, 'quantity': 1}],
            mode='payment',
            success_url=success_url,
            cancel_url=cancel_url,
        )
        return f"Checkout Session created: {session.url}"
    except Exception as e:
        return f"Error creating session: {str(e)}"

# ... (Existing tools)

# --- AGENT LOGIC ---

def aussie_router(user_input, history):
    # ...
    # Updated router instructions logic (should ideally be loaded from databank)
    # Adding Stripe to the dispatch logic
    system_instr = load_from_databank("router_instructions.md")
    # ...

@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."

@mcp.tool()
def generate_image(prompt: str) -> str:
    """Generate an image using a text-to-image model on Hugging Face."""
    try:
        # Using a popular model like FLUX or SDXL via InferenceClient
        image = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-schnell")
        os.makedirs("exports/images", exist_ok=True)
        image_path = f"exports/images/{hash(prompt)}.png"
        image.save(image_path)
        return f"Image generated and saved to {image_path}"
    except Exception as e:
        return f"Error generating image: {str(e)}"

# --- 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)