import streamlit as st
import pinecone
import os
import subprocess
import json
import requests
from pinecone import Pinecone, ServerlessSpec
#from streamlit_extras.add_vertical_space import add_vertical_space

# Set page config
st.set_page_config(page_title="StarCoder AI Code Generator", page_icon="🤖", layout="wide")

# Custom CSS for Grok-like UI
st.markdown("""
    <style>
        body {background-color: #121212; color: white;}
        .stApp {background: #181818; color: white; font-family: 'Arial', sans-serif;}
        .stTextArea textarea {background: #282828; color: white; border-radius: 10px;}
        .stTextInput input {background: #282828; color: white; border-radius: 10px;}
        .stButton>button {background: #3f51b5; color: white; border-radius: 8px; padding: 10px;}
        .stDownloadButton>button {background: #4caf50; color: white; border-radius: 8px; padding: 10px;}
        .stCode {background: #202020; padding: 15px; border-radius: 10px;}
        .stSidebar {background: #202020; color: white;}
        .stSidebar .stButton>button {background: #ff9800; color: white;}
    </style>
""", unsafe_allow_html=True)

# Pinecone API Key
PINECONE_API_KEY = "your pinecone api key here"
pc = Pinecone(api_key=PINECONE_API_KEY)
INDEX_NAME = "code-gen"
if INDEX_NAME not in pc.list_indexes().names():
    pc.create_index(name=INDEX_NAME, dimension=1536, metric='euclidean',
                    spec=ServerlessSpec(cloud='aws', region='us-east-1'))
index = pc.Index(INDEX_NAME)

# Hugging Face API Key
HUGGINGFACE_API_KEY = "your huggingface api key"

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "current_chat" not in st.session_state:
    st.session_state.current_chat = []


def new_chat():
    if st.session_state.current_chat:
        st.session_state.chat_history.append(st.session_state.current_chat)
    st.session_state.current_chat = []


# Sidebar - Chat History
st.sidebar.title("💬 Chat History")
for i, chat in enumerate(st.session_state.chat_history):
    if st.sidebar.button(f"Chat {i + 1} 🕒"):
        st.session_state.current_chat = chat

# Sidebar - New Chat Button
if st.sidebar.button("New Chat ➕"):
    new_chat()


# Function to query Hugging Face model
def generate_code(prompt, examples=[]):
    formatted_prompt = "".join(examples) + "\n" + prompt
    headers = {"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"}
    data = {"inputs": formatted_prompt, "parameters": {"temperature": 0.5, "max_length": 512}}
    response = requests.post("https://api-inference.huggingface.co/models/bigcode/starcoder2-15b", headers=headers,
                             json=data)
    return response.json()[0][
        'generated_text'].strip() if response.status_code == 200 else "Error: Unable to generate code"


def execute_python_code(code):
    try:
        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, timeout=5)
        return result.stdout if result.returncode == 0 else result.stderr
    except Exception as e:
        return str(e)


def store_in_pinecone(prompt, code):
    vector = [0.1] * 1536  # Placeholder vector
    index.upsert([(prompt, vector, {"code": code})])


def retrieve_from_pinecone(prompt):
    response = index.query(vector=[0.1] * 1536, top_k=2, include_metadata=True)
    return [r["metadata"]["code"] for r in response["matches"]]


# Main UI
st.title("🤖 StarCoder AI Code Generator")

# User Input
prompt = st.text_area("✍️ Enter your prompt:", height=100)
mode = st.selectbox("⚡ Choose prompting mode", ["One-shot", "Two-shot"])

generate_button = st.button("🚀 Generate Code")

if generate_button and prompt:
    examples = retrieve_from_pinecone(prompt) if mode == "Two-shot" else []
    generated_code = generate_code(prompt, examples)
    st.session_state.current_chat.append((prompt, generated_code))
    st.subheader("📜 Generated Code")
    st.code(generated_code, language="python")
    store_in_pinecone(prompt, generated_code)

    execute_button = st.button("▶️ Run Code")
    if execute_button:
        execution_result = execute_python_code(generated_code)
        st.subheader("🖥️ Execution Output")
        st.text(execution_result)

    st.download_button("📥 Download Code", generated_code, file_name="generated_code.py")