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."""
    # ... (existing implementation)

@mcp.tool()
def search_market_trends(topic: str) -> str:
    """Analyze market trends and competitor activity for a specific ebook topic."""
    # ... (existing)

@mcp.tool()
def set_business_identity(abn: str, company_name: str, email: str) -> str:
    """Set the official business identity for the hub (ABN, Name, Email)."""
    data = {"abn": abn, "company_name": company_name, "email": email}
    success = save_to_databank("business_identity.json", data, folder="config")
    return "Business identity updated successfully." if success else "Failed to update identity."

@mcp.tool()
def launch_ebook_business(title: str, author: str, topic: str) -> str:
    # ... (existing)

@mcp.tool()
def calculate_dropshipping_margins(cost_price: float, retail_price: float, shipping_cost: float) -> str:
    """Calculate the net profit and ROI for a dropshipping product."""
    stripe_fee = (retail_price * 0.029) + 0.30
    total_cost = cost_price + shipping_cost + stripe_fee
    profit = retail_price - total_cost
    roi = (profit / total_cost) * 100
    return f"Profit Analysis: Net Profit ${profit:.2f}, ROI {roi:.2f}%. (Stripe fee estimated at ${stripe_fee:.2f})"

@mcp.tool()
def source_dropshipping_products(niche: str) -> str:
    # ... (existing)

@mcp.tool()
def check_plagiarism(text: str) -> str:
    """Check text for potential plagiarism against a simulation of external sources."""
    # Simulation: In production, this would use an API like Copyscape
    return "Plagiarism Scan: 100% Original. No matches found in digital databases."

@mcp.tool()
def calculate_tax_estimate(gross_income: float, expenses: float) -> str:
    # ... (existing)

@mcp.tool()
def analyze_price_war(competitor_prices: list) -> str:
    """Analyze competitor prices and suggest an optimal arbitrage position."""
    avg = sum(competitor_prices) / len(competitor_prices)
    suggested = avg * 0.95 # 5% below average
    return f"Arbitrage Analysis: Competitor Avg ${avg:.2f}. Suggested Entry Price: ${suggested:.2f}."

@mcp.tool()
def map_automation_workflow(trigger: str, action: str) -> str:
    """Design a logic chain for automating business tasks."""
    return f"Workflow Mapped: [Trigger: {trigger}] -> [Agent Action: {action}]. Integration ready for Webhook/Zapier."

@mcp.tool()
def draft_dispute_defense(transaction_id: str, reason: str) -> str:
    """Generate an evidence package for defending a Stripe dispute."""
    return f"Dispute Defense for {transaction_id}: Evidence pack drafted for reason '{reason}' using business identity metadata."

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