import gradio as gr
import os
import requests

# Use Inference API with LLM only
MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
token = os.getenv("HUGGING_FACE_HUB_TOKEN")
ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"

def improve_description(user_input: str) -> str:
    """
    Rewrites a rough project description using LLM via router.huggingface.co OpenAI-compatible API.
    """
    if not user_input or not user_input.strip():
        return "Please provide a description to improve."

    user_text = user_input.strip()
    if len(user_text) < 10:
        return "The description is too short. Please add more details about your project."

    if not token:
        return "Error: HUGGING_FACE_HUB_TOKEN environment variable is not set. Please configure your Hugging Face token."

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": MODEL_ID,
        "messages": [
            {"role": "system", "content": "You are a technical writing assistant. Rewrite project descriptions to be clear, professional, and concise while preserving all technical details. Limit your output to a maximum of 3 sentences."},
            {"role": "user", "content": f"Rewrite this project description professionally, concisely, and in no more than 3 sentences:\n\n{user_text}"}
        ],
        "max_tokens": 300,
        "temperature": 0.3
    }
    try:
        resp = requests.post(ROUTER_URL, headers=headers, json=payload, timeout=60)
        if resp.status_code != 200:
            return f"Error: API request failed - {resp.status_code} {resp.reason}: {resp.text}"
        data = resp.json()
        improved_text = data["choices"][0]["message"]["content"].strip()
        return improved_text if improved_text else "Error: No response generated."
    except Exception as e:
        return f"Error: An unexpected error occurred - {str(e)}"


# --- Gradio Interface (no changes needed here) ---
with gr.Blocks(title="Technical Description Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 📝 Technical Description Assistant
        Instantly transform your rough project notes into a polished, professional
        technical description using a zero-shot AI model.
        """
    )

    with gr.Row():
        input_text = gr.Textbox(
            label="Your Rough Description",
            placeholder="e.g., 'my app uses python and ml stuff to make pics look better'",
            lines=10,
            scale=1,
        )
        output_text = gr.Textbox(
            label="Polished Technical Description",
            lines=10,
            scale=1,
        )

    submit_btn = gr.Button("✨ Generate Polished Description", variant="primary")

    gr.Markdown("### A few examples to get you started:")
    gr.Examples(
        examples=[
            "This is a web app that lets users upload pics and it uses AI to make them look better. Built with Python and some ML stuff.",
            "Made a tool that checks code for bugs. It's really fast and works with multiple languages. Uses AST parsing.",
            "API for weather data. Gets info from multiple sources and combines them. Has caching so it's faster.",
            "This project uses docker and cloud stuff to run machines.",
        ],
        inputs=input_text,
        outputs=output_text,
        fn=improve_description,
    )

    submit_btn.click(
        fn=improve_description,
        inputs=input_text,
        outputs=output_text,
        api_name="improve_description"
    )

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