src/streamlit_app.py

import streamlit as st
import pypdf
import sqlite3
import base64
import os
import shutil
import httpx
import ssl
from groq import Groq
from huggingface_hub import HfApi, hf_hub_download
from langchain_community.document_loaders import PyPDFLoader, UnstructuredPowerPointLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

# --- 0. CONFIGURATION & SSL ---
ssl._create_default_https_context = ssl._create_unverified_context
REPO_ID = "jesseys/private-ai-data"
DB_FILE = "chat_history.db"
VECTOR_INDEX_DIR = "faiss_index"
HF_TOKEN = os.environ.get("AdminToken")
GROQ_KEY = os.environ.get("GROQ_API_KEY")
PASS = os.environ.get("password")

api = HfApi(token=HF_TOKEN)
client = Groq(api_key=GROQ_KEY, http_client=httpx.Client(verify=False))
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

# --- 1. CORE FUNCTIONS ---
def sync_data(direction="pull"):
    files_to_sync = [DB_FILE, f"{VECTOR_INDEX_DIR}/index.faiss", f"{VECTOR_INDEX_DIR}/index.pkl"]
    for f_path in files_to_sync:
        try:
            if direction == "pull":
                downloaded = hf_hub_download(repo_id=REPO_ID, filename=f_path, repo_type="dataset", token=HF_TOKEN)
                os.makedirs(os.path.dirname(f_path) or ".", exist_ok=True)
                shutil.copy(downloaded, f_path)
            else:
                if os.path.exists(f_path):
                    api.upload_file(path_or_fileobj=f_path, path_in_repo=f_path, repo_id=REPO_ID, repo_type="dataset")
        except: pass

def clear_chat_history():
    """Fungsi untuk menghapus riwayat chat di DB dan Session"""
    if os.path.exists(DB_FILE):
        conn = sqlite3.connect(DB_FILE)
        conn.execute("DELETE FROM messages")
        conn.commit()
        conn.close()
    st.session_state.messages = []
    sync_data(direction="push")
    st.rerun()

def train_on_files(uploaded_files):
    all_docs = []
    for uploaded_file in uploaded_files:
        temp_path = f"temp_{uploaded_file.name}"
        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        loader = PyPDFLoader(temp_path) if uploaded_file.name.endswith(".pdf") else UnstructuredPowerPointLoader(temp_path)
        all_docs.extend(loader.load())
        os.remove(temp_path)
    
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    chunks = splitter.split_documents(all_docs)
    vectorstore = FAISS.from_documents(chunks, embeddings)
    vectorstore.save_local(VECTOR_INDEX_DIR)
    sync_data(direction="push")
    return True

# --- 2. DATABASE INIT ---
def init_db():
    conn = sqlite3.connect(DB_FILE)
    conn.execute('CREATE TABLE IF NOT EXISTS messages (role TEXT, content TEXT, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)')
    conn.commit(); conn.close()

def save_chat(role, content):
    conn = sqlite3.connect(DB_FILE)
    conn.execute("INSERT INTO messages (role, content) VALUES (?, ?)", (role, content))
    conn.commit(); conn.close()
    sync_data(direction="push")

# --- IMAGE ENCODER ---
def encode_image(uploaded_file):
    return base64.b64encode(uploaded_file.getvalue()).decode('utf-8')

#reset knowledge yang tadi sudah di train ke default model
def reset_knowledge_base():
    """Menghapus semua data vektor (pengetahuan PDF) secara lokal dan di cloud."""
    # 1. Hapus Folder Lokal
    if os.path.exists(VECTOR_INDEX_DIR):
        shutil.rmtree(VECTOR_INDEX_DIR)
    
    # 2. Hapus File di Hugging Face Dataset
    # Kita hapus index.faiss dan index.pkl dari repo agar balik ke default
    files_to_delete = [f"{VECTOR_INDEX_DIR}/index.faiss", f"{VECTOR_INDEX_DIR}/index.pkl"]
    for f_path in files_to_delete:
        try:
            api.delete_file(path_in_repo=f_path, repo_id=REPO_ID, repo_type="dataset")
        except:
            pass # Abaikan jika file memang tidak ada di repo
            
    st.success("Pengetahuan telah dikembalikan ke default!")
    st.rerun()

# --- 3. UI INTERFACE ---
st.set_page_config(page_title="Private AI", layout="wide")

if "initialized" not in st.session_state:
    sync_data(direction="pull")
    init_db()
    st.session_state.initialized = True

with st.sidebar:
    st.title("🛡️ Admin Panel")
    if st.toggle("Admin Access"):
        if st.text_input("Password", type="password") == PASS:
            # Fitur Training
            train_files = st.file_uploader("Upload Knowledge (PDF/PPTX)", accept_multiple_files=True, type=["pdf", "pptx"])
            if st.button("Start Training"):
                if train_on_files(train_files): st.success("Knowledge Updated!")
            
            st.divider()
            # Fitur Hapus Chat
            st.warning("Zona Berbahaya")
            if st.button("🗑️ Hapus Semua Riwayat Chat"):
                clear_chat_history()
                
            # Tombol reset pengetahuan
            if st.button("🗑️ Reset Semua Pengetahuan"):
                reset_knowledge_base()
        else: st.stop()

    st.divider()
    st.title("📂 Chat Context")
    chat_file = st.file_uploader("Upload Image/File for this chat only", type=["png", "jpg", "pdf"])
st.title("🚀 Knowledge Assistant")


# Load Chat History
if "messages" not in st.session_state:
    conn = sqlite3.connect(DB_FILE)
    st.session_state.messages = [{"role": r[0], "content": r[1]} for r in conn.execute("SELECT role, content FROM messages ORDER BY timestamp").fetchall()]
    conn.close()

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]): st.markdown(msg["content"])

# Chat Input
if prompt := st.chat_input("Tanyakan sesuatu..."):
    # Tampilkan pesan user segera
    st.chat_message("user").markdown(prompt)
    
    # --- 1. AMBIL KONTEKS TERBARU ---
    context = ""
    if os.path.exists(VECTOR_INDEX_DIR):
        try:
            # Muat ulang index setiap kali bertanya agar data 'training' terbaru terbaca
            vs = FAISS.load_local(VECTOR_INDEX_DIR, embeddings, allow_dangerous_deserialization=True)
            docs = vs.similarity_search(prompt, k=3)
            context = "\n".join([d.page_content for d in docs])
        except Exception as e:
            context = f"Terjadi kesalahan memuat dokumen: {e}"

    # --- 2. SETUP MODEL VISION ---
    vision_model = "meta-llama/llama-4-scout-17b-16e-instruct" 

    # --- 3. SIAPKAN KONTEN MULTIMODAL ---
    api_content = [{"type": "text", "text": prompt}]
    
    if chat_file and chat_file.type in ["image/png", "image/jpeg", "image/jpg"]:
        # Tampilkan indikator bahwa gambar sedang diproses
        st.sidebar.image(chat_file, caption="Gambar terdeteksi", width=150)
        base_64_image = encode_image(chat_file)
        api_content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{base_64_image}"}
        })

    # --- 4. KIRIM KE API DENGAN BATASAN RIWAYAT ---
    try:
        # Ambil maksimal 6 pesan terakhir agar konteks tetap tajam
        recent_messages = st.session_state.messages[-6:] if len(st.session_state.messages) > 6 else st.session_state.messages
        
        response = client.chat.completions.create(
            model=vision_model,
            messages=[
                {
                    "role": "system", 
                    "content": f"Anda adalah asisten cerdas. Gunakan konteks PDF ini untuk menjawab: {context}. Jawablah pertanyaan user dengan tepat."
                },
                *recent_messages,
                {"role": "user", "content": api_content}
            ]
        )
        
        answer = response.choices[0].message.content
        with st.chat_message("assistant"): st.markdown(answer)
        
        # Simpan ke memori dan sinkronkan
        save_chat("user", prompt)
        save_chat("assistant", answer)
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.messages.append({"role": "assistant", "content": answer})
        
    except Exception as e: 
        st.error(f"Error pada AI: {e}")