import streamlit as st
import requests
import os

# Load API token from environment
API_TOKEN = os.getenv("HF_API_TOKEN")

if not API_TOKEN:
    st.error("⚠️ Hugging Face API token is missing! Set `HF_API_TOKEN` in your environment variables.")
    st.stop()  # Stop execution if the token is missing

# Define model API endpoint
MODEL_ID = "bigcode/starcoder"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}

def translate_code(code_snippet, source_lang, target_lang):
    """Translates code from one language to another using Hugging Face API."""
    prompt = f"""### Task: Convert {source_lang} code to {target_lang}.
    
    {source_lang} Code:
    ```{source_lang.lower()}
    {code_snippet}
    ```

    Now convert it to {target_lang}:

    ```{target_lang.lower()}
    """

    try:
        response = requests.post(API_URL, headers=HEADERS, json={
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": 200,
                "temperature": 0.2,
                "top_k": 50,
            }
        })

        if response.status_code == 200:
            output = response.json()

            if isinstance(output, list) and len(output) > 0:
                generated_text = output[0].get("generated_text", "")

                # Extract translated code only
                translated_code = generated_text.split(f"```{target_lang.lower()}")[-1].strip()
                translated_code = translated_code.replace("```", "").strip()

                return translated_code if translated_code else "⚠️ Translation failed. No valid output received."

            else:
                return "⚠️ Unexpected response format from API."

        elif response.status_code == 400:
            return "⚠️ Error: Invalid request. Check input format."

        elif response.status_code == 401:
            return "⚠️ Error: Unauthorized. Check your API token."

        elif response.status_code == 403:
            return "⚠️ Error: Access forbidden. You may need special access to this model."

        elif response.status_code == 503:
            return "⚠️ Error: Model is loading. Please wait and try again."

        else:
            return f"⚠️ Error {response.status_code}: {response.text}"

    except requests.exceptions.RequestException as e:
        return f"⚠️ Network Error: {str(e)}"

# Streamlit UI
st.title("🔄 Code Translator using StarCoder")
st.write("Translate code between different programming languages using AI.")

# Define language options
languages = ["Python", "Java", "C", "C++"]

source_lang = st.selectbox("Select source language", languages)
target_lang = st.selectbox("Select target language", languages)
code_input = st.text_area("Enter your code here:", height=200)

if st.button("Translate"):
    if source_lang == target_lang:
        st.warning("⚠️ Source and target languages cannot be the same.")
    elif code_input.strip():
        with st.spinner("Translating..."):
            translated_code = translate_code(code_input, source_lang, target_lang)
            st.subheader(f"Translated {target_lang} Code:")
            st.code(translated_code, language=target_lang.lower())
    else:
        st.warning("⚠️ Please enter some code before translating.")