import torch
import torch.nn as nn
import torch.optim as optim
import streamlit as st
import sqlite3
import re
import urllib.request
import urllib.parse
from hashlib import sha256

# ===== 1. CONFIG & MODEL (สมอง AI) =====
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EMBED_DIM, HIDDEN_DIM = 128, 256

class BataBrain(nn.Module):
    def __init__(self, vocab_size):
        super(BataBrain, self).__init__()
        self.embedding = nn.Embedding(vocab_size, EMBED_DIM)
        self.lstm = nn.LSTM(EMBED_DIM, HIDDEN_DIM, batch_first=True)
        self.fc = nn.Linear(HIDDEN_DIM, vocab_size)

    def forward(self, x):
        embeds = self.embedding(x)
        out, _ = self.lstm(embeds)
        return self.fc(out[:, -1, :])

# ===== 2. AI CONTROLLER (ระบบคิดและสืบค้น) =====
class BataAI:
    def __init__(self):
        if 'w2i' not in st.session_state:
            st.session_state.w2i = {"<PAD>": 0, "<START>": 1, "<END>": 2, "<UNK>": 3}
            st.session_state.i2w = ["<PAD>", "<START>", "<END>", "<UNK>"]
        
        self.init_model()

    def init_model(self):
        v_size = len(st.session_state.i2w)
        self.model = BataBrain(v_size).to(DEVICE)
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)

    def search_web(self, query):
        try:
            q = urllib.parse.quote(query + " คืออะไร ข้อมูลทางการ")
            url = f"https://www.google.com/search?q={q}"
            req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
            with urllib.request.urlopen(req) as res:
                html = res.read().decode('utf-8', errors='ignore')
                clean = re.sub(r'<[^>]+>', ' ', html)
                snippets = re.findall(r'[ก-๙a-zA-Z\s]{30,200}', clean)
                return " ".join(snippets[5:10])
        except: return ""

    def generate_formal_response(self, prompt):
        info = self.search_web(prompt)
        prefix = "เรียน ท่านผู้ใช้งาน จากการตรวจสอบข้อมูลเบื้องต้นพบว่า "
        body = info[:200] if info else "ข้อมูลดังกล่าวยังไม่ปรากฏในฐานระบบขณะนี้"
        suffix = "\n\nจึงเรียนมาเพื่อโปรดพิจารณา หากต้องการข้อมูลเพิ่มเติมกรุณาแจ้งดิฉันค่ะ"
        return prefix + body + suffix

# ===== 3. DATABASE & AUTH (ระบบ Login) =====
def init_db():
    conn = sqlite3.connect('bata_web.db')
    c = conn.cursor()
    c.execute('CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, password TEXT)')
    conn.commit()
    return conn

# ===== 4. UI (หน้าเว็บ Streamlit) =====
def main():
    st.set_page_config(page_title="Bata AI Corporate", page_icon="💅")
    st.title("👁️👄👁️ Bata AI: Corporate Web")
    
    init_db()
    ai = BataAI()

    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False

    # --- LOGIN PAGE ---
    if not st.session_state.logged_in:
        tab1, tab2 = st.tabs(["เข้าสู่ระบบ", "ลงทะเบียน"])
        
        with tab1:
            u = st.text_input("ชื่อผู้ใช้งาน", key="login_u")
            p = st.text_input("รหัสผ่าน", type="password", key="login_p")
            if st.button("Login"):
                hp = sha256(p.encode()).hexdigest()
                conn = init_db()
                c = conn.cursor()
                c.execute("SELECT * FROM users WHERE username=? AND password=?", (u, hp))
                if c.fetchone():
                    st.session_state.logged_in = True
                    st.session_state.user = u
                    st.rerun()
                else: st.error("ข้อมูลไม่ถูกต้องค่ะแม่")

        with tab2:
            ru = st.text_input("ตั้งชื่อผู้ใช้งาน", key="reg_u")
            rp = st.text_input("ตั้งรหัสผ่าน", type="password", key="reg_p")
            if st.button("สมัครสมาชิก"):
                rhp = sha256(rp.encode()).hexdigest()
                try:
                    conn = init_db()
                    c = conn.cursor()
                    c.execute("INSERT INTO users VALUES (?,?)", (ru, rhp))
                    conn.commit()
                    st.success("ลงทะเบียนสำเร็จแล้วค่ะ! เชิญไป Login")
                except: st.error("ชื่อนี้มีคนใช้แล้วค่ะตัวมัม")

    # --- CHAT PAGE ---
    else:
        st.sidebar.write(f"👤 ผู้ใช้งาน: **{st.session_state.user}**")
        if st.sidebar.button("Logout"):
            st.session_state.logged_in = False
            st.rerun()

        st.subheader("สนทนาภาษาทางการ")
        user_msg = st.text_input("พิมพ์คำถามของท่าน...", placeholder="เช่น AI คืออะไร?")
        
        if st.button("ส่งคำถาม"):
            with st.spinner('ดิฉันกำลังใช้สมองคิดอยู่ค่ะ...'):
                response = ai.generate_formal_response(user_msg)
                st.info(f"**Bata AI:** {response}")

if __name__ == "__main__":
    main()