import streamlit as st
from dotenv import load_dotenv
import os
from crewai import Agent, Task, Crew
from langchain_google_genai import ChatGoogleGenerativeAI

# Load environment variables
load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# API key validation
if not GOOGLE_API_KEY:
    st.error("Google Gemini API key not found. Please set it in the .env file.")
    st.stop()

# Streamlit config
st.set_page_config(page_title="NeuraNexus: AI-Powered ML Assistant", page_icon="🧠", layout="wide")

# ==========================
# CSS & Animation
# ==========================
st.markdown("""
<style>
/* (Only showing part of the long CSS here to avoid repetition. Replace with your full custom CSS/JS from the previous version) */

body {
    background-color: #0F1C2E;
    color: #FFFFFF;
    font-family: 'Roboto', sans-serif;
}

h1 {
    font-family: 'Orbitron', sans-serif;
    color: #FF6B6B;
    text-align: center;
    font-size: 2.5em;
    text-shadow: 0 0 10px rgba(255, 107, 107, 0.7);
}

.result-container {
    background: rgba(255, 255, 255, 0.05);
    border: 1px solid rgba(74, 144, 226, 0.2);
    border-radius: 15px;
    padding: 20px;
    margin-top: 30px;
    color: #FFFFFF;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.2);
    backdrop-filter: blur(10px);
}

.by-theaimart {
    text-align: center;
    font-size: 20px;
    color: #50E3C2;
    margin-top: 20px;
    font-weight: bold;
}
</style>

<!-- Optional neural animation background -->
<div class="neural-network" id="neuralNetwork"></div>
<script>
function createNeuralNetwork() {
    const container = document.getElementById('neuralNetwork');
    const numNeurons = 20;
    const numSynapses = 30;
    for (let i = 0; i < numNeurons; i++) {
        const neuron = document.createElement('div');
        neuron.classList.add('neuron');
        neuron.style.left = `${Math.random() * 100}%`;
        neuron.style.top = `${Math.random() * 100}%`;
        neuron.style.width = `${Math.random() * 10 + 5}px`;
        neuron.style.height = neuron.style.width;
        container.appendChild(neuron);
    }
    for (let i = 0; i < numSynapses; i++) {
        const synapse = document.createElement('div');
        synapse.classList.add('synapse');
        synapse.style.left = `${Math.random() * 100}%`;
        synapse.style.top = `${Math.random() * 100}%`;
        synapse.style.width = `${Math.random() * 200 + 50}px`;
        synapse.style.transform = `rotate(${Math.random() * 360}deg)`;
        container.appendChild(synapse);
    }
}
document.addEventListener('DOMContentLoaded', createNeuralNetwork);
</script>
""", unsafe_allow_html=True)

# ==========================
# UI
# ==========================
st.markdown('<h1 class="glow">NeuraNexus: AI-Powered ML Assistant</h1>', unsafe_allow_html=True)
st.markdown('<h3 style="text-align:center;">Describe your ML challenge, and let NeuraNexus craft an innovative solution.</h3>', unsafe_allow_html=True)

problem_description = st.text_area("", height=150, placeholder="Enter your ML challenge here...")
analyze_button = st.button("SYNTHESIZE")

if 'analysis_result' not in st.session_state:
    st.session_state.analysis_result = ""

# ==========================
# Handle Analysis
# ==========================
if analyze_button:
    if not problem_description:
        st.warning("Please describe your ML challenge before synthesizing.")
    else:
        with st.spinner("NeuraNexus is synthesizing your solution..."):
            try:
                # Gemini LLM via LangChain
                llm = ChatGoogleGenerativeAI(
                    model="gemini-pro",
                    google_api_key=GOOGLE_API_KEY,
                    temperature=0.2
                )

                # Define Agent
                agent = Agent(
                    role="NeuraNexus - ML Architect",
                    goal="Design innovative ML strategies",
                    backstory="You are an advanced AI agent trained in modern and futuristic machine learning problem-solving.",
                    verbose=True,
                    allow_delegation=False,
                    llm=llm
                )

                # Define Task
                task = Task(
                    description=f"Analyze the following ML challenge and provide a detailed solution strategy: {problem_description}",
                    expected_output="An innovative, clear, and actionable ML solution strategy.",
                    agent=agent
                )

                # Crew
                crew = Crew(
                    agents=[agent],
                    tasks=[task],
                    verbose=False
                )

                result = crew.kickoff()
                st.session_state.analysis_result = str(result)
                st.success("Synthesis complete!")

            except Exception as e:
                st.session_state.analysis_result = f"❌ Error: {str(e)}"
                st.error("Gemini synthesis failed.")

# ==========================
# Display Output
# ==========================
if st.session_state.analysis_result:
    st.markdown("### NeuraNexus Synthesis Result")
    st.markdown(f"""
    <div class="result-container">
        {st.session_state.analysis_result}
    </div>
    """, unsafe_allow_html=True)

st.markdown('<p class="by-theaimart">By Theaimart</p>', unsafe_allow_html=True)

def main():
    pass

if __name__ == "__main__":
    main()