import gradio as gr
import os
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
pipe = pipeline(
    "text-generation",
    model=MODEL_ID,
    torch_dtype=torch.float32,
    device_map="auto"
)
print("Model loaded.")

SYSTEM = """You are an expert API architect. When given a product idea, recommend the best API stack with pricing and build order."""

def generate(prompt):
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": prompt}
    ]
    formatted = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    out = pipe(
        formatted,
        max_new_tokens=800,
        do_sample=True,
        temperature=0.7,
        return_full_text=False
    )
    return out[0]["generated_text"].strip()


def get_stack(idea, budget, level, market):
    if not idea.strip():
        return "⚠️ Please describe your product idea."
    try:
        prompt = f"""Product idea: {idea}
Budget: {budget} | Level: {level} | Market: {market}

Recommend 3-5 APIs using this format:

## 🎯 Idea Summary
(one line)

## 🔧 API Stack
### [API Name] — [role]
- Pricing: (free tier + paid)
- Docs: (URL)
- Difficulty: Easy/Medium/Hard

## 💰 Cost Estimate
- 100 users: $X/month
- 1,000 users: $X/month

## ⚡ Build Order
1. First: [API]
2. Second: [API]

## 🚨 Top Mistake
(one paragraph)"""
        return generate(prompt)
    except Exception as e:
        return f"❌ Error: {str(e)}"


def ask_followup(question, prev):
    if not question.strip():
        return ""
    if not prev or prev.startswith("⚠️") or prev.startswith("❌"):
        return "⚠️ Generate a recommendation first."
    try:
        prompt = f"Previous recommendation:\n{prev[:500]}\n\nFollow-up: {question}\n\nAnswer concisely."
        return generate(prompt)
    except Exception as e:
        return f"❌ Error: {str(e)}"


EXAMPLES = [
    ["A notes app with AI summarization", "$50/month", "Intermediate", "Students"],
    ["SMS marketing for restaurants", "$30/month", "Beginner", "Restaurant owners"],
    ["AI job matching platform", "$100/month", "Advanced", "Tech recruiters"],
    ["Podcast transcription tool", "$20/month", "Beginner", "Podcasters"],
    ["Crypto portfolio tracker", "$0 (free only)", "Intermediate", "Retail investors"],
]

css = """
.gradio-container { max-width: 900px !important; margin: auto; }
footer { display: none !important; }
"""

with gr.Blocks(css=css, title="API Stack Finder") as demo:
    gr.Markdown("""
# 🔧 API Stack Finder
**Describe your product idea → Get the perfect API stack**

Powered by **TinyLlama** — runs locally inside this Space, no external API needed.
""")

    with gr.Row():
        with gr.Column(scale=3):
            idea_input = gr.Textbox(
                label="Describe your product idea",
                placeholder="e.g. A freelancer marketplace with instant payments...",
                lines=4
            )
        with gr.Column(scale=1):
            budget_input = gr.Dropdown(
                choices=["$0 (free only)", "$10/month", "$30/month",
                         "$50/month", "$100/month", "$500/month", "Unlimited"],
                value="$30/month",
                label="Monthly API budget"
            )
            level_input = gr.Dropdown(
                choices=["Beginner", "Intermediate", "Advanced"],
                value="Intermediate",
                label="Technical level"
            )
            market_input = gr.Textbox(
                label="Target market",
                placeholder="e.g. Small business owners",
                value="General consumers"
            )

    gr.Markdown("**Try an example:**")
    with gr.Row():
        for ex in EXAMPLES:
            gr.Button(ex[0][:30] + "...", size="sm").click(
                lambda e=ex: (e[0], e[1], e[2], e[3]),
                outputs=[idea_input, budget_input, level_input, market_input]
            )

    submit_btn = gr.Button("🔍 Find My API Stack", variant="primary", size="lg")
    output = gr.Markdown("*Your API stack will appear here...*")

    gr.Markdown("---")
    gr.Markdown("### 💬 Ask a follow-up")
    with gr.Row():
        followup_input = gr.Textbox(
            label="Follow-up question",
            placeholder="e.g. Is there a free alternative to Stripe?",
            scale=4
        )
        followup_btn = gr.Button("Ask", scale=1, variant="secondary")
    followup_output = gr.Markdown()

    submit_btn.click(get_stack,
        inputs=[idea_input, budget_input, level_input, market_input],
        outputs=output)
    followup_btn.click(ask_followup,
        inputs=[followup_input, output],
        outputs=followup_output)

    gr.Markdown("---\n🤖 Model: `TinyLlama-1.1B-Chat` · Runs fully inside HuggingFace Space")

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