import os
import requests
import gradio as gr
from ase.io import read
import nglview as nv
import tempfile

# Retrieve the API key from the environment variable
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    raise ValueError("GROQ_API_KEY is missing! Set it in the Hugging Face Spaces 'Secrets'.")

# Define the API endpoint and headers
url = "https://api.groq.com/openai/v1/chat/completions"
headers = {"Authorization": f"Bearer {groq_api_key}"}

# Function to interact with Groq API
def chat_with_groq(user_input):
    body = {
        "model": "llama-3.1-8b-instant",
        "messages": [{"role": "user", "content": user_input}]
    }
    response = requests.post(url, headers=headers, json=body)
    if response.status_code == 200:
        return response.json()['choices'][0]['message']['content']
    else:
        return f"Error: {response.json()}"

# Function to generate and visualize CIF files
def generate_and_visualize_cifs(application):
    # Get response from Groq API
    response = chat_with_groq(f"Suggest 3 best materials for {application} with their properties and provide CIF data.")
    
    # Parse the response to extract materials and their CIF contents
    # This is a placeholder; actual implementation depends on response format
    materials = parse_response(response)
    
    views = []
    for material in materials:
        cif_content = material['cif_content']
        with tempfile.NamedTemporaryFile(mode='w+', suffix='.cif', delete=False) as tmp:
            tmp.write(cif_content)
            tmp_path = tmp.name
        atoms = read(tmp_path)
        view = nv.show_ase(atoms)
        views.append(view)
    return views

# Gradio interface
def interface_function(application):
    views = generate_and_visualize_cifs(application)
    return views

interface = gr.Interface(
    fn=interface_function,
    inputs=gr.Textbox(lines=2, placeholder="Enter material application..."),
    outputs="text",
    title="Material Selector and Visualizer",
    description="Enter a material application to get top 3 materials and visualize their atomic structures."
)

if __name__ == "__main__":
    interface.launch()