import os
import gradio as gr
from huggingface_hub import InferenceClient

# ----------------------
# TOKEN HANDLING (COLAB + SPACES)
# ----------------------
try:
    from google.colab import userdata
    HF_TOKEN = userdata.get("HF_TOKEN")   # Colab
except:
    HF_TOKEN = os.getenv("HF_TOKEN")      # Spaces

if not HF_TOKEN:
    raise ValueError("HF_TOKEN not found")

# ----------------------
# MODELS
# ----------------------
GEN_MODEL = "tiiuae/falcon-7b-instruct"
SUM_MODEL = "google/flan-t5-large"

client = InferenceClient(token=HF_TOKEN)

# ----------------------
# BACKEND FUNCTIONS
# ----------------------
def generate_tweet(topic, tone, max_length, history):
    if not topic.strip():
        return "⚠️ Please enter a topic.", history

    prompt = f"""
You are a world-class viral Twitter copywriter.

Write 3 highly engaging tweets about: {topic}

Rules:
- Tone: {tone}
- Strong hook
- Human-like writing
- Max 280 characters each
- Add 1-2 relevant hashtags
- Make them scroll-stopping
"""

    try:
        response = client.text_generation(
            model=GEN_MODEL,
            prompt=prompt,
            max_new_tokens=max_length,
            temperature=0.8
        )
        result = response.strip()
        history.append(f"Tweet | {topic}\n{result}\n")
        return result, history
    except Exception as e:
        return f"❌ Error: {str(e)}", history


def generate_blog_ideas(topic, tone, max_length, history):
    if not topic.strip():
        return "⚠️ Please enter a topic.", history

    prompt = f"""
You are an expert SEO strategist.

Generate 5 blog ideas about: {topic}

Rules:
- Tone: {tone}
- Specific and niche-focused
- Clear audience
- Click-worthy titles
"""

    try:
        response = client.text_generation(
            model=GEN_MODEL,
            prompt=prompt,
            max_new_tokens=max_length,
            temperature=0.7
        )
        result = response.strip()
        history.append(f"Blog | {topic}\n{result}\n")
        return result, history
    except Exception as e:
        return f"❌ Error: {str(e)}", history


def summarize_text(text, max_length, history):
    if not text.strip():
        return "⚠️ Please enter text.", history

    prompt = f"""
Summarize into:
- Bullet points
- Key insights
- Actionable takeaways

{text}
"""

    try:
        response = client.text_generation(
            model=SUM_MODEL,
            prompt=prompt,
            max_new_tokens=max_length,
            temperature=0.5
        )
        result = response.strip()
        history.append(f"Summary\n{result}\n")
        return result, history
    except Exception as e:
        return f"❌ Error: {str(e)}", history


def copy_text(text):
    return gr.update(value=text)

def show_history(history):
    return "\n".join(history)

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

    history_state = gr.State([])

    gr.Markdown("# 🤖 AI Productivity Assistant Pro")
    gr.Markdown("### Generate high-quality content in seconds")

    with gr.Tabs():

        # Tweet Generator
        with gr.Tab("🐦 Tweet Generator"):
            gr.Markdown("Generate viral tweets instantly")

            with gr.Row():
                with gr.Column():
                    topic = gr.Textbox(label="Topic")
                    tone = gr.Dropdown(
                        ["Professional", "Casual", "Funny", "Persuasive"],
                        value="Professional",
                        label="Tone"
                    )
                    max_len = gr.Slider(50, 300, value=150, label="Max Length")
                    btn = gr.Button("Generate")

                with gr.Column():
                    output = gr.Textbox(lines=10, label="Output")
                    copy_btn = gr.Button("Copy Output")

            btn.click(generate_tweet, [topic, tone, max_len, history_state], [output, history_state])
            copy_btn.click(copy_text, output, output)

        # Blog Ideas
        with gr.Tab("📝 Blog Ideas"):
            gr.Markdown("Generate high-converting blog ideas")

            with gr.Row():
                with gr.Column():
                    topic_b = gr.Textbox(label="Topic")
                    tone_b = gr.Dropdown(
                        ["Professional", "Casual", "Funny", "Persuasive"],
                        value="Professional",
                        label="Tone"
                    )
                    max_len_b = gr.Slider(100, 400, value=200, label="Max Length")
                    btn_b = gr.Button("Generate")

                with gr.Column():
                    output_b = gr.Textbox(lines=10, label="Output")
                    copy_btn_b = gr.Button("Copy Output")

            btn_b.click(generate_blog_ideas, [topic_b, tone_b, max_len_b, history_state], [output_b, history_state])
            copy_btn_b.click(copy_text, output_b, output_b)

        # Summarizer
        with gr.Tab("📄 Summarizer"):
            gr.Markdown("Summarize text into insights")

            with gr.Row():
                with gr.Column():
                    text = gr.Textbox(lines=8, label="Text")
                    max_len_s = gr.Slider(50, 300, value=150, label="Max Length")
                    btn_s = gr.Button("Summarize")

                with gr.Column():
                    output_s = gr.Textbox(lines=10, label="Output")
                    copy_btn_s = gr.Button("Copy Output")

            btn_s.click(summarize_text, [text, max_len_s, history_state], [output_s, history_state])
            copy_btn_s.click(copy_text, output_s, output_s)

        # History
        with gr.Tab("📜 History"):
            history_box = gr.Textbox(lines=20, label="History")
            refresh_btn = gr.Button("Refresh")
            refresh_btn.click(show_history, history_state, history_box)

    gr.Markdown("---")
    gr.Markdown("Built with Hugging Face | AI SaaS Prototype")

app.launch(share=True)