import os
import gradio as gr
from huggingface_hub import InferenceClient

# ----------------------
# Hugging Face Token
# ----------------------
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set")

# ----------------------
# Initialize Clients
# ----------------------
GEN_MODEL = "HuggingFaceH4/zephyr-7b-beta"  # Chat model
SUM_MODEL = "google/flan-t5-large"          # Text generation model

client_chat = InferenceClient(token=HF_TOKEN)
client_text = InferenceClient(token=HF_TOKEN)

# ----------------------
# Backend Functions
# ----------------------
def generate_tweet(topic, max_length):
    if not topic.strip():
        return "⚠️ Please enter a topic."
    try:
        response = client_chat.chat_completion(
            model=GEN_MODEL,
            messages=[
                {"role": "system", "content": "You are a professional social media copywriter."},
                {"role": "user", "content": f"Write a viral tweet about: {topic}\nKeep under 280 chars, engaging, punchy."}
            ],
            max_tokens=max_length,
            temperature=0.8
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"❌ Error: {str(e)}"

def generate_blog_ideas(topic, max_length):
    if not topic.strip():
        return "⚠️ Please enter a topic."
    try:
        response = client_chat.chat_completion(
            model=GEN_MODEL,
            messages=[
                {"role": "system", "content": "You are an expert content strategist and SEO blog writer."},
                {"role": "user", "content": f"Generate 5 high-quality blog post ideas about: {topic}\nNumbered list, specific, actionable, clickable."}
            ],
            max_tokens=max_length,
            temperature=0.7
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"❌ Error: {str(e)}"

def summarize_text(text, max_length):
    if not text.strip():
        return "⚠️ Please enter text to summarize."
    prompt = f"Summarize the following text into bullet points:\n{text}\nKeep concise, highlight key points."
    try:
        response = client_text.text_generation(
            model=SUM_MODEL,
            prompt=prompt,
            max_new_tokens=max_length,
            temperature=0.5
        )
        return response.strip()
    except Exception as e:
        return f"❌ Error: {str(e)}"

# ----------------------
# Frontend (Gradio Blocks)
# ----------------------
with gr.Blocks(title="AI Productivity Assistant Pro") as app:

    gr.Markdown("# 🤖 AI Productivity Assistant Pro")
    gr.Markdown("Generate Tweets, Blog Ideas, and Summarize Text effortlessly!")

    with gr.Tabs():

        # Tweet Generator
        with gr.Tab("🐦 Tweet Generator"):
            with gr.Column():
                topic_input = gr.Textbox(label="Enter a topic", placeholder="e.g., AI in education")
                max_len_tweet = gr.Slider(50, 280, value=150, step=10, label="Max Length")
                tweet_btn = gr.Button("Generate Tweet")
                tweet_output = gr.Textbox(label="Generated Tweet")
                tweet_btn.click(generate_tweet, inputs=[topic_input, max_len_tweet], outputs=tweet_output)

        # Blog Ideas
        with gr.Tab("📝 Blog Ideas"):
            with gr.Column():
                blog_input = gr.Textbox(label="Enter a topic", placeholder="e.g., Productivity hacks")
                max_len_blog = gr.Slider(100, 300, value=200, step=10, label="Max Length")
                blog_btn = gr.Button("Generate Blog Ideas")
                blog_output = gr.Textbox(label="Generated Blog Ideas")
                blog_btn.click(generate_blog_ideas, inputs=[blog_input, max_len_blog], outputs=blog_output)

        # Summarizer
        with gr.Tab("📄 Summarizer"):
            with gr.Column():
                text_input = gr.Textbox(label="Paste text to summarize", lines=10)
                max_len_sum = gr.Slider(50, 300, value=150, step=10, label="Max Length")
                sum_btn = gr.Button("Summarize")
                sum_output = gr.Textbox(label="Summary")
                sum_btn.click(summarize_text, inputs=[text_input, max_len_sum], outputs=sum_output)

    gr.Markdown("Built with ❤️ using Hugging Face")

app.launch()