import streamlit as st
import requests
import os

# Ensure the Hugging Face API Token is available
API_TOKEN = os.getenv("HF_API_TOKEN")
if not API_TOKEN:
    st.error("⚠️ API Token is missing! Please set HF_API_TOKEN as an environment variable.")
    st.stop()

# Define Hugging Face Model and 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):
    """Translate code between languages using Hugging Face API."""
    prompt = f"""You are an expert AI in code translation. Convert the following {source_lang} code to {target_lang}:
    
    {source_lang} Code:
    ```{source_lang.lower()}
    {code_snippet}
    ```

    Now, provide the equivalent {target_lang} code:
    """

    try:
        response = requests.post(API_URL, headers=HEADERS, json={"inputs": prompt})

        # Handle API response
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and result:
                generated_text = result[0].get("generated_text", "")
                translated_code = generated_text.split(f"```{target_lang.lower()}")[-1].strip("```").strip()
                return translated_code
            else:
                return "⚠️ Error: Unexpected API response format."

        elif response.status_code == 400:
            return "⚠️ Error: Bad request. Please check your input."

        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 model access."

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

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

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

# Streamlit UI
st.title("🔄 AI Code Translator")
st.write("Convert code between Python, Java, C++, and C.")

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:", height=200)

if st.button("Translate"):
    if source_lang == target_lang:
        st.warning("⚠️ Source and target languages must be different!")
    elif not code_input.strip():
        st.warning("⚠️ Please enter some code before translating.")
    else:
        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())