import os
import textwrap

# ─── CrewAI core ─────────────────────────────────────────────────────────────
from crewai import Agent, Task, Crew, Process, LLM
from crewai_tools import SerperDevTool, ScrapeWebsiteTool

# ─── UI ──────────────────────────────────────────────────────────────────────
import gradio as gr

# =============================================================================
# API KEYS
# =============================================================================
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
SERPER_API_KEY = os.environ.get("SERPER_API_KEY", "")

if not OPENAI_API_KEY:
    raise EnvironmentError(
        "❌ OPENAI_API_KEY not found. "
        "Add it in Space → Settings → Variables and secrets."
    )

if not SERPER_API_KEY:
    raise EnvironmentError(
        "❌ SERPER_API_KEY not found. "
        "Add it in Space → Settings → Variables and secrets."
    )

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
os.environ["SERPER_API_KEY"] = SERPER_API_KEY

# =============================================================================
# LLM
# =============================================================================
llm = LLM(
    model="gpt-4o-mini",
    temperature=0.4,
    max_tokens=2000,
)

# =============================================================================
# TOOLS
# =============================================================================
search_tool = SerperDevTool()
scrape_tool = ScrapeWebsiteTool()

# =============================================================================
# MEMORY
# =============================================================================
MEMORY_DIR = "/tmp/crewai_memory"
os.makedirs(MEMORY_DIR, exist_ok=True)

EMBEDDER_CONFIG = {
    "provider": "openai",
    "config": {
        "model": "text-embedding-3-small",
        "api_key": OPENAI_API_KEY,
    },
}

# =============================================================================
# AGENTS
# =============================================================================
researcher = Agent(
    role="Senior Research Analyst",
    goal=(
        "Find the most accurate, up-to-date, and relevant information on the given topic."
    ),
    backstory="Veteran research analyst with deep investigation skills.",
    tools=[search_tool, scrape_tool],
    llm=llm,
    verbose=True,
    allow_delegation=False,
    max_iter=5,
)

writer = Agent(
    role="Professional Content Writer",
    goal="Write a compelling structured article from research.",
    backstory="Expert in long-form and technical writing.",
    llm=llm,
    verbose=True,
    allow_delegation=False,
    max_iter=4,
)

editor = Agent(
    role="Senior Editor",
    goal="Polish article to publication quality.",
    backstory="Meticulous editor ensuring clarity and consistency.",
    llm=llm,
    verbose=True,
    allow_delegation=False,
    max_iter=3,
)

# =============================================================================
# TASK BUILDER
# =============================================================================
def build_tasks(topic: str, audience: str, tone: str, length: str):
    research_task = Task(
        description=f"Research topic: {topic}",
        expected_output="Structured research summary",
        agent=researcher,
    )

    writing_task = Task(
        description=f"Write article on: {topic}",
        expected_output="Full article",
        agent=writer,
        context=[research_task],
    )

    editing_task = Task(
        description="Edit article to final form",
        expected_output="Final polished article",
        agent=editor,
        context=[writing_task],
    )

    return [research_task, writing_task, editing_task]

# =============================================================================
# CREW RUNNER
# =============================================================================
def run_crew(topic: str, audience: str, tone: str, length: str):
    if not topic.strip():
        return "Please enter a topic."

    try:
        tasks = build_tasks(topic, audience, tone, length)

        crew = Crew(
            agents=[researcher, writer, editor],
            tasks=tasks,
            process=Process.sequential,
            verbose=True,
            memory=True,
            embedder=EMBEDDER_CONFIG,
            memory_config={"long_term": {"storage_path": MEMORY_DIR}},
        )

        result = crew.kickoff()
        return str(result)

    except Exception as e:
        return f"Error: {str(e)}"

# =============================================================================
# UI
# =============================================================================
with gr.Blocks() as demo:

    gr.Markdown("# Multi-Agent Article Generator")

    with gr.Row():

        with gr.Column():
            topic_input = gr.Textbox(label="Topic", lines=3)

            audience_input = gr.Dropdown(
                ["General Public", "Business", "Students"],
                value="General Public",
                label="Audience",
            )

            tone_input = gr.Dropdown(
                ["Informative", "Conversational", "Academic"],
                value="Informative",
                label="Tone",
            )

            length_input = gr.Radio(
                ["Short", "Medium", "Long"],
                value="Medium",
                label="Length",
            )

            submit_btn = gr.Button("Generate")
            clear_btn = gr.Button("Clear")

        with gr.Column():
            output_box = gr.Textbox(
                label="Output",
                lines=30,
                interactive=False,
            )

    submit_btn.click(
        fn=run_crew,
        inputs=[topic_input, audience_input, tone_input, length_input],
        outputs=output_box,
    )

    clear_btn.click(
        fn=lambda: ("", ""),
        inputs=[],
        outputs=[topic_input, output_box],
    )

# =============================================================================
# LAUNCH (UPDATED FOR GRADIO 6)
# =============================================================================
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate"),
    css="""
        footer { display: none !important; }
    """,
)