import streamlit as st
import os
import logging
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain.memory import ConversationBufferWindowMemory
import numpy as np

# Set up logging to help debug
logging.basicConfig(filename='app.log', level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')

# Check for Groq API key
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    st.error("Oops! The Groq API key is missing. Please add it in Hugging Face Secrets as GROQ_API_KEY.")
    logging.error("GROQ_API_KEY not found in environment variables.")
    st.stop()

# Initialize Groq API and memory
try:
    llm = ChatGroq(model="llama3-70b-8192", groq_api_key=groq_api_key, temperature=0.3)
    memory = ConversationBufferWindowMemory(k=3)
    logging.info("Successfully initialized Groq API and memory.")
except Exception as e:
    st.error(f"Error connecting to Groq: {str(e)}. Check your API key or try again later.")
    logging.error(f"Groq initialization failed: {str(e)}")
    st.stop()

# System prompt
system_prompt = """
You are an expert Electrical Engineering Assistant for professionals, students, and hobbyists. Answer questions about electrical circuits, components, and wiring with clear, accurate, and technical responses.

Capabilities:
1. Explain concepts like Ohm’s Law or Kirchhoff’s Laws with examples.
2. Guide users through circuit design step-by-step.
3. Recommend components (e.g., E12 series resistors).
4. Solve circuit unknowns with calculations.
5. Suggest safety measures (e.g., fuses, grounding).

Use this format:
---
**User Query Summary**: <Summarize the question>
**Technical Analysis**: <Explain or calculate step-by-step>
**Suggested Components/Practices**: <Recommend parts or methods>
**Further Reading or Simulation Advice**: <Suggest tools or learning resources>

Always show calculation steps if values are provided (e.g., voltage=12V, R1=100 ohms).
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}")
])
chain = prompt | llm

# Streamlit UI
st.title("⚡️ Electrical Engineering Assistant")
st.markdown("Ask about circuits, like 'Design a voltage divider for 12V to 5V'!")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []
    logging.info("Initialized chat history.")

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Clear chat history button
if st.button("Clear Chat History"):
    st.session_state.messages = []
    memory.clear()
    logging.info("Chat history cleared.")

# Input box
user_input = st.chat_input("Ask your question (e.g., 'What’s Ohm’s Law?')")

if user_input:
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)
    logging.info(f"Received user input: {user_input}")

    response_content = ""
    try:
        # Local calculation for voltage divider
        if "voltage divider" in user_input.lower():
            words = user_input.split()
            vin = None
            vout = None
            for word in words:
                if word.endswith("V"):
                    try:
                        value = float(word[:-1])
                        if vin is None:
                            vin = value
                        else:
                            vout = value
                    except ValueError:
                        pass
            if vin and vout and vout < vin:
                e12_series = [10, 12, 15, 18, 22, 27, 33, 39, 47, 56, 68, 82]
                ratio = vout / vin
                r1 = 1000
                r2 = r1 * ratio / (1 - ratio)
                r2 = min(e12_series, key=lambda x: abs(x * 1000 - r2)) * 1000
                vout_actual = vin * r2 / (r1 + r2)
                response_content += f"""
**Local Calculation**:
- Formula: Vout = Vin * R2 / (R1 + R2)
- Vin = {vin}V, Target Vout = {vout}V
- R1 = {r1}Ω, R2 = {r2}Ω (E12 series)
- Actual Vout = {vout_actual:.2f}V
"""
                logging.info(f"Performed local voltage divider calculation: Vin={vin}, Vout={vout}")
        # Get Groq response
        with st.spinner("Thinking..."):
            context = memory.load_memory_variables({})["history"]
            full_input = f"{context}\n\nUser: {user_input}"
            response = chain.invoke({"input": full_input})
            response_content += response.content
            memory.save_context({"input": user_input}, {"output": response.content})
            logging.info("Successfully received Groq response.")
    except Exception as e:
        response_content = f"Error: {str(e)}. Try again or check your question!"
        logging.error(f"Error processing input: {str(e)}")

    st.session_state.messages.append({"role": "assistant", "content": response_content})
    with st.chat_message("assistant"):
        st.markdown(response_content)