import os
import gradio as gr
from groq import Groq
from transformers import AutoTokenizer
import faiss
import numpy as np

# Authenticate Groq API using the API key stored on Hugging Face
os.environ["GROQ_API_KEY"] = "MY_API_KEY"
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Supported Microcontrollers/Processors
devices = {
    "Arduino": {
        "template": """
        // Arduino Syntax Template
        void setup() {
          // Setup code here, to run once:
        }

        void loop() {
          // Main code here, to run repeatedly:
        }
        """,
    },
    "PIC18": {
        "template": """
        // PIC18 Syntax Template
        #include <xc.h>
        void main() {
          // Initialization code
          while(1) {
            // Main program loop
          }
        }
        """,
    },
    "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
def generate_code(prompt, model="llama-3.3-70b-versatile"):
    try:
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model=model,
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Error during code generation: {str(e)}"

# 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

# Gradio Interface Logic
def code_generator(device, prompt):
    if device not in devices:
        return f"Error: {device} is not supported."

    # Show Syntax Template
    syntax_template = devices[device]["template"]

    # Generate Code
    generated_code = generate_code(f"Write {device} code for: {prompt}")

    # Handle if code generation fails
    if "Error" in generated_code:
        return generated_code  # Return the error message if the generation fails

    # Chunk, Tokenize, and Store Code in FAISS
    chunks = tokenize_and_chunk_code(generated_code)
    store_code_in_faiss(chunks)

    return f"### Syntax Template:\n{syntax_template}\n\n### Generated Code:\n{generated_code}"

# Gradio UI
device_dropdown = gr.Dropdown(choices=["Arduino", "PIC18", "8085"], label="Select Microcontroller/Processor")
prompt_input = gr.Textbox(label="Enter Your Code Requirements")
output_display = gr.Textbox(label="Generated Code", lines=20)

interface = gr.Interface(
    fn=code_generator,
    inputs=[device_dropdown, prompt_input],
    outputs=output_display,
    title="Microcontroller Code Generator",
    description="Generate code for Arduino, PIC18, or 8085 based on your requirements.",
)

# Run Gradio App
if __name__ == "__main__":
    interface.launch()