from crewai import Agent, Task, Crew
import gradio as gr
from gradio import ChatMessage
import asyncio
import re
import sys
from typing import List, Generator
import os
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

load_dotenv()

class ArticleCrew:
    def __init__(self):
        # Agent definitions remain the same
        self.planner = Agent(
            role="Content Planner",
            goal="Plan engaging and factually accurate content on {topic}",
            backstory="You're working on planning a blog article about the topic: {topic}. "
                     "You collect information that helps the audience learn something "
                     "and make informed decisions.",
            allow_delegation=False,
            verbose=True
        )
        
        self.writer = Agent(
            role="Content Writer",
            goal="Write insightful and factually accurate opinion piece about the topic: {topic}",
            backstory="You're working on writing a new opinion piece about the topic: {topic}. "
                     "You base your writing on the work of the Content Planner.",
            allow_delegation=False,
            verbose=True
        )
        
        self.editor = Agent(
            role="Editor",
            goal="Edit a given blog post to align with the writing style",
            backstory="You are an editor who receives a blog post from the Content Writer.",
            allow_delegation=False,
            verbose=True
        )
        
        self.output_parser = OutputParser()

    def create_tasks(self, topic: str):
        # Task definitions remain the same
        plan_task = Task(
            description=(
                f"1. Prioritize the latest trends, key players, and noteworthy news on {topic}.\n"
                f"2. Identify the target audience, considering their interests and pain points.\n"
                f"3. Develop a detailed content outline including introduction, key points, and call to action.\n"
                f"4. Include SEO keywords and relevant data or sources."
            ),
            expected_output="A comprehensive content plan document with an outline, audience analysis, SEO keywords, and resources.",
            agent=self.planner
        )

        write_task = Task(
            description=(
                "1. Use the content plan to craft a compelling blog post.\n"
                "2. Incorporate SEO keywords naturally.\n"
                "3. Sections/Subtitles are properly named in an engaging manner.\n"
                "4. Ensure proper structure with introduction, body, and conclusion.\n"
                "5. Proofread for grammatical errors."
            ),
            expected_output="A well-written blog post in markdown format, ready for publication.",
            agent=self.writer
        )

        edit_task = Task(
            description="Proofread the given blog post for grammatical errors and alignment with the brand's voice.",
            expected_output="A well-written blog post in markdown format, ready for publication.",
            agent=self.editor
        )

        return [plan_task, write_task, edit_task]

    async def process_article(self, topic: str) -> Generator[List[ChatMessage], None, None]:
        crew = Crew(
            agents=[self.planner, self.writer, self.editor],
            tasks=self.create_tasks(topic),
            verbose=2
        )

        class StreamCapture:
            def __init__(self):
                self.data = []
                self.current_chunk = ""
            
            def write(self, text):
                self.current_chunk += text
                if "\n" in text:
                    self.data.append(self.current_chunk)
                    self.current_chunk = ""
                return len(text)

        stream = StreamCapture()
        original_stdout = sys.stdout
        sys.stdout = stream
        
        try:
            result = crew.kickoff(inputs={"topic": topic})
            
            # Process intermediate outputs
            for chunk in stream.data:
                messages = self.output_parser.parse_output(chunk)
                if messages:
                    for msg in messages:
                        yield [msg]
            
            # Send final result
            yield [ChatMessage(
                role="assistant",
                content=result,
                metadata={"title": "📝 Final Article"}
            )]
        
        finally:
            sys.stdout = original_stdout

def create_demo():
    article_crew = ArticleCrew()
    
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 📝 AI Article Writing Crew")
        gr.Markdown("Watch as our AI crew collaborates to create your article!")
        
        chatbot = gr.Chatbot(
            label="Writing Process",
            avatar_images=(None, "🤖"),
            height=700,
            type="messages",
            show_label=True
        )
        
        topic = gr.Textbox(
            label="Article Topic",
            placeholder="Enter the topic you want an article about...",
            lines=2
        )
        
        async def process_input(topic, history):
            # Add user message as ChatMessage
            history.append(ChatMessage(role="user", content=f"Write an article about: {topic}"))
            yield history
            
            # Process and add agent messages
            async for messages in article_crew.process_article(topic):
                history.extend(messages)
                yield history
        
        btn = gr.Button("Write Article")
        btn.click(
            process_input,
            inputs=[topic, chatbot],
            outputs=[chatbot]
        )

    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(debug=True, share=True)