app.py

# -*- coding: utf-8 -*-
"""
🏛️ AI HÀNH CHÍNH CÔNG VIP PRO
Version 2.0 - Deploy on Hugging Face Spaces
"""

import os
import warnings
warnings.filterwarnings('ignore')

# Tạo thư mục
os.makedirs('data', exist_ok=True)

# ============================================================
# IMPORTS
# ============================================================
import torch
import gc
import json
import numpy as np
import pandas as pd
from datetime import datetime
import gradio as gr

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_core.documents import Document
from langchain_pinecone import PineconeVectorStore
from pinecone import Pinecone
from rank_bm25 import BM25Okapi
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import TfidfVectorizer

# ============================================================
# DỮ LIỆU DỊCH VỤ CÔNG (40+ THỦ TỤC)
# ============================================================
DICHVUCONG_DATA = {
    # NHÂN SỤ
    "KHAI_SINH": {"code":"KHAI_SINH","name":"Đăng ký khai sinh","docs":"Giấy khai sinh, CCCD cha mẹ, Sổ hộ khẩu","coquan":"UBND xã/phường","time":"1-3 ngày","lephi":"Miễn phí","note":"Làm trong 60 ngày kể từ khi sinh","steps":"1. Chuẩn bị: Giấy khai sinh, CCCD cha mẹ, Sổ hộ khẩu. 2. Đến UBND xã/phường nơi cư trú. 3. Nộp hồ sơ. 4. Nhận kết quả sau 1-3 ngày.","category":"NHAN_SU"},
    "KHAI_SINH_QUA_HAN": {"code":"KHAI_SINH_QUA_HAN","name":"Khai sinh quá hạn","docs":"Giấy khai sinh, CCCD, Sổ hộ khẩu, Đơn giải trình","coquan":"UBND xã/phường","time":"3-7 ngày","lephi":"Miễn phí","note":"Quá 60 ngày cần đơn giải trình","steps":"1. Viết đơn giải trình. 2. Xin xác nhận tổ dân phố. 3. Đến UBND nộp hồ sơ.","category":"NHAN_SU"},
    "KHAI_TU": {"code":"KHAI_TU","name":"Đăng ký khai tử","docs":"Giấy báo tử, CCCD người chết, Sổ hộ khẩu","coquan":"UBND xã/phường","time":"1-3 ngày","lephi":"Miễn phí","note":"Làm trong 60 ngày kể từ khi chết","steps":"1. Lấy giấy báo tử. 2. Chuẩn bị hồ sơ. 3. Đến UBND nơi người chết cư trú.","category":"NHAN_SU"},
    "KET_HON": {"code":"KET_HON","name":"Đăng ký kết hôn","docs":"CCCD cả 2 bên, Sổ hộ khẩu, Giấy xác nhận tình trạng hôn nhân","coquan":"UBND xã/phường","time":"1-3 ngày","lephi":"Miễn phí","note":"Trong 30 ngày có giấy xác nhận","steps":"1. Xin giấy xác nhận tình trạng hôn nhân. 2. Cả 2 đến UBND. 3. Đăng ký kết hôn.","category":"NHAN_SU"},
    "LY_HON": {"code":"LY_HON","name":"Ly hôn thuận tình","docs":"Đơn xin ly hôn, Giấy kết hôn, CCCD, Sổ hộ khẩu","coquan":"Tòa án","time":"2-4 tháng","lephi":"300.000đ","note":"Cả 2 bên đồng ý","steps":"1. Soạn đơn xin ly hôn. 2. Nộp tại Tòa án. 3. Tòa án thụ lý và ra quyết định.","category":"NHAN_SU"},
    "KET_HON_NGOAI": {"code":"KET_HON_NGOAI","name":"Kết hôn người nước ngoài","docs":"CCCD, Hộ chiếu, Giấy xác nhận, Giấy hợp pháp hóa lãnh sự","coquan":"UBND quận/huyện","time":"5-7 ngày","lephi":"Miễn phí","note":"Cần hợp pháp hóa lãnh sự","steps":"1. Hợp pháp hóa giấy tờ. 2. Xin xác nhận. 3. Đăng ký tại UBND quận.","category":"NHAN_SU"},
    # CCCD - HỘ KHẨU
    "CCCD_CAP": {"code":"CCCD_CAP","name":"Cấp CCCD lần đầu","docs":"Giấy khai sinh, Sổ hộ khẩu","coquan":"Công an quận/huyện","time":"7-15 ngày","lephi":"Miễn phí","note":"Từ 14 tuổi","steps":"1. Chuẩn bị giấy khai sinh, sổ hộ khẩu. 2. Đến Công an quận/huyện. 3. Làm thủ tục cấp CCCD.","category":"CCCD_HK"},
    "CCCD_DOI": {"code":"CCCD_DOI","name":"Đổi CCCD","docs":"CCCD cũ, Giấy khai sinh","coquan":"Công an quận/huyện","time":"7-15 ngày","lephi":"Miễn phí","note":"Đổi khi hư hỏng/sai thông tin","steps":"1. Mang CCCD cũ. 2. Đến Công an nơi thường trú. 3. Làm thủ tục đổi.","category":"CCCD_HK"},
    "CCCD_CAP_LAI": {"code":"CCCD_CAP_LAI","name":"Cấp lại CCCD khi mất","docs":"Giấy khai sinh, Sổ hộ khẩu, Đơn xin cấp lại","coquan":"Công an quận/huyện","time":"7-15 ngày","lephi":"Miễn phí","note":"Báo mất và xin cấp lại","steps":"1. Làm đơn báo mất. 2. Đến Công an. 3. Xin cấp lại CCCD.","category":"CCCD_HK"},
    "HO_KHAU_TACH": {"code":"HO_KHAU_TACH","name":"Tách sổ hộ khẩu","docs":"Sổ hộ khẩu, CCCD, Đơn xin tách, Giấy tờ nơi ở mới","coquan":"Công an quận/huyện","time":"3-7 ngày","lephi":"Miễn phí","note":"Cần đồng ý chủ hộ","steps":"1. Viết đơn xin tách hộ khẩu. 2. Xin chữ ký chủ hộ. 3. Đến Công án làm thủ tục.","category":"CCCD_HK"},
    "HO_KHAU_NHAP": {"code":"HO_KHAU_NHAP","name":"Nhập sổ hộ khẩu","docs":"Sổ hộ khẩu, CCCD, Đơn xin nhập, Giấy tờ nơi ở","coquan":"Công an quận/huyện","time":"3-7 ngày","lephi":"Miễn phí","note":"Cần đồng ý chủ hộ mới","steps":"1. Viết đơn xin nhập. 2. Xin đồng ý chủ hộ mới. 3. Đến Công an làm thủ tục.","category":"CCCD_HK"},
    "TAM_TRU": {"code":"TAM_TRU","name":"Đăng ký tạm trú","docs":"CCCD, Đơn xin tạm trú, Giấy tờ nơi ở","coquan":"Công an xã/phường","time":"3-5 ngày","lephi":"Miễn phí","note":"Tối đa 2 năm","steps":"1. Lấy giấy xác nhận của chủ nhà. 2. Viết đơn xin tạm trú. 3. Đến Công an xã/phường.","category":"CCCD_HK"},
    "TRU_SO": {"code":"TRU_SO","name":"Đăng ký trú quán","docs":"CCCD, Đơn xin trú quán, Xác nhận tổ dân phố","coquan":"Công an xã/phường","time":"3-5 ngày","lephi":"Miễn phí","note":"Cho người không có hộ khẩu","steps":"1. Xin xác nhận tổ dân phố. 2. Viết đơn xin trú quán. 3. Đến Công an.","category":"CCCD_HK"},
    "TAM_VANG": {"code":"TAM_VANG","name":"Đăng ký tạm vắng","docs":"CCCD, Sổ hộ khẩu, Đơn xin","coquan":"Công an","time":"1-3 ngày","lephi":"Miễn phí","note":"Khi đi khỏi nơi thường trú trên 30 ngày","steps":"1. Làm đơn xin báo tạm vắng. 2. Nộp tại Công an nơi thường trú.","category":"CCCD_HK"},
    # BẰNG LÁI
    "BANG_LAI_A1": {"code":"BANG_LAI_A1","name":"Cấp đổi bằng lái A1","docs":"CCCD, Giấy khai sinh (nếu dưới 18), 1 ảnh 3x4","coquan":"Sở GTVT","time":"7-10 ngày","lephi":"Miễn phí","note":"Từ 18 tuổi, đổi bằng không cần thi lại","steps":"1. Chuẩn bị CCCD, ảnh. 2. Đến Sở GTVT hoặc trung tâm sát hạch. 3. Nộp hồ sơ và nhận bằng mới.","category":"GIAO_THONG"},
    "BANG_LAI_A1_CAP_MOI": {"code":"BANG_LAI_A1_CAP_MOI","name":"Cấp mới bằng lái A1","docs":"CCCD, Giấy khai sinh, Học lý, Giấy khám sức khỏe","coquan":"Sở GTVT","time":"15-20 ngày","lephi":"110.000đ","note":"Phải thi sát hạch","steps":"1. Học lý thuyết và sa hình. 2. Đăng ký thi. 3. Thi sát hạch. 4. Nhận bằng khi đỗ.","category":"GIAO_THONG"},
    "BANG_LAI_A2": {"code":"BANG_LAI_A2","name":"Cấp bằng lái A2","docs":"CCCD, Học lý, Giấy khám sức khỏe","coquan":"Sở GTVT","time":"15-20 ngày","lephi":"150.000đ","note":"Xe mô tô trên 175cc","steps":"1. Học lý thuyết và sa hình A2. 2. Đăng ký thi. 3. Thi và nhận bằng.","category":"GIAO_THONG"},
    "BANG_LAI_B1": {"code":"BANG_LAI_B1","name":"Cấp bằng lái B1","docs":"CCCD, Học lý, Giấy khám sức khỏe","coquan":"Sở GTVT","time":"20-25 ngày","lephi":"350.000đ","note":"Ô tô dưới 9 chỗ","steps":"1. Học lý thuyết, sa hình, đường trường. 2. Đăng ký thi B1. 3. Thi và nhận bằng.","category":"GIAO_THONG"},
    "BANG_LAI_B2": {"code":"BANG_LAI_B2","name":"Cấp bằng lái B2","docs":"CCCD, Học lý, Giấy khám sức khỏe","coquan":"Sở GTVT","time":"20-25 ngày","lephi":"400.000đ","note":"Ô tô kinh doanh","steps":"1. Học lý thuyết, sa hình, đường trường. 2. Đăng ký thi B2. 3. Thi và nhận bằng.","category":"GIAO_THONG"},
    "BANG_LAI_DOI": {"code":"BANG_LAI_DOI","name":"Đổi bằng lái","docs":"GPLX cũ, CCCD, 1 ảnh 3x4","coquan":"Sở GTVT","time":"7-10 ngày","lephi":"135.000đ","note":"Đổi khi hết hạn","steps":"1. Chuẩn bị bằng cũ, CCCD, ảnh. 2. Đến Sở GTVT. 3. Đổi bằng mới.","category":"GIAO_THONG"},
    "BANG_LAI_CAP_LAI": {"code":"BANG_LAI_CAP_LAI","name":"Cấp lại bằng lái","docs":"CCCD, 1 ảnh 3x4, Đơn xin cấp lại","coquan":"Sở GTVT","time":"7-10 ngày","lephi":"135.000đ","note":"Khi bị mất","steps":"1. Làm đơn xin cấp lại. 2. Đến Sở GTVT. 3. Nhận bằng mới.","category":"GIAO_THONG"},
    # HỘ CHIẾU
    "HO_CHIEU_CAP": {"code":"HO_CHIEU_CAP","name":"Cấp hộ chiếu","docs":"CCCD, Sổ hộ khẩu, 4 ảnh 4x6","coquan":"Phòng xuất nhập cảnh","time":"7-15 ngày","lephi":"200.000đ","note":"Hộ chiếu gắn chip 10 năm","steps":"1. Chuẩn bị CCCD, hộ khẩu, ảnh. 2. Làm đơn xin cấp hộ chiếu. 3. Nộp tại Phòng XNC. 4. Nhận hộ chiếu.","category":"HANH_CHINH"},
    "HO_CHIEU_GIA_HAN": {"code":"HO_CHIEU_GIA_HAN","name":"Gia hạn hộ chiếu","docs":"Hộ chiếu cũ, CCCD, 1 ảnh 4x6","coquan":"Phòng xuất nhập cảnh","time":"5-10 ngày","lephi":"200.000đ","note":"Trước 6 tháng hết hạn","steps":"1. Kiểm tra hạn hộ chiếu. 2. Chuẩn bị hồ sơ gia hạn. 3. Nộp tại Phòng XNC.","category":"HANH_CHINH"},
    # LÝ LỊCH
    "LY_LICH_TU_PHAP": {"code":"LLTP","name":"Cấp lý lịch tư pháp","docs":"CCCD, Đơn xin cấp LLTP","coquan":"Sở Tư pháp","time":"3-5 ngày","lephi":"Miễn phí","note":"Kiểm tra án tích","steps":"1. Làm đơn xin cấp lý lịch tư pháp. 2. Nộp tại Sở Tư pháp. 3. Chờ 3-5 ngày và nhận kết quả.","category":"HANH_CHINH"},
    # NHÀ ĐẤT
    "SO_DO": {"code":"SO_DO","name":"Cấp sổ đỏ (Giấy chứng nhận quyền sử dụng đất)","docs":"Đơn xin cấp, Giấy tờ về quyền sử dụng đất","coquan":"Văn phòng đăng ký đất đai","time":"10-30 ngày","lephi":"80.000đ-500.000đ","note":"Cần giấy tờ hợp pháp","steps":"1. Làm đơn xin cấp sổ đỏ. 2. Chuẩn bị giấy tờ chứng minh quyền sử dụng. 3. Nộp tại VP Đăng ký đất đai. 4. Đóng phí và nhận sổ.","category":"NHADAT"},
    "SANG_TEN": {"code":"SANG_TEN","name":"Sang tên sổ đỏ","docs":"Sổ đỏ, CCCD, Hợp đồng chuyển nhượng","coquan":"Văn phòng đăng ký đất đai","time":"10-20 ngày","lephi":"0.5% giá trị","note":"Cần công chứng hợp đồng","steps":"1. Công chứng hợp đồng chuyển nhượng. 2. Cả 2 bên đến VP Đăng ký đất đai. 3. Làm thủ tục sang tên. 4. Đóng lệ phí trước bạ.","category":"NHADAT"},
    # GIAO THÔNG XE
    "DANG_KY_XE": {"code":"DANG_KY_XE","name":"Đăng ký xe ô tô","docs":"CCCD, Giấy chứng nhận quyền sở hữu xe, Giấy bảo hiểm","coquan":"Phòng đăng ký xe","time":"1-3 ngày","lephi":"1.500.000đ-3.000.000đ","note":"Trong 30 ngày kể từ khi mua","steps":"1. Chuẩn bị giấy tờ xe. 2. Mua bảo hiểm xe. 3. Đến Phòng đăng ký xe. 4. Nộp hồ sơ và nhận biển số.","category":"GIAO_THONG"},
    "CHUYEN_NHUONG_XE": {"code":"CHUYEN_NHUONG_XE","name":"Chuyển nhượng xe","docs":"CCCD cả 2 bên, Sổ đăng ký xe, Hợp đồng chuyển nhượng","coquan":"Phòng đăng ký xe","time":"2-5 ngày","lephi":"500.000đ-1.000.000đ","note":"Cả 2 cùng đến làm","steps":"1. Làm hợp đồng chuyển nhượng có công chứng. 2. Cả 2 bên đến Phòng đăng ký xe. 3. Làm thủ tục chuyển nhượng.","category":"GIAO_THONG"},
    # XÂY DỰNG
    "CAP_PHEP_XAY_DUNG": {"code":"PHEP_XD","name":"Cấp phép xây dựng nhà ở","docs":"Đơn xin, Giấy chứng nhận quyền sử dụng đất, Bản vẽ quy hoạch","coquan":"Phòng Đô thị","time":"15-20 ngày","lephi":"50.000đ-200.000đ","note":"Cần quy hoạch","steps":"1. Lấy bản vẽ quy hoạch 1/500. 2. Thiết kế bản vẽ xây dựng. 3. Làm đơn xin cấp phép. 4. Nộp tại Phòng Đô thị.","category":"XAYDUNG"},
    # KINH DOANH
    "CAP_PHEP_KINH_DOANH": {"code":"PHEP_KD","name":"Cấp giấy phép kinh doanh","docs":"Đơn xin cấp phép kinh doanh, CCCD, Địa chỉ kinh doanh","coquan":"Phòng Đăng ký kinh doanh hoặc UBND","time":"3-5 ngày","lephi":"Miễn phí","note":"Kinh doanh cá thể miễn phí","steps":"1. Làm đơn xin cấp phép kinh doanh. 2. Chuẩn bị CCCD, địa chỉ kinh doanh. 3. Nộp tại Phòng ĐKKD hoặc UBND.","category":"KINHDOANH"},
    # Y TẾ
    "BHYT_CAP": {"code":"BHYT_CAP","name":"Cấp thẻ BHYT","docs":"CCCD, Sổ hộ khẩu","coquan":"Bảo hiểm xã hội","time":"3-5 ngày","lephi":"Miễn phí","note":"Trẻ em dưới 6 tuổi miễn phí","steps":"1. Chuẩn bị CCCD, sổ hộ khẩu. 2. Đến BHXH nơi thường trú. 3. Làm thủ tục cấp thẻ BHYT.","category":"YTE"},
    "BHYT_DOI": {"code":"BHYT_DOI","name":"Đổi thẻ BHYT","docs":"Thẻ BHYT cũ, CCCD","coquan":"Bảo hiểm xã hội","time":"1-3 ngày","lephi":"Miễn phí","note":"Đổi khi thẻ hư hỏng hoặc sai thông tin","steps":"1. Mang thẻ BHYT cũ và CCCD. 2. Đến BHXH nơi đã cấp thẻ. 3. Làm thủ tục đổi thẻ.","category":"YTE"},
    # GIÁO DỤC
    "HOC_BONG": {"code":"HOC_BONG","name":"Đăng ký xét học bổng","docs":"Đơn xin học bổng, Bảng tổng kết, Học bạ","coquan":"Trường học hoặc Phòng GD-ĐT","time":"7-10 ngày","lephi":"Miễn phí","note":"Theo từng loại học bổng","steps":"1. Làm đơn xin xét học bổng. 2. Chuẩn bị hồ sơ (bảng tổng kết, học bạ). 3. Nộp tại trường hoặc Phòng GD-ĐT.","category":"GIAODUC"},
    # KHÁC
    "XIN_VIEC": {"code":"XIN_VIEC","name":"Đăng ký tìm việc làm","docs":"CCCD, Hồ sơ xin việc","coquan":"Trung tâm dịch vụ việc làm","time":"1-2 ngày","lephi":"Miễn phí","note":"","steps":"1. Chuẩn bị hồ sơ xin việc (CCCD, sơ yếu lý lịch). 2. Đến Trung tâm dịch vụ việc làm. 3. Đăng ký và được giới thiệu việc làm.","category":"KHAC"},
    "TRO_CAP": {"code":"TRO_CAP","name":"Đăng ký trợ cấp xã hội","docs":"Đơn xin trợ cấp, CCCD, Sổ hộ khẩu","coquan":"UBND xã/phường","time":"7-10 ngày","lephi":"Miễn phí","note":"","steps":"1. Làm đơn xin trợ cấp xã hội. 2. Xin xác nhận hoàn cảnh khó khăn. 3. Nộp tại UBND xã/phường.","category":"KHAC"},
}

# ============================================================
# SITUATIONS CLASSIFICATION
# ============================================================
SITUATIONS = {
    'khai_sinh_moi': ['khai sinh moi', 'sinh con', 'be moi sinh', 'dang ky khai sinh', 'lam khai sinh', 'khai be'],
    'khai_sinh_qua_han': ['khai sinh qua han', 'qua han khai sinh', 'muon khai sinh', 'qua han 60 ngay'],
    'ket_hon': ['ket hon', 'dang ky ket hon', 'dang ky cuoi', 'lay vo chong', 'hon nhan', 'lam cuoi'],
    'ly_hon': ['ly hon', 'ly di', 'don ly hon', 'ly dinh', 'dong y ly hon'],
    'khai_tu': ['khai tu', 'bao tu', 'nguoi chet', 'dang ky khai tu', 'giay bao tu'],
    'ket_hon_ngoai': ['ket hon nguoi nuoc ngoai', 'lay nguoi nuoc ngoai', 'vo chong nuoc ngoai'],
    'cccd_cap': ['cap cccd', 'lam cccd', 'cccd lan dau', 'can cuoc cong dan', 'cmnd', 'lam can cuoc'],
    'cccd_do': ['doi cccd', 'doi can cuoc', 'lam lai cccd', 'cccd het han', 'doi cmnd'],
    'cccd_cap_lai': ['cap lai cccd', 'cccd bi mat', 'lam lai can cuoc', 'mat cccd'],
    'ho_khau_tach': ['tach ho khau', 'tach khau', 'tach ra khoi ho', 'tach ho', 'tach so'],
    'ho_khau_nhap': ['nhap ho khau', 'nhap khau', 'vao ho khau', 'nhap vao ho', 'nhap ho'],
    'tam_tru': ['tam tru', 'dang ky tam tru', 'xin tam tru', 'lam tam tru', 'o thue'],
    'tru_so': ['tru so', 'tru qua han', 'tru qua 90 ngay', 'dang ky tru quan'],
    'tam_vang': ['tam vang', 'dang ky tam vang', 'bao tam vang', 'di tam vang'],
    'bang_lai_a1': ['bang lai a1', 'bang lai xe may', 'lai xe may', 'cap gplx a1', 'lam bang a1', 'bang xe may'],
    'bang_lai_a2': ['bang lai a2', 'bang lai xe cong nong', 'lai xe cong nong'],
    'bang_lai_b1': ['bang lai b1', 'bang lai oto duoi 9 cho', 'lai b1'],
    'bang_lai_b2': ['bang lai b2', 'bang lai oto', 'lai xe o to', 'bang lai dong so'],
    'bang_lai_do': ['doi bang lai', 'lam lai bang lai', 'bang lai mat', 'gia han bang lai'],
    'bang_lai_cap_lai': ['cap lai bang lai', 'bang lai bi mat', 'mat bang lai'],
    'thi_sat_hach': ['thi bang lai', 'thi sat hach', 'thi lai xe', 'dang ky thi bang lai'],
    'ho_chieu_cap': ['lam ho chieu', 'cap ho chieu', 'passport', 'lam passport', 'xin ho chieu'],
    'ho_chieu_gia_han': ['gia han ho chieu', 'doi ho chieu', 'ho chieu het han'],
    'ly_lich_tu_phap': ['ly lich tu phap', 'an tien an', 'so tu phap', 'xet nghiem tien an', 'xin lltp'],
    'so_do': ['so do', 'giay chu quyen', 'sodo', 'cap so do nha', 'lam so do', 'so do', 'giay so'],
    'sang_ten': ['sang ten', 'chuyen nhuong', 'chuyen ten', 'mua ban nha', 'sang ten nha'],
    'kinh_doanh': ['kinh doanh', 'mo cua hang', 'ban hang', 'giay phep kinh doanh', 'mo kinh doanh'],
    'dang_ky_xe': ['dang ky xe', 'dang ki xe o to', 'nghia vu xe', 'dang ky xe moi', 'dang ki xe'],
    'chuyen_nhuong_xe': ['chuyen nhuong xe', 'ban xe', 'mua xe', 'chuyen ten xe'],
    'xay_dung': ['xay dung', 'xay nha', 'phep xay dung', 'cong nhat', 'xin phep xay'],
    'bhyt_cap': ['cap the bhyt', 'lam bao hiem y te', 'dang ky bhyt', 'lam the bao hiem'],
    'bhyt_do': ['doi the bhyt', 'gia han bao hiem y te', 'bhyt het han', 'doi the bao hiem'],
    'hoc_bong': ['dang ky hoc bong', 'xin hoc bong', 'xet hoc bong', 'nop don hoc bong'],
    'xin_viec': ['dang ky tim viec lam', 'tim viec', 'xin viec', 'tim kiem viec lam'],
    'tro_cap': ['dang ky tro cap xa hoi', 'tro cap', 'ho tro', 'xin tro cap'],
    'hoi': ['hoi', 'hoi ve', 'tu van', 'giup doi', 'can hoi', 'tim hieu', 'xin tu van'],
}

# ============================================================
# LOAD DATASET & BUILD DOCUMENTS
# ============================================================
try:
    from datasets import load_dataset
    ds1 = load_dataset('hirine/dataset-thu-tuc-hanh-chinh-5733-samples', split='train')
    df1 = pd.DataFrame(ds1)
except Exception as e:
    df1 = pd.DataFrame({
        'id': range(100),
        'title': ['Khai sinh', 'Kết hôn', 'Làm CCCD', 'Tách hộ khẩu', 'Đổi bằng lái',
                  'Khai tử', 'Hộ chiếu', 'Lý lịch tư pháp', 'Sổ đỏ', 'Sang tên'] * 10,
        'text': ['Thủ tục hành chính cần giấy tờ...' * 5] * 100
    })

try:
    ds2_stream = load_dataset('phamson02/large-vi-legal-queries', split='train', streaming=True)
    data2 = []
    for i, item in enumerate(ds2_stream):
        data2.append(item)
        if i >= 999:
            break
    df2 = pd.DataFrame(data2)
except Exception as e:
    df2 = pd.DataFrame({
        'passage_id': range(500),
        'title': ['Hỏi về CCCD', 'Hỏi về hộ khẩu', 'Hỏi về bằng lái', 'Hỏi về khai sinh'] * 125,
        'context': ['Câu hỏi về thủ tục hành chính...' * 10] * 500
    })

# Build documents
all_docs = []
for i, row in df1.iterrows():
    truncated_text = row['text'][:10000]
    text = f"THỦ TỤC: {row['title']}\nNỘI DUNG: {truncated_text}"
    all_docs.append(Document(page_content=text, metadata={"source": "ds1", "id": i}))

ds2_limit = min(2000, len(df2))
for i, row in df2.head(ds2_limit).iterrows():
    truncated_context = row.get('context', 'Nội dung')[:10000]
    text = f"CÂU HỎI: {row.get('title', 'Hỏi')}\nTRẢ LỜI: {truncated_context}"
    all_docs.append(Document(page_content=text, metadata={"source": "ds2", "id": i}))

for code, proc in DICHVUCONG_DATA.items():
    text = f"""MÃ: {proc['code']}
THỦ TỤC: {proc['name']}
HỒ SƠ: {proc['docs']}
CƠ QUAN: {proc['coquan']}
THỜI GIAN: {proc['time']}
LỆ PHÍ: {proc['lephi']}
LƯU Ý: {proc['note']}
QUY TRÌNH: {proc['steps']}"""
    all_docs.append(Document(page_content=text, metadata={"source": "dvc", "code": code}))

# ============================================================
# EMBEDDINGS & CLASSIFIER
# ============================================================
device = 'cuda' if torch.cuda.is_available() else 'cpu'
embeddings = HuggingFaceEmbeddings(
    model_name='keepitreal/vietnamese-sbert',
    model_kwargs={'device': device},
    encode_kwargs={'normalize_embeddings': True}
)

# Train Classifier
train_text = []
train_label = []
for lbl, exs in SITUATIONS.items():
    for ex in exs:
        train_text.append(ex)
        train_label.append(lbl)

vec = TfidfVectorizer(ngram_range=(1, 2), min_df=1)
X = vec.fit_transform(train_text)
clf = VotingClassifier(
    estimators=[
        ('rf', RandomForestClassifier(n_estimators=50, random_state=42)),
        ('lr', LogisticRegression(max_iter=1000, random_state=42))
    ],
    voting='soft'
)
clf.fit(X, train_label)

# BM25
corpus = [doc.page_content for doc in all_docs]
tokenized_corpus = [doc.split() for doc in corpus]
bm25 = BM25Okapi(tokenized_corpus)

# ============================================================
# PINECONE SETUP
# ============================================================
PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY', '')
vector_db = None

if PINECONE_API_KEY:
    try:
        pc = Pinecone(api_key=PINECONE_API_KEY)
        index_name = 'ai-hanh-chinh-rag'
        dimension = 768

        if not pc.has_index(index_name):
            pc.create_index(
                name=index_name,
                dimension=dimension,
                metric='cosine',
                spec={'serverless': {'cloud': 'aws', 'region': 'us-east-1'}}
            )

        vector_db = PineconeVectorStore(
            index_name=index_name,
            embedding=embeddings,
            pinecone_api_key=PINECONE_API_KEY
        )

        index = pc.Index(index_name)
        stats = index.describe_index_stats()
        vector_count = stats.get('total_vector_count', 0)

        if vector_count == 0:
            ids = [f"doc_{doc.metadata.get('source', 'unk')}_{i}" for i, doc in enumerate(all_docs)]
            vector_db.add_documents(documents=all_docs, ids=ids)
    except Exception as e:
        vector_db = None

# ============================================================
# SEARCH FUNCTIONS
# ============================================================
def hybrid_search(query, k_final=5):
    if vector_db:
        dense_results = vector_db.similarity_search_with_score(query, k=10)
    else:
        dense_results = []

    tokenized_query = query.split()
    bm25_scores = bm25.get_scores(tokenized_query)
    top_indices = np.argsort(bm25_scores)[::-1][:10]

    combined = {}
    for doc, score in dense_results:
        doc_id = doc.page_content[:100]
        combined[doc_id] = {'doc': doc, 'score': 1 - score}

    for idx in top_indices:
        if idx < len(all_docs):
            doc_id = all_docs[idx].page_content[:100]
            if doc_id not in combined:
                combined[doc_id] = {'doc': all_docs[idx], 'score': 0.5}

    sorted_docs = sorted(combined.values(), key=lambda x: x['score'], reverse=True)[:k_final]
    return [item['doc'] for item in sorted_docs]

# ============================================================
# LLM SETUP
# ============================================================
MODEL_ID = 'Qwen/Qwen2.5-3B-Instruct'

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map='auto',
    quantization_config=bnb_config,
    trust_remote_code=True
)

llm = pipeline(
    'text-generation',
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    temperature=0.1,
    repetition_penalty=1.1
)

# ============================================================
# VOICE AI SETUP (DISABLED FOR PYTHON 3.13 COMPATIBILITY)
# ============================================================
# whisper_device = 'cuda' if torch.cuda.is_available() else 'cpu'
# whisper_model = whisper.load_model('base', device=whisper_device)

def speech_to_text(audio_path):
    return ""

async def text_to_speech(text, voice='vi-VN-HoaiMyNeural'):
    return None

# ============================================================
# AI PROCESS FUNCTION
# ============================================================
def ai_process(query):
    try:
        torch.cuda.empty_cache()
        gc.collect()

        situation = clf.predict(vec.transform([query]))[0]
        docs = hybrid_search(query, k_final=5)

        proc_info = ""
        for proc_code, proc in DICHVUCONG_DATA.items():
            if proc['code'] == situation.replace('_', '').upper() or proc['code'] in situation.upper():
                proc_info = f"""
**📋 THỦ TỤC:** {proc['name']}
**📝 HỒ SƠ:** {proc['docs']}
**🏢 NƠI LÀM:** {proc['coquan']}
**⏰ THỜI GIAN:** {proc['time']}
**💰 LỆ PHÍ:** {proc['lephi']}
**⚠️ LƯU Ý:** {proc['note']}
**📝 QUY TRÌNH:** {proc['steps']}
"""
                break

        context_parts = []
        for i, doc in enumerate(docs, 1):
            text = doc.page_content.strip()
            if len(text) > 500:
                text = text[:500] + "..."
            context_parts.append(f"[Tài liệu {i}] {text}")
        context = '\n\n'.join(context_parts)

        prompt = f"""Bạn là cán bộ hướng dẫn thủ tục hành chính công tại UBND. Hãy hỗ trợ người dân một cách chi tiết và thân thiện.

TÌNH HUỐNG: {situation.upper().replace('_', ' ')}
{proc_info}
TÀI LIỆU THAM KHẢO:
{context}

CÂU HỎI: {query}

HƯỚNG DẪN:
1. Chào hỏi thân thiện
2. Xác định thủ tục đúng
3. Liệt kê hồ sơ cần chuẩn bị (rõ ràng, từng món)
4. Nơi làm thủ tục
5. Thời gian và lệ phí
6. Lưu ý quan trọng

QUY ĐỊNH:
- Chỉ trả lời dựa trên tài liệu đã có
- Không bịa đặt thông tin
- Dùng ngôn ngữ dễ hiểu, lịch sự
- Nếu không có thông tin, hướng dẫn đến UBND

Bây giờ hãy trả lời:"""

        messages = [
            {"role": "system", "content": "Bạn là cán bộ hành chính công tận tâm. Hỗ trợ người dân chi tiết, chính xác dựa trên tài liệu có."},
            {"role": "user", "content": prompt}
        ]

        text_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        outputs = llm(
            text_prompt,
            max_new_tokens=512,
            temperature=0.1,
            repetition_penalty=1.1,
            do_sample=False
        )

        response_text = outputs[0]["generated_text"]

        if "<|im_start|>assistant\n" in response_text:
            response_text = response_text.split("<|im_start|>assistant\n")[-1]
        if "<|im_end|>" in response_text:
            response_text = response_text.split("<|im_end|")[0]
        response = response_text.strip()

        clean_lines = []
        for line in response.split('\n'):
            if any('\u4e00' <= char <= '\u9fff' for char in line):
                continue
            if "Infrastructure" in line or "Last Updated" in line:
                continue
            clean_lines.append(line)
        final_response = '\n'.join(clean_lines).strip()

        if len(final_response) < 20:
            final_response = "Dạ xin lỗi, hệ thống chưa cập nhật đủ thông tin về trường hợp này. Bác vui lòng mang giấy tờ cá nhân ra trực tiếp UBND Phường/Xã để được hỗ trợ ạ!"

        torch.cuda.empty_cache()
        gc.collect()

        return {
            "situation": situation,
            "response": final_response,
            "docs_count": len(docs)
        }

    except Exception as e:
        return {
            "situation": "error",
            "response": f"Dạ xin lỗi, hệ thống đang gặp sự cố. Vui lòng thử lại hoặc đến trực tiếp UBND để được hỗ trợ ạ!",
            "docs_count": 0
        }

# ============================================================
# GRADIO INTERFACE
# ============================================================

def process_chat(text, history):
    try:
        if history is None:
            history = []

        query = text.strip() if text else ""

        if len(query) < 3:
            return history, ''

        result = ai_process(query)

        response = f"""**🏛️ AI HÀNH CHÍNH CÔNG VIP PRO**

**📋 Tình huống:** `{result['situation']}`

{result['response']}

---
_📄 Tài liệu tham khảo: {result['docs_count']}_"""

        history.append((query, response))

        return history, ''

    except Exception as e:
        return history, str(e)

def clear_chat():
    return [], ''

pro_css = """
.gradio-container {
    font-family: 'Segoe UI', sans-serif;
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
}
.chatbot {
    border-radius: 20px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1);
}
.pro-badge {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 5px 15px;
    border-radius: 20px;
    font-weight: bold;
}
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue='purple', secondary_hue='indigo'),
               css=pro_css) as demo:

    gr.HTML("""
    <div style="text-align: center; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 20px; margin-bottom: 25px; box-shadow: 0 10px 30px rgba(0,0,0,0.2);">
        <h1 style="color: white; margin: 0; font-size: 2.8em; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);">🏛️ AI HÀNH CHÍNH CÔNG VIP PRO</h1>
        <p style="color: white; margin: 15px 0 0 0; font-size: 1.2em;">Version 2.0 - Tư vấn thủ tục hành chính bằng AI</p>
        <div style="margin-top: 15px;">
            <span class="pro-badge">Hybrid RAG</span>
            <span class="pro-badge">50+ Situations</span>
            <span class="pro-badge">Qwen 2.5 3B</span>
        </div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=3):
            chat = gr.Chatbot(height=500, label='💬 Lịch sử chat',
                             show_copy_button=True)

        with gr.Column(scale=1):
            msg_in = gr.Textbox(label='✏️ Nhập câu hỏi', lines=4,
                               placeholder='Ví dụ: Làm bằng lái xe máy cần giấy tờ gì?')

            with gr.Row():
                btn_send = gr.Button('🚀 GỬI', variant='primary', size='lg')
                btn_clear = gr.Button('🔄 LÀM MỚI', size='lg')

    gr.Examples(
        examples=[
            'Làm bằng lái xe máy cần giấy tờ gì?',
            'Khai sinh quá hạn phải làm sao?',
            'Đăng ký xe ô tô mới ở đâu?',
            'Tách hộ khẩu cần những gì?',
            'Làm hộ chiếu bao nhiêu tiền?'
        ],
        inputs=[msg_in]
    )

    btn_send.click(process_chat, [msg_in, chat],
                  [chat, msg_in])
    msg_in.submit(process_chat, [msg_in, chat],
                 [chat, msg_in])
    btn_clear.click(clear_chat, [], [chat, msg_in])

# ============================================================
# LAUNCH APP
# ============================================================
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)