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

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

st.title("💻 AI Code Generator (Stable Version)")

# ==============================
# LOAD MODEL (SAFE)
# ==============================
@st.cache_resource
def load_model():
    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"  # ✅ HF Free Safe

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32
    )

    model.eval()
    return tokenizer, model

# Load with spinner (IMPORTANT)
with st.spinner("🔄 Loading AI model... Please wait"):
    tokenizer, model = load_model()

st.success("✅ Model Loaded Successfully")

# ==============================
# CLEAN OUTPUT
# ==============================
def clean_code(code):
    code = code.strip()

    # Remove unwanted text
    remove_words = [
        "Explanation:",
        "Here is the code:",
        "Output:",
        "Answer:"
    ]

    for word in remove_words:
        code = code.replace(word, "")

    # Remove markdown
    code = code.replace("```python", "").replace("```", "")

    return code.strip()

# ==============================
# GENERATE CODE
# ==============================
def generate_code(prompt, language):

    full_prompt = f"""
You are an expert {language} programmer.

Write clean, correct, and complete code.

Rules:
- Only return code
- No explanations
- Complete solution

Task:
{prompt}

Code:
"""

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

    try:
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=200,
                do_sample=True,
                temperature=0.3,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id
            )

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

        if "Code:" in result:
            result = result.split("Code:")[-1]

        return clean_code(result)

    except Exception as e:
        return f"# ERROR: {str(e)}"

# ==============================
# SESSION STATE
# ==============================
if "history" not in st.session_state:
    st.session_state.history = []

# ==============================
# UI INPUT
# ==============================
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..."):
            code = generate_code(user_prompt, language)

        st.session_state.history.append((user_prompt, code))

# ==============================
# OUTPUT DISPLAY
# ==============================
if st.session_state.history:
    st.subheader("📌 Generated Results")

    for q, c in reversed(st.session_state.history):
        st.markdown(f"**🧑 Task:** {q}")
        st.code(c, language=language.lower())