import streamlit as st
import requests
from rdkit import Chem
import datetime
from db import get_database

# Function to validate SMILES string
def is_valid_smiles(smiles):
    """ Validate if the input is a valid SMILES string using RDKit """
    mol = Chem.MolFromSmiles(smiles)
    return mol is not None

# Function to save prediction to MongoDB
def save_to_db(smiles_input, predictions):
    db = get_database()
    collection = db["polymers"]  # your collection
    doc = {
        "smiles": smiles_input,
        "predictions": predictions,
        "timestamp": datetime.datetime.utcnow()
    }
    collection.insert_one(doc)

# Streamlit page for Polymer Property Prediction
def show():
    st.markdown("<h1 style='text-align: center; color: #4CAF50;'>🔬 Polymer Property Prediction</h1>", unsafe_allow_html=True)
    st.markdown("<hr style='border: 1px solid #ccc;'>", unsafe_allow_html=True)

    # Input box with placeholder inside
    input_text = st.text_input(
        label="",
        placeholder="Enter SMILES ",
        key="smiles_input"
    )

    # Show Predict button always
    predict_clicked = st.button("🚀 Predict", use_container_width=True)

    # Predict on button click OR on pressing Enter with input
    if (predict_clicked or input_text) and input_text.strip():
        with st.spinner("Predicting..."):
            if is_valid_smiles(input_text.strip()):
                try:
                    input_data = {
                        "smiles": input_text.strip()  # Only sending smiles
                    }
                    response = requests.post("http://127.0.0.1:8000/predict", json=input_data)
                    if response.status_code == 200:
                        result = response.json()
                        renamed_properties = {
                            "property1": "Tensile_strength (Mpa)",
                            "property2": "Ionization_Energy (eV)",
                            "property3": "Electron_Affinity (eV)",
                            "property4": "LogP",
                            "property5": "Refractive_Index",
                            "property6": "Molecular_Weight (g/mol)"
                        }
                        predictions = {}
                        for key, name in renamed_properties.items():
                            value = result.get(key, 'N/A')
                            st.markdown(f"<div style='font-size:18px; padding: 6px 0;'><strong>{name}:</strong> {value}</div>", unsafe_allow_html=True)
                            predictions[name] = value

                        # Save prediction to MongoDB
                        save_to_db(input_text.strip(), predictions)

                        st.success("Prediction saved successfully!")
                    else:
                        st.error("Prediction failed. Please try again.")
                except Exception as e:
                    st.error(f"Error: {e}")
            else:
                st.error("❌ Invalid SMILES input. Please enter a correct SMILES string.")