import streamlit as st
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import os
import re

# === Konfigurasi API Hugging Face Inference API ===

HF_TOKEN = os.environ.get("HF_TOKEN", "")  # simpan secret via HF Spaces
os.environ["OPENAI_API_KEY"] = HF_TOKEN
HF_BASE_URL = "https://api-inference.huggingface.com/v1"

# Fungsi koneksi database
def connectDatabase(username, port, host, password, database):
    try:
        db_uri = f"postgresql://{username}:{password}@{host}:{port}/{database}"
        st.session_state.db = SQLDatabase.from_uri(db_uri)
        st.success("✅ Database connected successfully.")
    except Exception as e:
        st.session_state.db = None
        st.error(f"❌ Failed to connect to database:\n{e}")

# Fungsi menjalankan query SQL
def runQuery(query):
    if "db" in st.session_state and st.session_state.db:
        try:
            return st.session_state.db.run(query)
        except Exception as e:
            return f"❌ Query execution failed: {e}"
    return "❌ Please connect to database first."

# Ambil skema database
def getDatabaseSchema():
    if "db" in st.session_state and st.session_state.db:
        try:
            return st.session_state.db.get_table_info()
        except Exception as e:
            return f"❌ Failed to fetch schema: {e}"
    return "❌ Please connect to database first."

# Prompt untuk menghasilkan SQL
def getQueryFromLLM(question):
    template = """
                Below is the schema of a PostgreSQL database. Read the schema carefully and answer the user's question using a valid SQL query. Be careful with table and column names (case-sensitive). Only provide the SQL query, and nothing else.

                Schema:
                {schema}

                Sekarang giliran Anda:
                Question: {question}
                SQL query:
                """
    try:
        schema = getDatabaseSchema()[:3000]
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | llm
        response = chain.invoke({"question": question, "schema": schema})
        query_only = re.sub(r"<think>.*?</think>", "", response.content, flags=re.DOTALL).strip()
        return query_only
    except Exception as e:
        st.error(f"❌ Model gagal merespons: {e}")
        return "Gagal membuat query SQL."

# Prompt untuk menjelaskan hasil
def getResponseForQueryResult(question, query, result):
    template = """
            Berdasarkan pertanyaan, SQL query, dan hasilnya, berikan jawaban dalam bahasa alami. Gunakan bahasa Indonesia jika pertanyaannya dalam bahasa Indonesia.

            Schema:
            {schema}

            Pertanyaan: {question}
            SQL query: {query}
            Hasil: {result}
            Jawaban:
            """
    prompt2 = ChatPromptTemplate.from_template(template)
    chain2 = prompt2 | llm
    response = chain2.invoke({
        "question": question,
        "schema": getDatabaseSchema(),
        "query": query,
        "result": result
    })
    return response.content.strip()

# Konfigurasi halaman
st.set_page_config(page_title="Chat with PostgreSQL DB", page_icon="🧠", layout="centered")
st.title("🧠 Chat with your Database")

# Inisialisasi state
if "chat" not in st.session_state:
    st.session_state.chat = []

# Sidebar
with st.sidebar:
    st.subheader("🔗 Koneksi Database")
    host = st.text_input("Host", value="aws-0-ap-southeast-1.pooler.supabase.com")
    port = st.text_input("Port", value="6543")
    username = st.text_input("Username", value="postgres")
    password = st.text_input("Password", type="password", value="password")
    database = st.text_input("Database", value="postgres")
    if st.button("Connect"):
        connectDatabase(username, port, host, password, database)

# Inisialisasi LLM
llm = ChatOpenAI(
    model="qwen/qwen3-1.7b",
    base_url=HF_BASE_URL,
    api_key=HF_TOKEN,
    temperature=0
)

# Input pertanyaan
question = st.chat_input("Tanyakan sesuatu tentang database Anda...")

# Proses pertanyaan
if question:
    if "db" not in st.session_state or st.session_state.db is None:
        st.error("❌ Please connect to database first.")
    else:
        st.session_state.chat.append({"role": "user", "content": question})

        with st.spinner("💡 Sedang berpikir..."):
            query = getQueryFromLLM(question)
            result = runQuery(query)
            response = getResponseForQueryResult(question, query, result)

        st.session_state.chat.append({"role": "assistant", "content": response})

        with st.expander("🧾 Generated SQL Query"):
            st.code(query, language="sql")

        with st.expander("📊 Raw Result"):
            st.write(result)

# Tampilkan chat
for chat in st.session_state.chat:
    st.chat_message(chat["role"]).markdown(chat["content"])