import gradio as gr
import os
import time
import tempfile
import requests
from PIL import Image
from io import BytesIO
import markdown
import re
import json
import random
from transformers import pipeline
from huggingface_hub import HfApi
from linkedin_api import Linkedin

# Initialize models
try:
    text_generation = pipeline(
        "text-generation",
        model="mistralai/Mistral-7B-Instruct-v0.2",
        max_length=4096,
        temperature=0.7
    )
    image_generation = pipeline("text-to-image", model="runwayml/stable-diffusion-v1-5")
except Exception as e:
    print(f"Error loading models: {e}")
    # Fallback to smaller models if needed
    try:
        text_generation = pipeline(
            "text-generation",
            model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", 
            max_length=2048
        )
        image_generation = None
    except:
        text_generation = None
        image_generation = None

# Function to generate blog content
def generate_blog_content(topic, tone="professional", length="medium"):
    if not text_generation:
        return "Error: Text generation model could not be loaded."
    
    length_words = {
        "short": "500-800",
        "medium": "1000-1500",
        "long": "2000-2500"
    }[length]
    
    tone_instructions = {
        "professional": "Use a formal, business-like tone with industry terminology.",
        "casual": "Write in a conversational, friendly tone as if talking to a peer.",
        "technical": "Include detailed technical information and analysis.",
        "storytelling": "Structure the content as a narrative with examples and stories."
    }[tone]
    
    prompt = f"""
    Write a complete blog post about {topic}.
    
    {tone_instructions}
    
    The blog post should be approximately {length_words} words and include:
    - An attention-grabbing headline
    - An engaging introduction
    - 3-5 well-structured sections with subheadings
    - Practical insights and takeaways
    - A conclusion
    
    Format the blog in markdown with proper headings, bullet points, and emphasis.
    """
    
    try:
        result = text_generation(prompt, max_length=4096)[0]['generated_text']
        # Clean up the output - extract just the blog post
        blog_content = result.split(prompt)[-1].strip()
        return blog_content
    except Exception as e:
        return f"Error generating blog content: {str(e)}"

# Function to generate image
def generate_featured_image(topic):
    if not image_generation:
        return None, "Image generation not available. Using default image."
    
    prompt = f"Professional illustration for blog about {topic}, digital art, high quality"
    try:
        image = image_generation(prompt)
        if isinstance(image, list):
            image = image[0] if image else None
        
        temp_img_path = f"temp_image_{random.randint(1000, 9999)}.png"
        if hasattr(image, 'save'):
            image.save(temp_img_path)
        else:
            # Handle different return types
            if isinstance(image, dict) and 'images' in image:
                image = Image.fromarray(image['images'][0])
                image.save(temp_img_path)
                
        return temp_img_path, "Image generated successfully"
    except Exception as e:
        return None, f"Error generating image: {str(e)}"

# Function to post to LinkedIn
def post_to_linkedin(content, image_path=None, linkedin_username=None, linkedin_password=None):
    if not linkedin_username or not linkedin_password:
        return "Error: LinkedIn credentials are required."
    
    try:
        # Extract title from markdown
        title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE)
        title = title_match.group(1) if title_match else "New Blog Post"
        
        # Convert markdown to plain text for LinkedIn
        # Remove markdown formatting for LinkedIn post
        plain_content = content
        plain_content = re.sub(r'^#+\s+', '', plain_content, flags=re.MULTILINE)  # Remove headings
        plain_content = re.sub(r'\*\*(.*?)\*\*', r'\1', plain_content)  # Remove bold
        plain_content = re.sub(r'\*(.*?)\*', r'\1', plain_content)  # Remove italic
        
        # Shorten for LinkedIn
        if len(plain_content) > 1300:  # LinkedIn character limit
            plain_content = plain_content[:1297] + "..."
        
        # Add a title and link to full blog if available
        post_text = f"{title}\n\n{plain_content}"
        
        # Initialize LinkedIn API
        api = Linkedin(linkedin_username, linkedin_password)
        
        # Post to LinkedIn
        if image_path and os.path.exists(image_path):
            # Upload image first
            media_id = api.upload_image(image_path)
            # Post with image
            post_response = api.create_post(post_text, media_ids=[media_id])
        else:
            # Text-only post
            post_response = api.create_post(post_text)
        
        # Clean up temporary image file
        if image_path and os.path.exists(image_path):
            try:
                os.remove(image_path)
            except:
                pass
                
        return f"Successfully posted to LinkedIn: {title}"
    except Exception as e:
        return f"Error posting to LinkedIn: {str(e)}"

# Function to save as Hugging Face Space
def save_as_blog(content, title, author, image_path=None, hf_token=None):
    if not hf_token:
        return "Error: Hugging Face token is required to save blog."
    
    try:
        # Process content
        html_content = markdown.markdown(content)
        
        # Create a simple HTML template
        blog_html = f"""
        <!DOCTYPE html>
        <html lang="en">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>{title}</title>
            <style>
                body {{ font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }}
                h1, h2, h3 {{ color: #333; }}
                img {{ max-width: 100%; height: auto; border-radius: 8px; margin: 20px 0; }}
                .author {{ font-style: italic; color: #555; margin-bottom: 20px; }}
                .content {{ margin-top: 30px; }}
            </style>
        </head>
        <body>
            <h1>{title}</h1>
            <div class="author">By {author}</div>
            {f'<img src="featured_image.png" alt="{title}">' if image_path else ''}
            <div class="content">{html_content}</div>
        </body>
        </html>
        """
        
        # Create a safe repo name
        repo_name = f"blog-{re.sub(r'[^a-z0-9-]', '-', title.lower())}"
        
        # Initialize the Hugging Face API
        hf_api = HfApi(token=hf_token)
        
        # Create the Space if it doesn't exist
        try:
            hf_api.create_repo(
                repo_id=f"spaces/{repo_name}",
                repo_type="space",
                space_sdk="static"
            )
        except Exception as e:
            print(f"Space might already exist: {e}")
        
        # Upload the HTML file
        hf_api.upload_file(
            path_or_fileobj=blog_html.encode(),
            path_in_repo="index.html",
            repo_id=f"spaces/{repo_name}",
            repo_type="space"
        )
        
        # Upload image if available
        if image_path and os.path.exists(image_path):
            hf_api.upload_file(
                path_or_fileobj=image_path,
                path_in_repo="featured_image.png",
                repo_id=f"spaces/{repo_name}",
                repo_type="space"
            )
        
        return f"Blog published successfully: https://huggingface.co/spaces/{repo_name}"
    except Exception as e:
        return f"Error saving blog: {str(e)}"

# Main app function
def generate_blog(topic, tone, length, author_name, publish_option, linkedin_username=None, linkedin_password=None, hf_token=None):
    status_updates = []
    status_updates.append(f"Generating blog content for topic: {topic}")
    
    # Generate the blog content
    start_time = time.time()
    blog_content = generate_blog_content(topic, tone, length)
    generation_time = time.time() - start_time
    status_updates.append(f"Content generated in {generation_time:.2f} seconds")
    
    # Extract title from content
    title_match = re.search(r'^#\s+(.+)$', blog_content, re.MULTILINE)
    title = title_match.group(1) if title_match else topic
    
    # Generate image
    status_updates.append("Generating featured image...")
    image_path, image_message = generate_featured_image(topic)
    status_updates.append(image_message)
    
    # Handle publishing
    if publish_option == "linkedin" and linkedin_username and linkedin_password:
        status_updates.append("Posting to LinkedIn...")
        linkedin_result = post_to_linkedin(blog_content, image_path, linkedin_username, linkedin_password)
        status_updates.append(linkedin_result)
    
    if publish_option in ["huggingface", "both"] and hf_token:
        status_updates.append("Saving as Hugging Face blog...")
        hf_result = save_as_blog(blog_content, title, author_name, image_path, hf_token)
        status_updates.append(hf_result)
    
    # Clean up
    if image_path and os.path.exists(image_path):
        try:
            os.remove(image_path)
        except:
            pass
    
    return blog_content, title, "\n".join(status_updates)

# Gradio interface
with gr.Blocks(title="Blog Generator & Publisher") as app:
    gr.Markdown("# AI Blog Generator & LinkedIn Publisher")
    gr.Markdown("Generate professional blog content and publish directly to LinkedIn or save to Hugging Face Spaces.")
    
    with gr.Tab("Generate Blog"):
        with gr.Row():
            with gr.Column():
                topic_input = gr.Textbox(label="Blog Topic", placeholder="Enter the topic of your blog post")
                tone_input = gr.Dropdown(
                    label="Writing Tone", 
                    choices=["professional", "casual", "technical", "storytelling"],
                    value="professional"
                )
                length_input = gr.Dropdown(
                    label="Content Length",
                    choices=["short", "medium", "long"],
                    value="medium"
                )
                author_input = gr.Textbox(label="Author Name", placeholder="Your name")
                
                with gr.Accordion("Publishing Options", open=False):
                    publish_option = gr.Radio(
                        label="Publish To",
                        choices=["none", "linkedin", "huggingface", "both"],
                        value="none"
                    )
                    with gr.Group():
                        linkedin_username = gr.Textbox(label="LinkedIn Username", visible=False)
                        linkedin_password = gr.Textbox(label="LinkedIn Password", type="password", visible=False)
                    
                    hf_token = gr.Textbox(label="Hugging Face Token", type="password", visible=False)
                    
                    def update_visibility(option):
                        linkedin_visible = option in ["linkedin", "both"]
                        hf_visible = option in ["huggingface", "both"]
                        return {
                            linkedin_username: gr.update(visible=linkedin_visible),
                            linkedin_password: gr.update(visible=linkedin_visible),
                            hf_token: gr.update(visible=hf_visible)
                        }
                    
                    publish_option.change(update_visibility, inputs=[publish_option], outputs=[linkedin_username, linkedin_password, hf_token])
                
                generate_btn = gr.Button("Generate Blog", variant="primary")
            
            with gr.Column():
                title_output = gr.Textbox(label="Blog Title")
                blog_output = gr.Markdown(label="Blog Content")
                status_output = gr.Textbox(label="Status", lines=5)
        
        generate_btn.click(
            generate_blog,
            inputs=[topic_input, tone_input, length_input, author_input, publish_option, linkedin_username, linkedin_password, hf_token],
            outputs=[blog_output, title_output, status_output]
        )
    
    with gr.Tab("About"):
        gr.Markdown("""
        ## About This Tool
        
        This application uses AI to generate professional blog content that you can publish directly to LinkedIn or save as a Hugging Face Space.
        
        ### Features:
        
        - Generate blog posts on any topic
        - Choose from different writing tones and length options
        - Create featured images automatically
        - Publish directly to LinkedIn
        - Save as a Hugging Face Space blog
        
        ### How to Use:
        
        1. Enter your blog topic
        2. Select your preferred tone and length
        3. Enter your author name
        4. Choose publishing options (if desired)
        5. Click "Generate Blog"
        
        ### Credits:
        
        This app was created using:
        - Hugging Face's Transformers library
        - Mistral and Stable Diffusion models
        - Gradio for the interface
        """)

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