import os
import requests
import gradio as gr
from pymatgen.core import Structure, Lattice
from pymatgen.io.cif import CifWriter
import tempfile
import shutil

# 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 get material recommendations from Groq API
def get_material_recommendations(application):
    body = {
        "model": "llama-3.1-8b-instant",
        "messages": [{"role": "user", "content": f"Recommend three best materials for {application} with their key properties."}]
    }
    
    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 create dummy structures for visualization
def create_dummy_structure(material_name):
    # This function creates a simple cubic structure as a placeholder
    lattice = Lattice.cubic(3.5)
    structure = Structure(
        lattice,
        ["Si"],
        [[0, 0, 0]]
    )
    return structure

# Function to generate CIF files and visualize structures
def generate_and_visualize(application):
    # Get material recommendations
    recommendations = get_material_recommendations(application)
    
    # Parse the recommendations to extract material names
    # For simplicity, we'll assume the material names are listed as:
    # 1. Material A
    # 2. Material B
    # 3. Material C
    lines = recommendations.split('\n')
    material_names = []
    for line in lines:
        if line.strip().startswith(('1.', '2.', '3.')):
            parts = line.split('.')
            if len(parts) > 1:
                material_name = parts[1].strip().split(' ')[0]
                material_names.append(material_name)
    
    # Create temporary directory to store CIF files
    temp_dir = tempfile.mkdtemp()
    cif_file_paths = []
    for name in material_names:
        structure = create_dummy_structure(name)
        cif_path = os.path.join(temp_dir, f"{name}.cif")
        writer = CifWriter(structure)
        writer.write_file(cif_path)
        cif_file_paths.append(cif_path)
    
    # Display the recommendations and provide CIF files for download
    output_text = recommendations
    download_links = [gr.File(cif_path, label=os.path.basename(cif_path)) for cif_path in cif_file_paths]
    
    return output_text, download_links

# Create Gradio interface
interface = gr.Interface(
    fn=generate_and_visualize,
    inputs=gr.Textbox(lines=2, placeholder="Enter application (e.g., aerospace structural component)"),
    outputs=[gr.Textbox(label="Material Recommendations"), gr.File(label="Download CIF Files")],
    title="Materials Science Expert",
    description="Enter a specific application to get the top 3 material recommendations with their properties and download their CIF files."
)

# Launch Gradio app
if __name__ == "__main__":
    interface.launch()