import streamlit as st
import warnings
import os
from pathlib import Path

# Suppress warnings yang tidak perlu
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", message=".*torchvision.*")
warnings.filterwarnings("ignore", message=".*UNEXPECTED.*")

from google import genai
from google.genai import types   # ← Penting
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document

# ================== CONFIG ==================
st.set_page_config(
    page_title="Human+ Lab-to-Protocol AI",
    page_icon="🧬",
    layout="centered"
)

st.title("🧬 Human+ Lab-to-Protocol AI")
st.markdown("""
**Upload hasil lab Anda → Dapatkan rekomendasi IV drip & protocol personal**  
*Semua data diproses sementara dan hilang setelah sesi selesai (privacy-first)*
""")

# ================== GEMINI CLIENT ==================
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not GEMINI_API_KEY:
    st.error("❌ GEMINI_API_KEY tidak ditemukan. Tambahkan di Settings → Secrets.")
    st.stop()

try:
    client = genai.Client(api_key=GEMINI_API_KEY)
    MODEL_NAME = "gemini-2.5-flash"
    st.success("✅ Gemini client siap")
except Exception as e:
    st.error(f"Gagal menginisialisasi Gemini: {str(e)}")
    st.stop()

# ================== LOAD KNOWLEDGE BASE ==================
@st.cache_resource(show_spinner=False, ttl=7200)
def load_knowledge_base():
    with st.spinner("⏳ Memuat knowledge base Human+..."):
        try:
            knowledge_dir = Path("/app/knowledge")
            if not knowledge_dir.exists():
                st.error(f"❌ Folder knowledge tidak ditemukan di: {knowledge_dir}")
                return None

            documents = []
            for file_path in knowledge_dir.glob("*.md"):
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read().strip()
                    if content:
                        documents.append(
                            Document(
                                page_content=content,
                                metadata={"source": file_path.name}
                            )
                        )

            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=800,
                chunk_overlap=150,
                separators=["\n\n## ", "\n\n### ", "\n\n", "\n", " ", "."]
            )
            chunks = text_splitter.split_documents(documents)

            embeddings = HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-mpnet-base-v2"
            )

            vector_store = FAISS.from_documents(chunks, embeddings)
            
            st.success(f"✅ Knowledge base siap ({len(documents)} file • {len(chunks)} chunks)")
            return vector_store

        except Exception as e:
            st.error(f"Gagal memuat knowledge base: {str(e)}")
            return None


vector_store = load_knowledge_base()
if vector_store is None:
    st.stop()

# ================== UPLOAD & ANALYSIS ==================
uploaded_file = st.file_uploader(
    "Upload hasil lab (PDF)", 
    type=["pdf"],
    help="Maksimal 10MB. Semua data hanya diproses sementara."
)

if uploaded_file:
    if uploaded_file.size > 10 * 1024 * 1024:
        st.error("❌ File terlalu besar. Maksimal 10MB.")
        st.stop()

    pdf_bytes = uploaded_file.read()

    with st.spinner("Menganalisis hasil lab menggunakan Gemini..."):
        try:
            pdf_part = types.Part.from_bytes(
                data=pdf_bytes,
                mime_type="application/pdf"
            )

            # Perbaikan utama: gunakan config=GenerateContentConfig
            response = client.models.generate_content(
                model=MODEL_NAME,
                contents=[
                    "Extract semua biomarker penting dari hasil lab PDF ini dalam format JSON yang terstruktur. "
                    "Sertakan nama tes, nilai, unit, reference range, dan flag (high/low/normal). "
                    "Fokus pada marker longevity: hs-CRP, Homocysteine, Vitamin D, HbA1c, Testosterone, "
                    "Magnesium, Zinc, B12, Ferritin, ApoB, dll.",
                    pdf_part
                ],
                config=types.GenerateContentConfig(
                    temperature=0.1,
                    max_output_tokens=4096
                )
            )

            extracted_text = response.text
            st.success("✅ Lab berhasil dianalisis!")
            
            with st.expander("📋 Hasil Ekstraksi Biomarker"):
                st.markdown(extracted_text)

        except Exception as e:
            st.error(f"Gagal menganalisis PDF: {str(e)}")
            st.stop()

    # ================== GENERATE PROTOCOL ==================
    if st.button("🔬 Generate Human+ Personalized Protocol", type="primary"):
        with st.spinner("Membuat rekomendasi protocol hyperpersonalized..."):
            try:
                retriever = vector_store.as_retriever(
                    search_type="mmr",
                    search_kwargs={"k": 8, "fetch_k": 20, "lambda_mult": 0.7}
                )
                
                relevant_docs = retriever.invoke(extracted_text)
                
                context = "\n\n".join([
                    f"[Sumber: {doc.metadata.get('source', 'unknown')}] {doc.page_content}"
                    for doc in relevant_docs
                ])

                prompt = f"""
Kamu adalah AI Longevity Specialist resmi dari **Human+ Bali**, didirikan oleh Benjamin White (@bennywhitethatsright).

Filosofi utama: **"From Injured → Optimized"**

Hasil lab pasien:
{extracted_text}

Knowledge base Human+:
{context}

Buatlah rekomendasi protocol yang jelas, positif, empowering, dan actionable dengan struktur berikut:

1. **Ringkasan Temuan Utama** (apa yang paling perlu dioptimalkan + motivasi)
2. **Rekomendasi IV Drip** (The All-In atau varian custom + komposisi)
3. **Suplemen Harian** (dosis & timing)
4. **Lifestyle & Recovery Protocol** (ice bath, structured water, breathing, sunlight)
5. **Jadwal Retest & Next Step**

Gunakan bahasa Indonesia yang mudah dipahami dan sesuai voice Benjamin White. Akhiri dengan semangat "From Injured → Optimized".
"""

                response = client.models.generate_content(
                    model=MODEL_NAME,
                    contents=[prompt],
                    config=types.GenerateContentConfig(
                        temperature=0.4,
                        max_output_tokens=4096
                    )
                )
                
                protocol = response.text

                st.subheader("📋 Rekomendasi Protocol Personal dari Human+")
                st.markdown(protocol)

                report_text = f"Human+ Lab-to-Protocol Report\n\n{protocol}"
                st.download_button(
                    label="📥 Download Report",
                    data=report_text,
                    file_name="Human+_Protocol_Report.txt",
                    mime="text/plain"
                )

            except Exception as e:
                st.error(f"Gagal generate protocol: {str(e)}")

# Footer
st.divider()
st.caption("""
**Privacy Note**: Aplikasi ini 100% stateless. Tidak ada data yang disimpan di server.  
Dibuat untuk Human+ Bali • From Injured → Optimized
""")