import streamlit as st
import pandas as pd
import hashlib
import os
from datetime import datetime
from typing import Optional, Dict, Any
from dataclasses import dataclass
from rapidfuzz import process, fuzz
from pymatgen.core import Composition
from langchain_google_genai import ChatGoogleGenerativeAI

# ==========================================
# 1. APP CONFIG & STYLING
# ==========================================
st.set_page_config(
    page_title="NACE MR0175 Auditor",
    page_icon="🛡️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Professional UI
st.markdown("""
<style>
    .stButton>button { width: 100%; border-radius: 5px; }
    .success-box { padding: 15px; background-color: #d1e7dd; border-left: 5px solid #198754; color: #0f5132; }
    .fail-box { padding: 15px; background-color: #f8d7da; border-left: 5px solid #dc3545; color: #842029; }
    .warn-box { padding: 15px; background-color: #fff3cd; border-left: 5px solid #ffc107; color: #664d03; }
    .guide-box { padding: 15px; background-color: #e2e3e5; border-radius: 10px; margin-bottom: 20px; }
</style>
""", unsafe_allow_html=True)

@dataclass
class MaterialRecord:
    name: str
    uns: str
    formula: str
    source: str
    confidence: str 

# ==========================================
# 2. DATA & SERVICE LAYERS
# ==========================================
@st.cache_data
def load_database():
    try:
        return pd.read_csv("materials.csv")
    except:
        return pd.DataFrame() 

df_materials = load_database()

class MaterialService:
    def __init__(self, api_key):
        self.api_key = api_key
        self.llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", google_api_key=api_key, temperature=0)

    def search_database(self, query: str) -> Optional[MaterialRecord]:
        if df_materials.empty: return None
        choices = df_materials['Name'].tolist() + df_materials['UNS'].tolist()
        match = process.extractOne(query, choices, scorer=fuzz.WRatio)
        
        if match and match[1] > 85: 
            found_val = match[0]
            record = df_materials[(df_materials['Name'] == found_val) | (df_materials['UNS'] == found_val)].iloc[0]
            return MaterialRecord(record['Name'], record['UNS'], record['Formula'], f"Internal DB ({record['Standard']})", "HIGH")
        return None

    def ask_ai_for_formula(self, query: str) -> MaterialRecord:
        prompt = f"""
        Act as a Metallurgist. Identify the alloy '{query}'.
        Return JSON ONLY: {{"name": "Standard Name", "uns": "UNS Code", "formula": "ChemicalString"}}
        """
        try:
            response = self.llm.invoke(prompt)
            import json
            clean_json = response.content.replace('```json', '').replace('```', '').strip()
            data = json.loads(clean_json)
            return MaterialRecord(data.get('name', 'Unknown'), data.get('uns', 'Unknown'), data.get('formula', ''), "AI Estimate (Verify Formula)", "LOW")
        except:
            return None

class PhysicsEngine:
    @staticmethod
    def calculate_pren(formula: str) -> Dict[str, Any]:
        try:
            comp = Composition(formula)
            cr = comp.get_wt_fraction("Cr") * 100
            mo = comp.get_wt_fraction("Mo") * 100
            n = comp.get_wt_fraction("N") * 100
            
            pren = cr + (3.3 * mo) + (16 * n)
            
            if pren > 40: verdict, color = "PASS", "green"
            elif pren > 30: verdict, color = "CONDITIONAL", "orange"
            else: verdict, color = "FAIL", "red"
                
            return {
                "pren": round(pren, 2),
                "breakdown": {"Cr": round(cr,1), "Mo": round(mo,1), "N": round(n,2)},
                "verdict": verdict,
                "color": color
            }
        except Exception as e:
            return {"error": str(e)}

# ==========================================
# 3. FRONTEND UI
# ==========================================
api_key = os.environ.get("GOOGLE_API_KEY")
if not api_key: st.stop()

# Initialize Session State
if 'search_query' not in st.session_state: st.session_state.search_query = ""
if 'material_record' not in st.session_state: st.session_state.material_record = None

# --- TITLE & ONBOARDING ---
st.title("🛡️ NACE MR0175 Compliance Auditor")
st.caption("AI-Assisted Material Verification for Sour Service ($H_2S$) Environments")

with st.expander("👋 New here? Read the Quick Start Guide", expanded=True):
    st.markdown("""
    **What does this tool do?**  
    It verifies if an alloy is safe to use in Sour Gas wells (where $H_2S$ causes cracking).
    
    **How to use it:**
    1.  **Search:** Type a material name (e.g., *Inconel 625*) or UNS Code (e.g., *N06625*).
    2.  **Verify:** The app will pull the chemical formula from our database (or ask AI if unknown).
    3.  **Audit:** You check the formula matches your MTR (Mill Cert), then click **Run Analysis**.
    4.  **Result:** It calculates the **PREN** score to determine Pass/Fail status.
    """)

# --- INPUT SECTION ---
st.markdown("### 1. Select Material")

c1, c2, c3, c4 = st.columns(4)
# Quick Action Buttons
if c1.button("Inconel 625"): st.session_state.search_query = "UNS N06625"
if c2.button("316L SS"): st.session_state.search_query = "UNS S31603"
if c3.button("Duplex 2205"): st.session_state.search_query = "Duplex 2205"
if c4.button("Clear"): 
    st.session_state.search_query = ""
    st.session_state.material_record = None

# Search Bar
col1, col2 = st.columns([3, 1])
with col1:
    query = st.text_input("Search Database", value=st.session_state.search_query, placeholder="e.g. UNS N06625, Super Duplex, Monel...")
with col2:
    st.markdown("<br>", unsafe_allow_html=True)
    search_btn = st.button("🔍 Search Database", type="primary")

# Search Logic
if search_btn and query:
    material_service = MaterialService(api_key)
    with st.spinner("Searching Internal Approved List..."):
        result = material_service.search_database(query)
        if result:
            st.success(f"Match Found: {result.name}")
            st.session_state.material_record = result
        else:
            with st.spinner("Not in DB. Consulting AI Standards..."):
                result = material_service.ask_ai_for_formula(query)
                if result:
                    st.warning(f"AI Estimation: {result.name}")
                    st.session_state.material_record = result
                else:
                    st.error("Material not identified.")

# --- RESULTS SECTION ---
if st.session_state.material_record:
    rec = st.session_state.material_record
    st.markdown("---")
    
    # 2. Verification
    c1, c2 = st.columns(2)
    with c1:
        st.subheader("2. Specification Check")
        st.info(f"**Name:** {rec.name}\n\n**Code:** `{rec.uns}`\n\n**Source:** {rec.source}")
    with c2:
        st.subheader("3. Verify Chemistry (Human-in-Loop)")
        confirmed_formula = st.text_input("Review Formula (Edit if MTR differs)", rec.formula)
        
        run_calc = st.button("Run NACE Calculation ⚡", type="primary", use_container_width=True)

    # 3. Calculation & Logic
    if run_calc:
        res = PhysicsEngine.calculate_pren(confirmed_formula)
        
        if "error" in res:
            st.error("Invalid Chemical Formula. Please check input.")
        else:
            st.markdown("---")
            st.subheader("4. Audit Result")
            
            # The Verdict
            if res['verdict'] == "PASS":
                st.markdown(f'<div class="success-box"><h3>✅ APPROVED for Severe Sour Service</h3>PREN is > 40. Material typically resists Pitting and Stress Corrosion Cracking in H2S.</div>', unsafe_allow_html=True)
            elif res['verdict'] == "CONDITIONAL":
                st.markdown(f'<div class="warn-box"><h3>⚠️ CONDITIONAL USE</h3>PREN is between 30-40. Acceptable for mild sour service only. Check NACE tables for partial pressure limits.</div>', unsafe_allow_html=True)
            else:
                st.markdown(f'<div class="fail-box"><h3>🛑 REJECTED for Sour Service</h3>PREN is < 30. High risk of catastrophic failure in H2S environments. Use for Sweet Service only.</div>', unsafe_allow_html=True)
            
            st.markdown("<br>", unsafe_allow_html=True)
            
            # Metrics with Tooltips
            k1, k2, k3 = st.columns(3)
            k1.metric("PREN Score", res['pren'], help="Pitting Resistance Equivalent Number. Formula: %Cr + 3.3(%Mo) + 16(%N). Target > 40.")
            k2.metric("Molybdenum (Mo)", f"{res['breakdown']['Mo']}%", help="Primary element for resisting pitting corrosion.")
            k3.metric("Chromium (Cr)", f"{res['breakdown']['Cr']}%", help="Forms the passive oxide layer.")
            
            # Provenance
            timestamp = datetime.now().isoformat()
            token = hashlib.sha256(f"{rec.uns}{res['pren']}{timestamp}".encode()).hexdigest()[:12]
            
            with st.expander("📄 View Digital Provenance (MTR Hash)"):
                st.caption(f"SHA-256 Signature: {token}")
                st.caption(f"Timestamp: {timestamp}")
                st.caption("Algorithm: Pymatgen Core v2024")
                st.json(res)