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

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

# --- LLM TOOL WORKER (The Intelligence Engine) ---

def llm_worker(prompt, system_prompt="You are a specialized business assistant."):
    """Helper to route tool intelligence through OpenAI or Free Fallback."""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    try:
        # Try Primary Intelligence (OpenAI)
        response = client.chat.completions.create(model="gpt-4o-mini", messages=messages)
        return response.choices[0].message.content
    except Exception:
        # Fallback to Free Intelligence (Hugging Face Llama 3.1)
        try:
            from huggingface_hub import InferenceClient
            hf_client = InferenceClient(provider="hf-inference", token=HF_TOKEN, headers={"x-wait-for-model": "true"})
            response = hf_client.chat_completion(model="meta-llama/Meta-Llama-3.1-8B-Instruct", messages=messages, max_tokens=1500)
            return response.choices[0].message.content
        except Exception as e:
            return f"Intelligence Error: {str(e)}"

# --- REAL MCP TOOLS ---

@mcp.tool()
def create_stripe_checkout_session(price_id: str, success_url: str, cancel_url: str) -> str:
    """Create a real 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)}"

@mcp.tool()
def create_stripe_product_with_price(name: str, description: str, unit_amount_cents: int, currency: str = "aud") -> str:
    """Create a real Product and Price in Stripe."""
    try:
        product = stripe.Product.create(name=name, description=description)
        price = stripe.Price.create(product=product.id, unit_amount=unit_amount_cents, currency=currency)
        return f"Product Created: {name} (ID: {product.id}). Price Created (ID: {price.id}) for {unit_amount_cents/100:.2f} {currency.upper()}."
    except Exception as e:
        return f"Error creating Stripe product: {str(e)}"

@mcp.tool()
def generate_ebook(title: str, author: str, chapters: list) -> str:
    """Generate professional EPUB and PDF files with branded metadata."""
    epub_path, pdf_path = create_ebook_files(title, author, chapters)
    return f"Ebook generated successfully: {epub_path}, {pdf_path}"

@mcp.tool()
def generate_image(prompt: str) -> str:
    """Generate a branded image with multiple free fallbacks."""
    try:
        from gradio_client import Client
        import shutil
        business_name = os.environ.get("BUSINESS_NAME", "Fair Dinkum Publishing")
        owner = os.environ.get("BUSINESS_OWNER", "BRETT SJOBERG")
        brand_context = f"Professional brand asset for {business_name} (Owner: {owner}). Style: Modern, clean, high-quality. "
        full_prompt = brand_context + prompt
        
        try:
            client = Client("mrfakename/Z-Image-Turbo", token=HF_TOKEN)
            result = client.predict(prompt=full_prompt, height=1024, width=1024, num_inference_steps=9, seed=42, randomize_seed=True, api_name="/generate_image")
            temp_image_path = result[0] if isinstance(result, (list, tuple)) else result
            model_used = "Z-Image-Turbo"
        except Exception:
            client = Client("black-forest-labs/FLUX.1-schnell", token=HF_TOKEN)
            result = client.predict(prompt=full_prompt, seed=0, randomize_seed=True, width=1024, height=1024, num_inference_steps=4, api_name="/infer")
            temp_image_path = result[0] if isinstance(result, (list, tuple)) else result
            model_used = "FLUX.1-schnell"
        
        os.makedirs("exports/images", exist_ok=True)
        final_path = f"exports/images/{abs(hash(prompt))}.png"
        shutil.copy(temp_image_path, final_path)
        return f"Branded Image Generated using {model_used}: {final_path}"
    except Exception as e:
        return f"Image Error: {str(e)}"

@mcp.tool()
def search_market_trends(topic: str) -> str:
    """Deeply analyze market trends, competition, and pricing for any niche."""
    prompt = f"Conduct a professional market research analysis for the niche: '{topic}'. Suggest a pricing strategy and identify potential competitors."
    return llm_worker(prompt, system_prompt="You are an expert Ebook and Dropshipping Market Analyst.")

@mcp.tool()
def source_dropshipping_products(niche: str) -> str:
    """Sourcing high-demand products for a dropshipping niche."""
    prompt = f"Find and describe 3 high-demand, high-margin dropshipping products for the niche: '{niche}'. Include estimated cost and retail price."
    return llm_worker(prompt, system_prompt="You are an expert E-commerce Sourcing Agent.")

@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 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:
    """Automated sequence for ebook business generation."""
    chapters = [{"title": "Introduction", "content": f"A guide to {topic}."}]
    epub_path, pdf_path = create_ebook_files(title, author, chapters, base_name=title.lower().replace(" ", "_"))
    return f"Business Launched: '{title}' created. Files: {epub_path}, {pdf_path}. Ready for launch."

@mcp.tool()
def audit_store_cro(url: str = "Preview Mode") -> str:
    """Audit a storefront for Conversion Rate Optimization (CRO) and speed."""
    prompt = f"Perform a detailed CRO and user experience audit for the storefront: {url}. Suggest 3 actionable improvements."
    return llm_worker(prompt, system_prompt="You are a Conversion Rate Optimization Expert.")

@mcp.tool()
def generate_store_layout(niche: str, store_type: str = "Dropshipping") -> str:
    """Generate a high-conversion store layout/wireframe draft."""
    prompt = f"Design a high-conversion store layout for a {store_type} business in the '{niche}' niche. Include sections for hero, social proof, and product grids."
    return llm_worker(prompt, system_prompt="You are an E-commerce Store Architect.")

@mcp.tool()
def post_to_business_platforms(title: str, content: str, platforms: list) -> str:
    """Distribute blog content to popular business platforms."""
    # Simulation: Log the distribution
    return f"Multi-Platform Distribution: '{title}' posted to {', '.join(platforms)}."

@mcp.tool()
def check_plagiarism(text: str) -> str:
    """Audit content for original integrity and potential copyright issues."""
    prompt = f"Perform a deep plagiarism and original integrity audit on the following text. Highlight any sections that seem derivative: \n\n{text}"
    return llm_worker(prompt, system_prompt="You are a professional Content Auditor and Plagiarism Specialist.")

@mcp.tool()
def map_automation_workflow(trigger: str, action: str) -> str:
    """Design a technical logic chain for cross-platform business automation."""
    prompt = f"Design a robust automation workflow for the following: [Trigger: {trigger}] -> [Action: {action}]. Provide technical steps for Zapier or Make.com."
    return llm_worker(prompt, system_prompt="You are a Senior Workflow Integration Architect.")

@mcp.tool()
def draft_dispute_defense(transaction_id: str, reason: str) -> str:
    """Draft a professional, evidence-backed defense package for a payment dispute."""
    prompt = f"Draft a professional response to a Stripe dispute. Transaction ID: {transaction_id}, Reason: {reason}. Use business identity Fair Dinkum Publishing."
    return llm_worker(prompt, system_prompt="You are a Risk Mitigation and Dispute Specialist.")

@mcp.tool()
def generate_personalized_response(customer_name: str, issue: str) -> str:
    """Create an empathetic, helpful Aussie-style support response."""
    prompt = f"Write a helpful, witty, and empathetic Aussie customer support response for {customer_name} who is experiencing: '{issue}'."
    return llm_worker(prompt, system_prompt="You are a Fair Dinkum Customer Success Agent.")

@mcp.tool()
def create_blogger_post(title: str, topic: str) -> str:
    """Draft a full, SEO-optimized blog post for Fair Dinkum Publishing."""
    prompt = f"Draft a comprehensive, SEO-optimized blog post titled '{title}' about the topic '{topic}'. Include clear CTAs and an Aussie flair."
    return llm_worker(prompt, system_prompt="You are a Professional Blogger and SEO Copywriter.")

@mcp.tool()
def generate_ad_copy(platform: str, product_name: str) -> str:
    """Draft high-converting ad copy for social media platforms."""
    prompt = f"Draft high-converting, high-CTR ad copy for {platform} promoting the product '{product_name}'. Use psychological triggers and clear CTAs."
    return llm_worker(prompt, system_prompt="You are a Precision Paid Acquisition Expert.")

@mcp.tool()
def script_to_video_hook(topic: str, product_link: str) -> str:
    """Generate viral video hooks and storyboard outlines for multimedia content."""
    prompt = f"Create 3 viral video hooks and a short storyboard outline for a video about '{topic}'. Mention the link: {product_link}."
    return llm_worker(prompt, system_prompt="You are a Viral Multimedia Strategist.")

@mcp.tool()
def draft_automated_sequence(niche: str, goal: str) -> str:
    """Draft a multi-step high-conversion email marketing funnel."""
    prompt = f"Draft a 7-day automated email funnel for the niche '{niche}' with the primary goal: '{goal}'. Include subject lines and body copy."
    return llm_worker(prompt, system_prompt="You are a Master Email Marketing Architect.")

@mcp.tool()
def audit_email_infrastructure(domain: str) -> str:
    """Perform a technical audit of DNS and deliverability infrastructure."""
    prompt = f"Analyze the current email infrastructure for {domain}. Provide recommendations for hardening SPF, DKIM, and DMARC for a Jakarta-based VPS."
    return llm_worker(prompt, system_prompt="You are a Senior Email Deliverability Engineer.")

@mcp.tool()
def estimate_empire_valuation(monthly_profit: float, growth_rate: float) -> str:
    """Provide a professional valuation estimate for the digital portfolio."""
    multiple = 24 if growth_rate < 0.05 else 36
    valuation = monthly_profit * multiple
    prompt = f"Provide a detailed strategic rationale for an empire valuation of ${valuation:,.2f} based on ${monthly_profit}/mo profit and {growth_rate*100}% growth."
    return llm_worker(prompt, system_prompt="You are a Portfolio Valuation and Exit Strategist.")

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

def aussie_router(user_input, history):
    system_instr = load_from_databank("router_instructions.md") or "You are the Aussie Domain Router."
    messages = [{"role": "system", "content": system_instr}]
    for h in history:
        if h["role"] == "user": messages.append({"role": "user", "content": h["content"]})
        if h["role"] == "assistant": messages.append({"role": "assistant", "content": h["content"]})
    messages.append({"role": "user", "content": user_input})
    
    return llm_worker(user_input, system_prompt=system_instr)

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

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)