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

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

# ==============================
# LOAD MODEL
# ==============================
@st.cache_resource
def load_model():
    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"   # 🔥 BEST MODEL

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32,
        device_map="auto"   # auto uses GPU if available
    )

    return tokenizer, model

tokenizer, model = load_model()

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

    remove_words = [
        "Explanation:",
        "Here is the code:",
        "Output:",
        "Answer:"
    ]

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

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

    return code.strip()

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

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

Generate clean, correct, and complete code.

Rules:
- Only return code
- No explanation
- Proper syntax
- Complete working solution

Task:
{prompt}

Code:
"""

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

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=350,
            do_sample=True,
            temperature=0.2,
            top_p=0.9,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id,
            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)

# ==============================
# SESSION STATE (CHAT HISTORY)
# ==============================
if "history" not in st.session_state:
    st.session_state.history = []

# ==============================
# UI
# ==============================
st.title("💻 AI Code Generator (Advanced)")

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 high-quality code..."):
            code = generate_code(user_prompt, language)

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

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