import streamlit as st
import os
import asyncio
import time
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain.prompts import PromptTemplate
from langchain.schema import AIMessage, HumanMessage

# ------------------------------- CONFIG -------------------------------
st.set_page_config(page_title="Intelligent Coding Agent", layout="wide")

# Inject custom CSS for full-width layout and animations
st.markdown("""
    <style>
        .appview-container .main .block-container {
            max-width: 100%;
            padding: 0 2rem;
        }
        .stTextArea, .stTextInput, .stButton {
            width: 100% !important;
        }
        .stCodeBlock {
            white-space: pre-wrap !important;
            width: 100%;
        }
        .branding {
            text-align: center;
            font-weight: bold;
            font-size: 14px;
            color: #888;
            padding: 10px;
        }
        .loader {
            text-align: center;
            font-size: 18px;
            font-weight: bold;
            animation: fadeIn 1s infinite alternate;
        }
        @keyframes fadeIn {
            from {opacity: 0.3;}
            to {opacity: 1;}
        }
    </style>
""", unsafe_allow_html=True)

# ------------------------------- CACHING -------------------------------
@st.cache_resource
def get_openai_model():
    """Load OpenAI GPT-3.5-Turbo model once (cached) for efficiency."""
    return ChatOpenAI(model_name="gpt-4o-mini-2024-07-18", temperature=0.7)

@st.cache_resource
def get_vectorstore():
    """Load ChromaDB once (cached) for fast retrieval."""
    return Chroma(
        embedding_function=OpenAIEmbeddings(),
        persist_directory="/home/user/chroma_db"
    )

# Initialize models
llm = get_openai_model()
vectorstore = get_vectorstore()

# ------------------------------- PROMPT TEMPLATE -------------------------------
code_prompt = PromptTemplate(
    input_variables=["language", "description"],
    template="""
    You are an **expert AI coding assistant** specializing in **ALL programming languages**, including:
    - Python, JavaScript, TypeScript, PHP, C#, Java, C++, Go, Rust, Swift, Kotlin, R, SQL, Bash
    - .NET, .NET Core, ASP.NET, Node.js, Django, Flask, Spring Boot, Ruby on Rails
    - Web technologies: HTML, CSS, WebAssembly, WebSockets, Web3.js, GraphQL, REST APIs

    **Task:** Generate a well-structured, optimized, and fully functional code snippet.

    **Language:** {language}  
    **User Request:** {description}  

    - Follow industry best practices and clean code principles.
    - Include relevant comments for clarity.
    - If applicable, provide setup instructions.
    - Ensure security and performance optimizations.
    """
)

# ------------------------------- STREAMLIT UI -------------------------------
st.title("🧠💻 Intelligent Coding Agent")
st.markdown("### Generate, Enhance, and Modify Code Seamlessly!")

languages = [
    "Python", "JavaScript", "TypeScript", "PHP", "C#", "Java", "C++", "Go", "Rust", "Swift", 
    "Kotlin", "R", "SQL", "Bash", "HTML/CSS/JS", ".NET", ".NET Core"
]
language = st.selectbox("Select a programming language:", languages)
description = st.text_area("What do you need?", placeholder="E.g., Build an authentication system in Flask", height=150)

# ------------------------------- SESSION STATE -------------------------------
if "generated_code" not in st.session_state:
    st.session_state.generated_code = ""

# ------------------------------- ANIMATED LOADER -------------------------------
def show_loader():
    """Shows an animated loading effect."""
    loader_text = "⏳ Generating your code..."
    for _ in range(3):
        st.markdown(f"<div class='loader'>{loader_text}</div>", unsafe_allow_html=True)
        time.sleep(0.6)

# ------------------------------- ASYNC FUNCTIONS -------------------------------
async def generate_code():
    """Generates optimized code while ensuring proper formatting."""
    if not description.strip():
        st.warning("Please enter a description!")
        return
    
    show_loader()

    # Format prompt
    prompt = code_prompt.format(language=language, description=description)
    
    # Generate code asynchronously
    response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])

    # Store generated code in session state
    st.session_state.generated_code = response.content

    # Display properly formatted code block
    st.subheader("Here’s your generated code:")
    st.code(st.session_state.generated_code, language=language.lower())

    # Save to vectorstore for reference
    vectorstore.add_texts([st.session_state.generated_code])

async def enhance_code():
    """Enhances existing code with better efficiency and readability."""
    if not st.session_state.generated_code:
        st.warning("Generate code first before enhancing!")
        return

    st.info("🔄 Improving your code...")

    prompt = f"Improve the following code for better efficiency and readability:\n\n{st.session_state.generated_code}"
    response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])

    st.session_state.generated_code = response.content
    st.subheader("Here’s the improved version:")
    st.code(st.session_state.generated_code, language=language.lower())

async def modify_code():
    """Modifies code based on user feedback."""
    modification = st.text_area("What changes do you want?", placeholder="E.g., Add input validation")

    if st.button("Apply Changes"):
        if not st.session_state.generated_code:
            st.warning("Generate code first before modifying!")
            return
        if not modification.strip():
            st.warning("Please enter modification details!")
            return

        st.info("🔄 Making modifications...")

        prompt = f"Modify the following code based on the request '{modification}':\n\n{st.session_state.generated_code}"
        response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])

        st.session_state.generated_code = response.content
        st.subheader("Here’s the modified code:")
        st.code(st.session_state.generated_code, language=language.lower())

# ------------------------------- BUTTONS -------------------------------
# ------------------------------- BUTTONS -------------------------------
col1, col2 = st.columns(2)  # Ensures both buttons are of equal width

with col1:
    if st.button("Generate Code", use_container_width=True):
        asyncio.run(generate_code())

with col2:
    if st.button("Enhance Code", use_container_width=True):
        asyncio.run(enhance_code())

# ------------------------------- DOWNLOAD OPTION -------------------------------
if st.session_state.generated_code:
    st.download_button("Download Code", st.session_state.generated_code, "generated_code.txt", "text/plain", use_container_width=True)

# ------------------------------- BRANDING -------------------------------
st.markdown('<p class="branding">🚀 Powered by WaysAhead Global</p>', unsafe_allow_html=True)

st.markdown("---")
st.markdown("🔹 **Optimized with OpenAI GPT-3.5, LangChain, and ChromaDB** 🔹")