import os
import streamlit as st
from groq import Groq
from transformers import AutoTokenizer
import faiss
import numpy as np

# Set the Hugging Face API Key
GrOQ_API_KEY = "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"
client = Groq(api_key=GrOQ_API_KEY)

# Supported Microcontrollers/Processors
devices = {
    "Arduino": {
        "template": """
        // Arduino Syntax Template (for Arduino Uno)
        void setup() {
          // Setup code here, to run once:
          pinMode(13, OUTPUT);
        }
        void loop() {
          // Main code here, to run repeatedly:
          digitalWrite(13, HIGH);   // Turn on LED
          delay(1000);               // Wait for a second
          digitalWrite(13, LOW);    // Turn off LED
          delay(1000);               // Wait for a second
        }
        """,
    },
    "PIC18": {
        "template": """
        // PIC18 Syntax Template
        // C Language:
        #include <xc.h>
        void main() {
          TRISB = 0; // Set PORTB as output
          while(1) {
            LATB = 0xFF;  // Turn on all LEDs on PORTB
            __delay_ms(500);  // Delay for 500ms
            LATB = 0x00;  // Turn off all LEDs
            __delay_ms(500);  // Delay for 500ms
          }
        }
        // Assembly Language:
        ; PIC18 Assembly Language Template
        ORG 0x00
        START:
            BCF STATUS, RP0   ; Select Bank 0
            MOVLW 0xFF        ; Load 0xFF into WREG (all LEDs on)
            MOVWF PORTB       ; Move WREG value to PORTB
            CALL Delay        ; Call delay subroutine
            CLRF PORTB        ; Clear PORTB (turn off LEDs)
            CALL Delay        ; Call delay subroutine
            GOTO START        ; Repeat
        Delay:
            MOVLW 0xFF        ; Load WREG with delay value
            MOVWF 0x30        ; Store in memory location
        DELAY_LOOP:
            DECFSZ 0x30, F    ; Decrement and check if zero
            GOTO DELAY_LOOP   ; Repeat until zero
            RETURN
        """,
    },
    "8085": {
        "template": """
        ; 8085 Assembly Language Template
        START:  MVI A, 00H ; Load Accumulator with 0
                OUT PORT1  ; Output data to port
                HLT        ; Halt program
        """,
    },
}

# Function to Generate Code using Groq API
def generate_code(prompt, model="llama-3.3-70b-versatile"):
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model=model,
    )
    return chat_completion.choices[0].message.content

# Function to Chunk and Tokenize Code
def tokenize_and_chunk_code(code, model_name="gpt2", max_length=512):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokens = tokenizer(code, truncation=True, return_tensors="pt")
    chunks = [
        tokens.input_ids[0][i: i + max_length]
        for i in range(0, len(tokens.input_ids[0]), max_length)
    ]
    return chunks

# Function to Store Code in FAISS
def store_code_in_faiss(chunks):
    dimension = 512
    embeddings = [np.random.rand(dimension) for _ in chunks]  # Replace with actual embeddings
    faiss_index = faiss.IndexFlatL2(dimension)
    faiss_index.add(np.array(embeddings, dtype=np.float32))
    return faiss_index

# Streamlit App
def main():
    # Custom styling for the app using markdown
    st.markdown("""
        <style>
            body {
                background: linear-gradient(to right, #ff7e5f, #feb47b, #6a11cb, #2575fc);
                font-family: Arial, sans-serif;
            }
            .title {
                font-size: 40px;
                font-weight: bold;
                color: #FF6347;
            }
            .subheader {
                font-size: 24px;
                color: #4CAF50;
            }
            .generated-code {
                background-color: #f0f0f0;
                padding: 10px;
                border-radius: 10px;
            }
            .sidebar .sidebar-content {
                background-color: #f4f4f4;
            }
            .stButton button {
                background-color: #008CBA;
                color: white;
                border-radius: 8px;
                padding: 10px;
                font-size: 16px;
            }
        </style>
    """, unsafe_allow_html=True)
    
    st.title("Microcontroller Code Generator", anchor="title")

    # Step 1: Display Custom Image
    # Use the image URL from Imgur
    image_url = "https://imgur.com/CBm5exE.png"
    st.image(image_url, caption="Custom Image", use_container_width=True)

    st.write("Select a microcontroller and provide your code requirements to generate code.")

    # Sidebar with different selection options for users
    with st.sidebar:
        st.header("Microcontroller Selection")
        device = st.selectbox("Select a microcontroller/processor", options=["Arduino", "PIC18", "8085"])
        st.write("Your selected microcontroller: ", device)

    # Step 1: Show Syntax Template based on selected device
    st.subheader("Coding Syntax Template", anchor="subheader")
    st.code(devices[device]["template"], language="c" if device != "8085" else "asm")

    # Step 2: Get User Prompt and Generate Code
    prompt = st.text_area("Enter your code requirements:", height=100)

    if prompt:
        if st.button("Generate Code"):
            st.write("Generating code, please wait...")
            code = generate_code(f"Write {device} code for: {prompt}")
            
            # Step 3: Chunk, Tokenize, and Store Code in FAISS
            chunks = tokenize_and_chunk_code(code)
            store_code_in_faiss(chunks)
            
            # Step 4: Display Generated Code
            st.subheader("Generated Code", anchor="generated-code")
            st.code(code)

# Run the Streamlit App
if __name__ == "__main__":
    main()