import gradio as gr
import json
from transformers import pipeline
from marketing_graph import build_marketing_graph
from langchain_core.tracers.context import tracing_v2_enabled

# ---------------- Load Qwen Model ----------------
# Using Qwen 2.5B instruct model (or another Qwen variant you have access to)
generator = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-0.5B-Instruct",
    device_map="auto"
)

# ---------------- Marketing Agent ----------------
def run_marketing_agent(product_name, audience, goal):
    graph = build_marketing_graph()

    state = {
        "input_data": f"""
**Product:** {product_name}  
**Target Audience:** {audience}  
**Goal:** {goal}
"""
    }

    executed_nodes = []
    current_node = "MarketResearch"

    # Enable LangSmith tracing
    with tracing_v2_enabled(project_name="MarketingAgentGraph"):
        while current_node:
            if current_node not in graph.nodes:
                break

            executed_nodes.append(current_node)
            node_func = graph.nodes[current_node]

            try:
                # Use Qwen generator for text-heavy nodes
                if current_node in ["MarketResearch", "ContentPlanner", "PostGenerator", "AnalyticsAgent"]:
                    input_text = state.get("input_data", "") + "\n\n" + state.get(current_node, "")
                    qwen_output = generator(input_text, max_new_tokens=300, do_sample=True, temperature=0.7)
                    output_text = qwen_output[0]["generated_text"]
                    output = {current_node: output_text}
                else:
                    output = node_func(state)

            except Exception as e:
                # Self-healing if node fails
                fixed = manager.marketing_fix_node(
                    current_node,
                    state.get("input_data", ""),
                    str(e),
                )
                output = {current_node: fixed}

            state.update(output)

            # Determine next node
            next_nodes = [to_node for (from_node, to_node) in graph.edges if from_node == current_node]
            current_node = next_nodes[0] if next_nodes and next_nodes[0] in graph.nodes else None

    market_result = state.get("MarketResearch", "No Market Research generated")

    emoji_map = {
        "MarketResearch": "📝 Market Research",
        "ContentPlanner": "📅 Content Plan",
        "PostGenerator": "📣 Generated Posts",
        "AnalyticsAgent": "📊 Analytics & Suggestions"
    }

    sections = []
    for node in executed_nodes:
        content = state.get(node, "")
        if content:
            sections.append(f"## {emoji_map.get(node, node)}")
            sections.append(content)

    sections.append("## ✨ FINAL MARKETING RESULT ✨")
    sections.append(f"```json\n{json.dumps(state, indent=2)}\n```")
    full_strategy = "\n\n".join(sections)

    return market_result, full_strategy

# --- Gradio Interface ---
with gr.Blocks(title="🧠 Self-Healing Marketing Agent (Qwen Model)") as demo:
    gr.Markdown("## Enter your product info to generate Market Research and a full Marketing Strategy")

    with gr.Row():
        with gr.Column(scale=1):
            product_name_input = gr.Textbox(label="Product Name", value="Smart fitness band")
            audience_input = gr.Textbox(label="Target Audience", value="young professionals")
            goal_input = gr.Textbox(label="Goal", value="Increase brand awareness")
            submit_btn = gr.Button("🚀 Generate Strategy", variant="primary")
        with gr.Column(scale=2):
            market_output_md = gr.Markdown()
            full_output_md = gr.Markdown()

    submit_btn.click(
        run_marketing_agent,
        inputs=[product_name_input, audience_input, goal_input],
        outputs=[market_output_md, full_output_md]
    )

if __name__ == "__main__":
    demo.launch()