import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ==============================
# 🔐 HIDE STREAMLIT MENU
# ==============================
st.markdown("""
<style>
#MainMenu {visibility: hidden;}
header {visibility: hidden;}
footer {visibility: hidden;}
.stDeployButton {display:none;}
</style>
""", unsafe_allow_html=True)

# ==============================
# PAGE CONFIG
# ==============================
st.set_page_config(page_title="💻 AI Code Generator", layout="wide")

# ==============================
# LOAD MODEL (FAST + STABLE)
# ==============================
@st.cache_resource
def load_model():
    model_name = "microsoft/phi-2"

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32,   # CPU friendly
        device_map=None
    )

    return tokenizer, model

tokenizer, model = load_model()

# ==============================
# CODE GENERATION FUNCTION
# ==============================
def generate_code(prompt, language):

    full_prompt = f"""### Instruction:
Write ONLY valid {language} code.

Do not include explanations.
Do not include special tokens.

Task:
{prompt}

### Response:
"""

    inputs = tokenizer(full_prompt, return_tensors="pt")

    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        temperature=0.2,
        top_p=0.9,
        do_sample=False
    )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract response
    if "### Response:" in result:
        result = result.split("### Response:")[-1]

    # Remove unwanted tokens
    unwanted_tokens = [
        "<|endoftext|>",
        "<|file_separator|>",
        "<|assistant|>",
        "<|system|>"
    ]

    for token in unwanted_tokens:
        result = result.replace(token, "")

    return result.strip()

# ==============================
# UI
# ==============================
st.title("💻 AI Code Generator (Fast & Clean)")

col1, col2 = st.columns(2)

with col1:
    user_prompt = st.text_area("Describe your task", height=200)

with col2:
    language = st.selectbox(
        "Select Programming Language",
        ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
    )

# ==============================
# GENERATE BUTTON
# ==============================
if st.button("🚀 Generate Code"):
    if not user_prompt.strip():
        st.warning("Please enter a task")
    else:
        with st.spinner("Generating code..."):
            try:
                code = generate_code(user_prompt, language)

                st.success("✅ Generated Code")
                st.code(code, language=language.lower())

            except Exception as e:
                st.error(f"❌ Error: {str(e)}")