import streamlit as st
from transformers import RobertaTokenizer, T5ForConditionalGeneration
import torch

# =========================================================
# PAGE CONFIG
# =========================================================

st.set_page_config(
    page_title="Multilingual Code Comment Generator",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# =========================================================
# CUSTOM CSS
# =========================================================

st.markdown(
    """
    <style>

    .main {
        background-color: #0e1117;
        color: white;
    }

    textarea {
        font-size: 15px !important;
        font-family: 'Consolas', monospace !important;
    }

    .stButton>button {
        width: 100%;
        background: linear-gradient(90deg, #4F46E5, #9333EA);
        color: white;
        border-radius: 12px;
        height: 3em;
        font-size: 16px;
        font-weight: 600;
        border: none;
    }

    .stButton>button:hover {
        background: linear-gradient(90deg, #4338CA, #7E22CE);
        color: white;
    }

    .result-box {
        background-color: #161b22;
        padding: 20px;
        border-radius: 12px;
        border: 1px solid #30363d;
        color: #e6edf3;
        font-size: 16px;
        line-height: 1.6;
    }

    .header-title {
        font-size: 42px;
        font-weight: 800;
        color: white;
    }

    .sub-text {
        color: #9ca3af;
        font-size: 18px;
    }

    </style>
    """,
    unsafe_allow_html=True
)

# =========================================================
# SIDEBAR
# =========================================================

with st.sidebar:

    st.title("⚡ Project Info")

    st.markdown("---")

    st.markdown("### 🤖 Model")
    st.write("Salesforce CodeT5")

    st.markdown("### 📊 BLEU Score")
    st.write("24.44")

    st.markdown("### 🧠 Architecture")
    st.write("Encoder-Decoder Transformer")

    st.markdown("### 🔥 Frameworks")
    st.write("PyTorch + HuggingFace")

    st.markdown("### 🌐 Supported Languages")
    st.write("Python")
    st.write("Java (Upcoming)")

    st.markdown("---")

    st.markdown("### 🚀 Features")

    st.write("✅ AI-powered code summarization")
    st.write("✅ Automatic docstring generation")
    st.write("✅ Beam-search decoding")
    st.write("✅ Transformer fine-tuning")
    st.write("✅ HuggingFace model deployment")

# =========================================================
# MODEL CONFIG
# =========================================================

MODEL_NAME = "Eren18/multilingual-code-comment-generator-v2"

# =========================================================
# MODEL LOADING
# =========================================================

@st.cache_resource
def load_model():

    tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME)

    model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME)

    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = model.to(device)

    return tokenizer, model, device


tokenizer, model, device = load_model()

# =========================================================
# INFERENCE FUNCTION
# =========================================================

def generate_comment(code):

    inputs = tokenizer(
        code,
        return_tensors="pt",
        truncation=True,
        max_length=256
    ).to(device)

    outputs = model.generate(
        **inputs,
        max_length=128,
        num_beams=4,
        no_repeat_ngram_size=2,
        early_stopping=True
    )

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

    return generated_comment

# =========================================================
# HEADER
# =========================================================

st.markdown(
    """
    <div class='header-title'>
        🤖 Multilingual Code Comment Generator
    </div>
    """,
    unsafe_allow_html=True
)

st.markdown(
    """
    <div class='sub-text'>
        Generate AI-powered human-readable docstrings from raw source code using a fine-tuned CodeT5 Transformer.
    </div>
    """,
    unsafe_allow_html=True
)

st.markdown("---")

# =========================================================
# MAIN LAYOUT
# =========================================================

left_col, right_col = st.columns(2)

# =========================================================
# LEFT PANEL
# =========================================================

with left_col:

    st.subheader("💻 Input Source Code")

    language = st.selectbox(
        "Select Programming Language",
        ["Python", "Java"]
    )

    sample_code = '''def calculate_discount(price, percent):
    return price - (price * percent / 100)'''

    code_input = st.text_area(
        "Paste your function here",
        value=sample_code,
        height=400
    )

    generate_button = st.button("⚡ Generate Comment")

# =========================================================
# RIGHT PANEL
# =========================================================

with right_col:

    st.subheader("📝 Generated Documentation")

    if generate_button:

        if code_input.strip() == "":
            st.error("Please enter source code.")

        else:

            with st.spinner("Generating intelligent documentation..."):

                generated = generate_comment(code_input)

            st.success("Comment generated successfully!")

            st.markdown(
                f"""
                <div class='result-box'>
                {generated}
                </div>
                """,
                unsafe_allow_html=True
            )

            st.download_button(
                label="📥 Download Comment",
                data=generated,
                file_name="generated_comment.txt",
                mime="text/plain"
            )

# =========================================================
# FOOTER
# =========================================================

st.markdown("---")

st.markdown(
    """
    <center>
        <h4>Built With ❤️ Using</h4>
        <p>
            🤗 HuggingFace &nbsp;&nbsp;|&nbsp;&nbsp;
            🔥 PyTorch &nbsp;&nbsp;|&nbsp;&nbsp;
            ⚡ Streamlit &nbsp;&nbsp;|&nbsp;&nbsp;
            🧠 Transformers
        </p>
    </center>
    """,
    unsafe_allow_html=True
)