import gradio as gr
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langgraph.graph import StateGraph, END
from typing import Dict, TypedDict, Annotated, List, Tuple, Union, Optional
import json
from langchain_chroma import Chroma
from langchain.schema import Document
from datetime import datetime
from langchain.text_splitter import RecursiveCharacterTextSplitter
from pypdf import PdfReader

# Load environment variables
load_dotenv(verbose=True)

# Verify OpenAI API key
if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OpenAI API key not found.")

# Define state types
class ProcessState(TypedDict):
    pdf_file: str
    content: str
    enhanced: str
    linkedin_post: str
    verification: dict
    error: str
    status: str
    verification_score: float
    enhancement_attempts: int
    needs_improvement: bool
    research_context: str

def extract_pdf_content(pdf_file: str) -> str:
    """Extract text content from PDF file."""
    try:
        reader = PdfReader(pdf_file)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text.strip()
    except Exception as e:
        raise Exception(f"Error extracting PDF content: {str(e)}")

def get_content(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Get content from PDF file."""
    try:
        progress(0.25, desc="Extracting PDF content...")
        content = extract_pdf_content(state["pdf_file"])
        state["content"] = content
        state["status"] = "✅ PDF content extracted"
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error extracting PDF content: {str(e)}"
        state["status"] = "❌ Failed to extract PDF content"
        return state

def get_chroma_collection():
    """Get or create a Chroma collection using OpenAI embeddings."""
    try:
        collection = Chroma(
            collection_name="youtube_videos",
            embedding_function=OpenAIEmbeddings(model="text-embedding-3-small"),
            persist_directory="./chroma_db"
        )
        return collection
    except Exception as e:
        raise Exception(f"Error creating Chroma collection: {str(e)}")

def enhance_content(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Enhance the PDF content with semantic search and similarity analysis."""
    try:
        if not state["content"]:
            return state
        
        progress(0.50, desc="Enhancing content...")
        
        # Get similar content from the vector store
        collection = get_chroma_collection()
        similar_docs = collection.similarity_search(
            state["content"],
            k=3
        )
        
        # Initialize LLM for content generation
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert content enhancer. Transform this PDF content into engaging content:

1. Identify and emphasize key points
2. Add context and examples
3. Make it more engaging and professional
4. Keep it concise (max 3000 characters)
5. Maintain factual accuracy

Content:
{content}

Similar Content for Context:
{similar_content}
"""),
            ("human", "Enhance this content for a professional audience.")
        ])
        
        chain = prompt | llm | StrOutputParser()
        state["enhanced"] = chain.invoke({
            "content": state["content"],
            "similar_content": "\n".join([doc.page_content for doc in similar_docs])
        })
        state["status"] = "✅ Content enhanced"
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error enhancing content: {str(e)}"
        state["status"] = "❌ Failed to enhance content"
        return state

def format_linkedin_post(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Format content as a LinkedIn post."""
    try:
        if not state["enhanced"]:
            return state
        
        progress(0.75, desc="Formatting for LinkedIn...")
        
        # Initialize LLM for formatting
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Create an engaging LinkedIn post from this content. The post should be:

1. Natural and conversational - write like a real person sharing insights
2. Focused on value - emphasize practical takeaways and actionable insights
3. Authentic - avoid overused phrases or corporate speak
4. Visually clean - use line breaks and emojis sparingly and purposefully
5. Under 1500 characters

Content Preservation Rules:
- MUST maintain the exact same topic and subject matter
- MUST keep all specific examples, techniques, and exercises mentioned
- MUST preserve the original context and purpose
- MUST include all key points from the original content
- MUST maintain the same level of technical detail
- MUST keep the same target audience in mind
- MUST preserve any specific terminology or jargon that's important to the topic
- MUST maintain the same tone and expertise level

Formatting Guidelines:
- Start with a hook that grabs attention
- Share insights in a natural flow
- Use 2-3 relevant hashtags maximum
- End with a genuine call to action
- Avoid numbered lists unless absolutely necessary
- Don't use section headers or dividers
- Don't use bullet points or emoji bullets
- Don't use multiple hashtag groups

Content to transform:
{content}

Remember: The goal is to make the content more engaging while keeping ALL the original information, examples, and technical details intact."""),
            ("human", "Create a natural, engaging LinkedIn post that preserves all the original content and context.")
        ])
        
        chain = prompt | llm | StrOutputParser()
        state["linkedin_post"] = chain.invoke({"content": state["enhanced"]})
        state["status"] = "✅ LinkedIn post formatted"
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error formatting LinkedIn post: {str(e)}"
        state["status"] = "❌ Failed to format LinkedIn post"
        return state

def verify_content(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Verify the enhanced content against the original using semantic similarity."""
    try:
        if not state["enhanced"] or not state["content"]:
            return state
        
        progress(1.0, desc="Verifying content...")
        
        # Initialize enhancement attempts if not present
        if "enhancement_attempts" not in state:
            state["enhancement_attempts"] = 0
        
        # Calculate semantic similarity using Chroma
        collection = get_chroma_collection()
        similar_docs = collection.similarity_search(
            state["enhanced"],
            k=1
        )
        similarity_score = 0.0
        if similar_docs:
            # Chroma returns a list of Document objects with a score attribute
            # But the default similarity_search does not return scores, so we just check if content is similar
            similarity_score = 1.0 if similar_docs[0].page_content == state["content"] else 0.0
        
        # Initialize LLM for verification
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Verify the enhanced content against the original:

1. Check factual accuracy
2. Ensure key messages are preserved
3. Look for any misrepresentations

Return JSON in this format:
{{
    "verified": boolean,
    "score": float between 0-1,
    "feedback": string with details
}}

Original:
{original}

Enhanced:
{enhanced}

Semantic Similarity Score: {similarity_score}"""),
            ("human", "Verify this content.")
        ])
        
        chain = prompt | llm | StrOutputParser()
        verification_result = json.loads(chain.invoke({
            "original": state["content"],
            "enhanced": state["enhanced"],
            "similarity_score": similarity_score
        }))
        
        # Update state with verification results
        state["verification"] = verification_result
        state["verification_score"] = verification_result["score"]
        
        # Trigger agent decision if score is below threshold
        if verification_result["score"] < 0.85 and state["enhancement_attempts"] < 3:
            state["needs_improvement"] = True
            # Create improvement plan
            state = agent_decide(state)
            state["status"] = f"🔄 Planning improvements (Attempt {state['enhancement_attempts'] + 1}/3)"
        else:
            state["needs_improvement"] = False
            if verification_result["score"] >= 0.85:
                state["status"] = "✅ Content quality threshold met"
            else:
                state["status"] = "⚠️ Max enhancement attempts reached"
        
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error verifying content: {str(e)}"
        state["status"] = "❌ Failed to verify content"
        return state

def should_continue(state: ProcessState) -> bool:
    """Determine if processing should continue."""
    return not state.get("error", "")

def create_workflow() -> StateGraph:
    """Create the LangGraph workflow."""
    workflow = StateGraph(ProcessState)
    
    # Add nodes
    workflow.add_node("get_content", get_content)
    workflow.add_node("enhance_content", enhance_content)
    workflow.add_node("format_linkedin", format_linkedin_post)
    workflow.add_node("verify_content", verify_content)
    workflow.add_node("agent_decide", agent_decide)
    workflow.add_node("research_content", research_content)
    workflow.add_node("enhance_again", enhance_again)
    
    # Set entry point
    workflow.set_entry_point("get_content")
    
    # Add edges for main flow
    workflow.add_edge("get_content", "enhance_content")
    workflow.add_edge("enhance_content", "format_linkedin")
    workflow.add_edge("format_linkedin", "verify_content")
    workflow.add_edge("verify_content", "agent_decide")
    
    # Add conditional edges for agentic flow
    workflow.add_conditional_edges(
        "agent_decide",
        lambda x: x["needs_improvement"],
        {
            True: "research_content",
            False: END
        }
    )
    
    # Add edges for enhancement loop
    workflow.add_edge("research_content", "enhance_again")
    workflow.add_edge("enhance_again", "verify_content")
    
    # Add conditional edges for error handling
    workflow.add_conditional_edges(
        "get_content",
        should_continue,
        {
            True: "enhance_content",
            False: END
        }
    )
    workflow.add_conditional_edges(
        "enhance_content",
        should_continue,
        {
            True: "format_linkedin",
            False: END
        }
    )
    workflow.add_conditional_edges(
        "format_linkedin",
        should_continue,
        {
            True: "verify_content",
            False: END
        }
    )
    workflow.add_conditional_edges(
        "verify_content",
        should_continue,
        {
            True: "agent_decide",
            False: END
        }
    )
    workflow.add_conditional_edges(
        "research_content",
        should_continue,
        {
            True: "enhance_again",
            False: END
        }
    )
    workflow.add_conditional_edges(
        "enhance_again",
        should_continue,
        {
            True: "verify_content",
            False: END
        }
    )
    
    return workflow

def process_video(video_url: str, progress=gr.Progress()) -> tuple:
    """Process YouTube video and generate LinkedIn post."""
    try:
        # Input validation
        if not video_url:
            return (
                "⚠️ Please enter a YouTube URL",  # error
                "❌ Failed: No URL provided",      # status
                "",                               # transcript
                "",                               # enhanced
                "",                               # linkedin
                ""                                # verification
            )
        
        if "youtube.com" not in video_url and "youtu.be" not in video_url:
            return (
                "⚠️ Invalid URL. Please enter a YouTube URL",  # error
                "❌ Failed: Invalid URL",                      # status
                "",                                           # transcript
                "",                                           # enhanced
                "",                                           # linkedin
                ""                                            # verification
            )
        
        # Initialize state
        initial_state = ProcessState(
            video_url=video_url,
            transcript="",
            enhanced="",
            linkedin_post="",
            verification={},
            error="",
            status="Starting..."
        )
        
        # Create and run workflow
        workflow = create_workflow()
        app = workflow.compile()
        final_state = app.invoke(initial_state)
        
        # Format verification text
        if final_state.get("verification"):
            verification_text = f"""Verification Results:
• Status: {"✅ Verified" if final_state["verification"]["verified"] else "❌ Not Verified"}
• Accuracy Score: {final_state["verification"]["score"]:.2f}
• Feedback: {final_state["verification"]["feedback"]}"""
        else:
            verification_text = ""
        
        return (
            final_state.get("error", ""),           # error
            final_state.get("status", ""),          # status
            final_state.get("transcript", ""),      # transcript
            final_state.get("enhanced", ""),        # enhanced
            final_state.get("linkedin_post", ""),   # linkedin
            verification_text                       # verification
        )
        
    except Exception as e:
        return (
            f"⚠️ Error: {str(e)}",  # error
            "❌ Processing failed",  # status
            "",                     # transcript
            "",                     # enhanced
            "",                     # linkedin
            ""                      # verification
        )

def process_from_stage(state: ProcessState, start_stage: str, progress=gr.Progress()) -> tuple:
    """Process content from a specific stage onwards."""
    try:
        # Select appropriate workflow based on stage
        if start_stage == "enhance":
            workflow = create_workflow()
            if not state["content"]:
                return (
                    "⚠️ No content available to enhance",
                    "❌ Failed: No content",
                    state.get("content", ""),
                    "",
                    "",
                    ""
                )
        elif start_stage == "format":
            workflow = create_workflow()
            if not state["enhanced"]:
                return (
                    "⚠️ No enhanced content available to format",
                    "❌ Failed: No enhanced content",
                    state.get("content", ""),
                    state.get("enhanced", ""),
                    "",
                    ""
                )
        else:
            workflow = create_workflow()
        
        app = workflow.compile()
        final_state = app.invoke(state)
        
        # Format verification text
        if final_state.get("verification"):
            verification_text = f"""Verification Results:
• Status: {"✅ Verified" if final_state["verification"]["verified"] else "❌ Not Verified"}
• Accuracy Score: {final_state["verification"]["score"]:.2f}
• Feedback: {final_state["verification"]["feedback"]}"""
        else:
            verification_text = ""
        
        return (
            final_state.get("error", ""),
            final_state.get("status", ""),
            final_state.get("content", ""),
            final_state.get("enhanced", ""),
            final_state.get("linkedin_post", ""),
            verification_text
        )
        
    except Exception as e:
        return (
            f"⚠️ Error: {str(e)}",
            "❌ Processing failed",
            state.get("content", ""),
            state.get("enhanced", ""),
            state.get("linkedin_post", ""),
            ""
        )

def format_verification_text(verification: dict) -> str:
    """Format verification results into a readable string."""
    if not verification:
        return ""
        
    return f"""Verification Results:
• Status: {"✅ Verified" if verification.get("verified") else "❌ Not Verified"}
• Accuracy Score: {verification.get("score", 0):.2f}
• Feedback: {verification.get("feedback", "No feedback available")}"""

def safe_json_loads(json_str: str, default: dict = None) -> dict:
    """Safely parse JSON string with error handling."""
    if default is None:
        default = {}
    try:
        return json.loads(json_str) if json_str else default
    except json.JSONDecodeError:
        return default

def format_improvement_plan(plan: dict) -> str:
    """Format the improvement plan into a readable string."""
    if not plan:
        return "No improvement plan available"
        
    text = "📋 Improvement Plan:\n\n"
    
    # Improvement Areas
    if "improvement_areas" in plan:
        text += "🎯 Priority Areas:\n"
        for area in plan["improvement_areas"]:
            text += f"• {area.get('area', 'N/A')} (Priority: {area.get('priority', 'N/A')}/5)\n"
            text += f"  Strategy: {area.get('strategy', 'N/A')}\n"
            text += f"  Research Focus: {area.get('research_focus', 'N/A')}\n\n"
    
    # Research Priorities
    if "research_priorities" in plan:
        text += "🔍 Research Priorities:\n"
        for topic in plan["research_priorities"]:
            text += f"• {topic.get('topic', 'N/A')}\n"
            text += f"  Reason: {topic.get('reason', 'N/A')}\n"
            text += f"  Expected Impact: {topic.get('expected_impact', 'N/A')}\n\n"
    
    # Enhancement Strategy
    if "enhancement_strategy" in plan:
        text += "⚡ Enhancement Strategy:\n"
        strategy = plan["enhancement_strategy"]
        text += f"• Approach: {strategy.get('approach', 'N/A')}\n"
        text += f"• Key Focus: {strategy.get('key_focus', 'N/A')}\n"
        text += "• Expected Improvements:\n"
        for imp in strategy.get("expected_improvements", []):
            text += f"  - {imp}\n"
    
    return text

def format_research_results(research: dict) -> str:
    """Format the research results into a readable string."""
    if not research:
        return "No research results available"
        
    text = "📚 Research Results:\n\n"
    
    # Focused Research
    if "focused_research" in research:
        text += "🎯 Focused Research by Area:\n"
        for area, data in research["focused_research"].items():
            text += f"• {area} (Priority: {data.get('priority', 'N/A')}/5)\n"
            text += f"  Strategy: {data.get('strategy', 'N/A')}\n"
            text += "  Key Findings:\n"
            for content in data.get("content", [])[:1]:  # Show first finding
                text += f"  - {content[:200]}...\n\n"
    
    # Additional Research
    if research.get("similar_content"):
        text += "📖 Additional Research:\n"
        for content in research["similar_content"][:2]:  # Show first two
            text += f"• {content[:200]}...\n\n"
    
    return text

def create_ui():
    with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
        current_state = gr.State({
            "pdf_file": "",
            "content": "",
            "enhanced": "",
            "linkedin_post": "",
            "verification": {},
            "error": "",
            "status": "",
            "improvement_plan": {},
            "research_context": "{}",
            "enhancement_attempts": 0,
            "needs_improvement": False
        })
        
        gr.Markdown(
            """
            # PDF to LinkedIn Post Converter
            Transform your PDF documents into professional LinkedIn posts with AI content enhancement.
            
            ### 📄 How to Use
            1. Upload a PDF file
            2. Click "Generate Post"
            3. Review the enhanced content
            4. Copy your LinkedIn-ready post
            """
        )
        
        with gr.Row():
            with gr.Column():
                pdf_file = gr.File(
                    label="PDF File",
                    file_types=[".pdf"],
                    type="filepath"
                )
                convert_btn = gr.Button("🚀 Generate from PDF", variant="primary", size="lg")

        status = gr.Textbox(
            label="Status",
            value="Ready to process...",
            interactive=False
        )
        
        error = gr.Textbox(
            label="Error",
            visible=False,
            interactive=False
        )

        with gr.Tabs() as tabs:
            with gr.TabItem("📝 Content"):
                with gr.Row():
                    with gr.Column():
                        content = gr.TextArea(
                            label="📄 Raw Content",
                            interactive=False,
                            show_copy_button=True,
                            lines=8
                        )
                    with gr.Column():
                        enhanced = gr.TextArea(
                            label="✨ Enhanced Content",
                            interactive=False,
                            show_copy_button=True,
                            lines=8
                        )
                
                with gr.Row():
                    with gr.Column():
                        linkedin = gr.TextArea(
                            label="🔗 LinkedIn Post",
                            interactive=False,
                            show_copy_button=True,
                            lines=6
                        )
                
                with gr.Row():
                    with gr.Column():
                        verification = gr.TextArea(
                            label="✓ Verification Results",
                            interactive=False,
                            lines=4
                        )
                
                with gr.Row():
                    with gr.Column():
                        improvement_plan = gr.TextArea(
                            label="📋 Improvement Plan",
                            interactive=False,
                            show_copy_button=True,
                            lines=8,
                            visible=True,
                            value="Waiting for verification..."
                        )
                
                with gr.Row():
                    with gr.Column():
                        research_results = gr.TextArea(
                            label="🔍 Research Results",
                            interactive=False,
                            show_copy_button=True,
                            lines=8,
                            visible=True,
                            value="Waiting for research..."
                        )
                
                with gr.Row():
                    with gr.Column():
                        improved_linkedin = gr.TextArea(
                            label="🚀 Improved LinkedIn Post Final",
                            interactive=False,
                            show_copy_button=True,
                            lines=6,
                            visible=True,
                            value="Waiting for improvements..."
                        )

                # Loading indicators
                with gr.Row(visible=False) as loading_indicators:
                    content_loading = gr.Markdown("🔄 Extracting content...")
                    enhanced_loading = gr.Markdown("🔄 Enhancing content...")
                    linkedin_loading = gr.Markdown("🔄 Formatting for LinkedIn...")
                    verify_loading = gr.Markdown("🔄 Verifying content...")
                    plan_loading = gr.Markdown("🔄 Creating improvement plan...")
                    research_loading = gr.Markdown("🔄 Researching content...")
                    improved_loading = gr.Markdown("🔄 Creating improved post...")
            
            with gr.TabItem("ℹ️ Help"):
                gr.Markdown(
                    """
                    ### How to Use
                    1. **Input**: Upload a PDF file
                    2. **Process**: Click the "Generate Post" button
                    3. **Wait**: The system will process your PDF through multiple steps
                    4. **Review**: Check the generated content in each tab
                    5. **Copy**: Use the copy button to grab your LinkedIn post
                    
                    ### 🔄 Regeneration Options
                    - Click 🔄 next to "Enhanced Content" to regenerate from the enhancement stage
                    - Click 🔄 next to "LinkedIn Post" to regenerate from the formatting stage
                    
                    ### 💡 Tips for Best Results
                    - Use well-formatted PDFs with clear text
                    - Optimal length: 2-10 pages
                    - Ensure PDFs have readable text (not scanned images)
                    - Review and personalize the post before sharing
                    - Consider your target audience when selecting content
                    """
                )

        def update_loading_state(stage: str):
            """Update loading indicators based on current stage."""
            states = {
                "content": [True, False, False, False, False, False, False],
                "enhance": [False, True, False, False, False, False, False],
                "format": [False, False, True, False, False, False, False],
                "verify": [False, False, False, True, False, False, False],
                "plan": [False, False, False, False, True, False, False],
                "research": [False, False, False, False, False, True, False],
                "improved": [False, False, False, False, False, False, True],
                "done": [False, False, False, False, False, False, False]
            }
            
            # Loading messages for each stage
            loading_messages = {
                "content": "🔄 Extracting content...\n⏳ Please wait...",
                "enhance": "✨ Enhancing content...\n⚡ AI is working its magic...",
                "format": "🎨 Formatting for LinkedIn...\n📝 Creating engaging post...",
                "verify": "🔍 Verifying content...\n⚖️ Checking accuracy...",
                "plan": "🔄 Creating improvement plan...",
                "research": "🔎 Researching content...\n📚 Finding relevant information...",
                "improved": "🚀 Creating improved LinkedIn post...\n✨ Applying enhancements..."
            }
            
            # Get current stage message
            current_message = loading_messages.get(stage, "")
            
            # Return loading states and message
            return [
                gr.update(visible=state) for state in states.get(stage, [False] * 7)
            ], current_message

        def process_with_loading(pdf_path, state):
            """Process PDF with loading indicators."""
            try:
                # Initialize state if needed
                if "improvement_plan" not in state:
                    state["improvement_plan"] = {}
                if "research_context" not in state:
                    state["research_context"] = "{}"
                if "enhancement_attempts" not in state:
                    state["enhancement_attempts"] = 0
                if "needs_improvement" not in state:
                    state["needs_improvement"] = False
                    
                # Show loading indicators
                loading_states, message = update_loading_state("content")
                yield [
                    "",  # error
                    "Processing...",  # status
                    message,  # content (loading)
                    "",  # enhanced
                    "",  # linkedin
                    "",  # verification
                    "Waiting for verification...",  # improvement plan
                    "Waiting for research...",  # research results
                    "Waiting for improvements...",  # improved linkedin
                    state,  # current_state
                    *loading_states  # loading indicators
                ]
                
                # Get content
                state["pdf_file"] = pdf_path
                content_text = get_content(state)["content"]
                
                # Show enhancing state
                loading_states, message = update_loading_state("enhance")
                yield [
                    "",
                    "Enhancing content...",
                    content_text,
                    message,  # enhanced (loading)
                    "",
                    "",
                    "",
                    "",
                    "",
                    state,
                    *loading_states
                ]
                
                # Enhance content
                state["content"] = content_text
                enhanced_state = enhance_content(state)
                enhanced_text = enhanced_state["enhanced"]
                
                # Show formatting state
                loading_states, message = update_loading_state("format")
                yield [
                    "",
                    "Formatting for LinkedIn...",
                    content_text,
                    enhanced_text,
                    message,  # linkedin (loading)
                    "",
                    "",
                    "",
                    "",
                    state,
                    *loading_states
                ]
                
                # Format LinkedIn post
                state["enhanced"] = enhanced_text
                linkedin_state = format_linkedin_post(state)
                linkedin_text = linkedin_state["linkedin_post"]
                
                # Show verifying state
                loading_states, message = update_loading_state("verify")
                yield [
                    "",
                    "Verifying content...",
                    content_text,
                    enhanced_text,
                    linkedin_text,
                    "🔍 Verifying...\n⚖️ Analyzing accuracy...",  # verification (loading)
                    "",
                    "",
                    "",
                    state,
                    *loading_states
                ]
                
                # Verify content
                state["linkedin_post"] = linkedin_text
                final_state = verify_content(state)
                verification_text = format_verification_text(final_state.get("verification", {}))
                
                # Update improvement plan and research results
                improvement_plan_text = format_improvement_plan(final_state.get("improvement_plan", {}))
                research_results_text = format_research_results(safe_json_loads(final_state.get("research_context", "{}")))
                
                # Check if enhancement is needed
                if final_state.get("needs_improvement", False):
                    # Show planning state
                    loading_states, message = update_loading_state("plan")
                    yield [
                        "",
                        f"Creating improvement plan (Attempt {final_state.get('enhancement_attempts', 1)}/3)...",
                        content_text,
                        enhanced_text,
                        linkedin_text,
                        verification_text,
                        improvement_plan_text,
                        research_results_text,
                        "",
                        state,
                        *loading_states
                    ]
                    
                    # Show researching state
                    loading_states, message = update_loading_state("research")
                    yield [
                        "",
                        f"Researching content (Attempt {final_state.get('enhancement_attempts', 1)}/3)...",
                        content_text,
                        enhanced_text,
                        linkedin_text,
                        verification_text,
                        improvement_plan_text,
                        research_results_text,
                        "",
                        state,
                        *loading_states
                    ]
                    
                    # Research content
                    state = research_content(state)
                    research_results_text = format_research_results(safe_json_loads(state.get("research_context", "{}")))
                    
                    # Show enhancing again state
                    loading_states, message = update_loading_state("enhance")
                    yield [
                        "",
                        f"Enhancing content again (Attempt {final_state.get('enhancement_attempts', 1)}/3)...",
                        content_text,
                        enhanced_text,
                        linkedin_text,
                        verification_text,
                        improvement_plan_text,
                        research_results_text,
                        "",
                        state,
                        *loading_states
                    ]
                    
                    # Enhance again
                    state = enhance_again(state)
                    enhanced_text = state["enhanced"]
                    
                    # Update LinkedIn post
                    state["enhanced"] = enhanced_text
                    linkedin_state = format_linkedin_post(state)
                    linkedin_text = linkedin_state["linkedin_post"]
                    
                    # Verify again
                    state["linkedin_post"] = linkedin_text
                    final_state = verify_content(state)
                    verification_text = format_verification_text(final_state.get("verification", {}))
                    improvement_plan_text = format_improvement_plan(final_state.get("improvement_plan", {}))
                    research_results_text = format_research_results(safe_json_loads(final_state.get("research_context", "{}")))
                
                # After research and enhancement, create improved LinkedIn post
                if final_state.get("needs_improvement", False):
                    # Show improved post loading state
                    loading_states, message = update_loading_state("improved")
                    yield [
                        "",
                        f"Creating improved LinkedIn post (Attempt {final_state.get('enhancement_attempts', 1)}/3)...",
                        content_text,
                        enhanced_text,
                        linkedin_text,
                        verification_text,
                        improvement_plan_text,
                        research_results_text,
                        message,  # improved linkedin (loading)
                        state,
                        *loading_states
                    ]
                    
                    # Create improved LinkedIn post
                    improved_state = format_linkedin_post(final_state)
                    improved_text = improved_state["linkedin_post"]
                    
                    # Update final state
                    final_state["improved_linkedin"] = improved_text
                
                # Complete
                loading_states, _ = update_loading_state("done")
                yield [
                    "",
                    "✅ Processing complete!",
                    content_text,
                    enhanced_text,
                    linkedin_text,
                    verification_text,
                    improvement_plan_text,
                    research_results_text,
                    final_state.get("improved_linkedin", "No improvements needed"),
                    final_state,
                    *loading_states
                ]
                
            except Exception as e:
                loading_states, _ = update_loading_state("done")
                yield [
                    f"⚠️ Error: {str(e)}",
                    "❌ Processing failed",
                    state.get("content", ""),
                    state.get("enhanced", ""),
                    state.get("linkedin_post", ""),
                    "",
                    "Error occurred during processing",
                    "Error occurred during processing",
                    "Error occurred during processing",
                    state,
                    *loading_states
                ]

        # Set up event handlers
        convert_btn.click(
            fn=process_with_loading,
            inputs=[pdf_file, current_state],
            outputs=[
                error,
                status,
                content,
                enhanced,
                linkedin,
                verification,
                improvement_plan,
                research_results,
                improved_linkedin,
                current_state,
                content_loading,
                enhanced_loading,
                linkedin_loading,
                verify_loading,
                plan_loading,
                research_loading,
                improved_loading
            ],
            show_progress=True,  # Show progress bar
            api_name="convert"  # Name the API endpoint
        )
        
        # Update error visibility with immediate feedback
        error.change(
            lambda x: gr.update(visible=bool(x), value=x),  # Update both visibility and value
            error,
            error,
            queue=False  # Process immediately
        )

        # Add loading state visibility updates
        def update_loading_visibility(is_loading):
            return {
                loading: gr.update(visible=is_loading)
                for loading in [
                    content_loading,
                    enhanced_loading,
                    linkedin_loading,
                    verify_loading,
                    plan_loading,
                    research_loading,
                    improved_loading
                ]
            }

        convert_btn.click(
            lambda: update_loading_visibility(True),
            None,
            [content_loading, enhanced_loading, linkedin_loading, 
             verify_loading, plan_loading, research_loading, improved_loading],
            queue=False
        )

    return demo

def agent_decide(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Agent decides whether to enhance content further based on verification score and creates an improvement plan."""
    try:
        progress(0.95, desc="Analyzing content quality and planning improvements...")
        
        # Get verification score and attempts
        score = state.get("verification", {}).get("score", 0)
        attempts = state.get("enhancement_attempts", 0)
        feedback = state.get("verification", {}).get("feedback", "")
        
        # Initialize LLM for agentic decision making
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert content strategist. Analyze the content quality and create an improvement plan.

Current Content:
{content}

Verification Results:
- Score: {score}
- Feedback: {feedback}
- Previous Attempts: {attempts}

Create a detailed improvement plan in JSON format:
{{
    "needs_improvement": boolean,
    "improvement_areas": [
        {{
            "area": string,
            "priority": number (1-5),
            "strategy": string,
            "research_focus": string
        }}
    ],
    "research_priorities": [
        {{
            "topic": string,
            "reason": string,
            "expected_impact": string
        }}
    ],
    "enhancement_strategy": {{
        "approach": string,
        "key_focus": string,
        "expected_improvements": [string]
    }}
}}

Consider:
1. Content quality and engagement
2. Information accuracy and completeness
3. Target audience needs
4. Previous enhancement attempts
5. Available research context"""),
            ("human", "Analyze this content and create an improvement plan.")
        ])
        
        chain = prompt | llm | StrOutputParser()
        plan = json.loads(chain.invoke({
            "content": state["enhanced"],
            "score": score,
            "feedback": feedback,
            "attempts": attempts
        }))
        
        # Update state with plan
        state["verification_score"] = score
        state["enhancement_attempts"] = attempts
        state["needs_improvement"] = plan["needs_improvement"]
        state["improvement_plan"] = plan
        
        # Create detailed status message
        if plan["needs_improvement"] and attempts < 3:
            status = f"🔄 Planning improvements (Attempt {attempts + 1}/3)\n"
            status += "Key focus areas:\n"
            for area in plan["improvement_areas"][:2]:  # Show top 2 priorities
                status += f"• {area['area']} (Priority: {area['priority']})\n"
            state["status"] = status
        else:
            if score >= 0.95:
                state["status"] = "✅ Content quality threshold met"
            else:
                state["status"] = "⚠️ Max enhancement attempts reached"
        
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error in agent decision: {str(e)}"
        state["status"] = "❌ Failed to analyze content"
        return state

def research_content(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Research additional context based on the improvement plan."""
    try:
        progress(0.96, desc="Researching based on improvement plan...")
        
        # Get improvement plan
        plan = state.get("improvement_plan", {})
        if not plan:
            raise Exception("No improvement plan found")
        
        # Initialize research results
        research_results = {
            "similar_content": [],
            "focused_research": {},
            "verification_feedback": state.get("verification", {}).get("feedback", "")
        }
        
        # Get similar content from vector store
        collection = get_chroma_collection()
        
        # Research each priority area
        for area in plan["improvement_areas"]:
            # Search for content related to this area
            similar_docs = collection.similarity_search(
                f"{area['area']} {area['research_focus']}",
                k=2
            )
            
            # Store research results
            research_results["focused_research"][area["area"]] = {
                "content": [doc.page_content for doc in similar_docs],
                "priority": area["priority"],
                "strategy": area["strategy"]
            }
        
        # Research specific topics from research_priorities
        for topic in plan["research_priorities"]:
            topic_docs = collection.similarity_search(
                topic["topic"],
                k=1
            )
            if topic_docs:
                research_results["similar_content"].extend([doc.page_content for doc in topic_docs])
        
        # Store research results
        state["research_context"] = json.dumps(research_results)
        state["status"] = "✅ Research completed based on improvement plan"
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error researching content: {str(e)}"
        state["status"] = "❌ Failed to research content"
        return state

def enhance_again(state: ProcessState, progress=gr.Progress()) -> ProcessState:
    """Enhance content using research and improvement plan."""
    try:
        progress(0.97, desc="Enhancing content based on research and plan...")
        
        # Get research context and improvement plan
        research_context = json.loads(state["research_context"])
        plan = state.get("improvement_plan", {})
        if not plan:
            raise Exception("No improvement plan found")
        
        # Initialize LLM for enhancement
        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert content enhancer. Improve the content based on the research and improvement plan while maintaining the original topic and key messages.

Current Content:
{content}

Improvement Plan:
{plan}

Research Results:
{research}

Enhancement Strategy:
{strategy}

Create enhanced content that:
1. Maintains the original topic and key messages
2. Addresses each improvement area according to its priority
3. Incorporates relevant research findings
4. Follows the enhancement strategy
5. Improves engagement and clarity
6. Keeps the same core subject matter and examples

Important:
- DO NOT change the main topic or subject matter
- DO NOT replace specific examples with generic ones
- DO NOT lose the original context or purpose
- DO NOT generate content about a different topic
- DO preserve and enhance the original message"""),
            ("human", "Enhance this content while maintaining its original topic and key messages.")
        ])
        
        chain = prompt | llm | StrOutputParser()
        enhanced = chain.invoke({
            "content": state["enhanced"],
            "plan": json.dumps(plan),
            "research": json.dumps(research_context),
            "strategy": json.dumps(plan["enhancement_strategy"])
        })
        
        # Update state
        state["enhanced"] = enhanced
        state["enhancement_attempts"] = state.get("enhancement_attempts", 0) + 1
        state["status"] = f"✅ Content enhanced with research (Attempt {state['enhancement_attempts']}/3)"
        return state
    except Exception as e:
        state["error"] = f"⚠️ Error enhancing content: {str(e)}"
        state["status"] = "❌ Failed to enhance content"
        return state

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