# -*- coding: utf-8 -*-
"""
Chatbot tra cứu Bộ pháp điển Việt Nam trên Hugging Face Spaces.
Dataset mặc định: tmquan/phapdien-moj-gov-vn / subset articles.

Thiết kế an toàn pháp lý:
- Không tự bịa quy định.
- Luôn trả lời kèm điều/đề mục/nguồn/URL.
- Có thể chạy hoàn toàn không cần API key: SQLite FTS5 + Gradio.
"""

from __future__ import annotations

import html
import json
import os
import re
import sqlite3
import textwrap
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence

import gradio as gr
from datasets import load_dataset
from tqdm import tqdm

DATASET_ID = os.getenv("DATASET_ID", "tmquan/phapdien-moj-gov-vn")
DATASET_CONFIG = os.getenv("DATASET_CONFIG", "articles")
DATASET_SPLIT = os.getenv("DATASET_SPLIT", "train")

# Trên Hugging Face Spaces, /data có thể được mount persistent nếu bật persistent storage.
# Nếu không có, dùng thư mục cache chuẩn của HF.
CACHE_ROOT = Path(os.getenv("SPACE_CACHE_DIR", os.getenv("HF_HOME", "/tmp"))) / "phapdien_law_chatbot"
CACHE_ROOT.mkdir(parents=True, exist_ok=True)
INDEX_PATH = CACHE_ROOT / "phapdien_fts.sqlite"

MAX_INDEX_CONTENT_CHARS = int(os.getenv("MAX_INDEX_CONTENT_CHARS", "50000"))
MAX_ANSWER_CONTEXT_CHARS = int(os.getenv("MAX_ANSWER_CONTEXT_CHARS", "1200"))
DEFAULT_TOP_K = int(os.getenv("DEFAULT_TOP_K", "5"))

STOPWORDS = {
    "và", "hoặc", "là", "của", "có", "cho", "về", "theo", "được", "trong", "khi",
    "với", "một", "các", "những", "này", "đó", "tôi", "muốn", "hỏi", "quy", "định",
    "pháp", "luật", "điều", "khoản", "nếu", "thì", "không", "ai", "gì", "như", "nào",
}

APP_CSS = """
.gradio-container { max-width: 1180px !important; }
.source-card {
    border: 1px solid #e5e7eb;
    border-radius: 14px;
    padding: 12px 14px;
    margin: 10px 0;
    background: #ffffff;
}
.source-title { font-weight: 700; font-size: 1.02rem; }
.source-meta { color: #4b5563; font-size: 0.92rem; margin-top: 4px; }
.source-text { margin-top: 8px; line-height: 1.55; }
.notice { padding: 12px 14px; border-radius: 14px; background: #fff7ed; border: 1px solid #fed7aa; }
"""


@dataclass
class SearchResult:
    id: int
    subject_title: str
    topic_title: str
    article_title: str
    chapter_title: str
    source_note_text: str
    content_text: str
    source_url: str
    score: float


def clean_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False)
    return str(value).strip()


def compact_text(value: str) -> str:
    return re.sub(r"\s+", " ", clean_text(value)).strip()


def tokenize_query(query: str, *, max_terms: int = 14) -> List[str]:
    text = query.lower().strip()
    tokens = re.findall(r"[0-9a-zA-ZÀ-ỹĐđ]+", text, flags=re.UNICODE)
    filtered: List[str] = []
    for token in tokens:
        token = token.strip()
        if len(token) < 2:
            continue
        if token in STOPWORDS:
            continue
        if token not in filtered:
            filtered.append(token)
    return filtered[:max_terms]


def build_fts_query(query: str) -> str:
    """Tạo query FTS5 đơn giản, an toàn hơn cho tiếng Việt.

    FTS5 hỗ trợ prefix query dạng token*. Ta lọc ký tự đặc biệt để tránh lỗi MATCH.
    """
    tokens = tokenize_query(query)
    if not tokens:
        fallback = re.sub(r'[^0-9a-zA-ZÀ-ỹĐđ\s]', ' ', query, flags=re.UNICODE).strip()
        tokens = tokenize_query(fallback, max_terms=6)
    if not tokens:
        return '""'
    return " OR ".join(f"{token}*" for token in tokens)


def connect_db() -> sqlite3.Connection:
    conn = sqlite3.connect(str(INDEX_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL;")
    conn.execute("PRAGMA synchronous=NORMAL;")
    return conn


def index_exists() -> bool:
    if not INDEX_PATH.exists() or INDEX_PATH.stat().st_size < 4096:
        return False
    try:
        conn = connect_db()
        count = conn.execute("SELECT COUNT(*) FROM articles").fetchone()[0]
        conn.close()
        return count > 100
    except Exception:
        return False


def init_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        DROP TABLE IF EXISTS articles;
        DROP TABLE IF EXISTS articles_fts;

        CREATE TABLE articles (
            id INTEGER PRIMARY KEY,
            subject_id TEXT,
            topic_id TEXT,
            topic_number INTEGER,
            topic_title TEXT,
            subject_number INTEGER,
            subject_title TEXT,
            article_anchor TEXT,
            article_title TEXT,
            chapter_title TEXT,
            source_note_text TEXT,
            source_links TEXT,
            related_note_text TEXT,
            content_text TEXT,
            content_char_len INTEGER,
            content_word_count INTEGER,
            source_url TEXT,
            scraped_at TEXT
        );

        CREATE VIRTUAL TABLE articles_fts USING fts5(
            article_title,
            subject_title,
            topic_title,
            chapter_title,
            source_note_text,
            content_text,
            content='articles',
            content_rowid='id',
            tokenize='unicode61 remove_diacritics 2'
        );
        """
    )
    conn.commit()


def build_index(force: bool = False) -> str:
    if index_exists() and not force:
        return f"Đã có chỉ mục tại: {INDEX_PATH}"

    conn = connect_db()
    init_schema(conn)

    ds = load_dataset(DATASET_ID, DATASET_CONFIG, split=DATASET_SPLIT)
    total = len(ds) if hasattr(ds, "__len__") else None

    insert_article_sql = """
        INSERT INTO articles (
            id, subject_id, topic_id, topic_number, topic_title,
            subject_number, subject_title, article_anchor, article_title,
            chapter_title, source_note_text, source_links, related_note_text,
            content_text, content_char_len, content_word_count, source_url, scraped_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """
    insert_fts_sql = """
        INSERT INTO articles_fts(
            rowid, article_title, subject_title, topic_title, chapter_title, source_note_text, content_text
        ) VALUES (?, ?, ?, ?, ?, ?, ?)
    """

    for idx, row in enumerate(tqdm(ds, total=total, desc="Indexing phapdien"), start=1):
        content = compact_text(row.get("content_text", ""))
        if len(content) > MAX_INDEX_CONTENT_CHARS:
            content = content[:MAX_INDEX_CONTENT_CHARS] + " …"

        values = (
            idx,
            clean_text(row.get("subject_id")),
            clean_text(row.get("topic_id")),
            row.get("topic_number") if row.get("topic_number") is not None else None,
            compact_text(row.get("topic_title")),
            row.get("subject_number") if row.get("subject_number") is not None else None,
            compact_text(row.get("subject_title")),
            clean_text(row.get("article_anchor")),
            compact_text(row.get("article_title")),
            compact_text(row.get("chapter_title")),
            compact_text(row.get("source_note_text")),
            clean_text(row.get("source_links")),
            compact_text(row.get("related_note_text")),
            content,
            row.get("content_char_len") if row.get("content_char_len") is not None else len(content),
            row.get("content_word_count") if row.get("content_word_count") is not None else len(content.split()),
            clean_text(row.get("source_url")),
            clean_text(row.get("scraped_at")),
        )
        conn.execute(insert_article_sql, values)
        conn.execute(
            insert_fts_sql,
            (
                idx,
                values[8],
                values[6],
                values[4],
                values[9],
                values[10],
                values[13],
            ),
        )
        if idx % 500 == 0:
            conn.commit()

    conn.commit()
    conn.execute("INSERT INTO articles_fts(articles_fts) VALUES('optimize')")
    conn.commit()
    count = conn.execute("SELECT COUNT(*) FROM articles").fetchone()[0]
    conn.close()
    return f"Đã xây chỉ mục {count:,} điều/khoản tại: {INDEX_PATH}"


def like_fallback(conn: sqlite3.Connection, query: str, top_k: int) -> List[sqlite3.Row]:
    tokens = tokenize_query(query, max_terms=5)
    if not tokens:
        return []
    where = " OR ".join(["(content_text LIKE ? OR article_title LIKE ? OR subject_title LIKE ?)" for _ in tokens])
    params: List[str] = []
    for token in tokens:
        pattern = f"%{token}%"
        params.extend([pattern, pattern, pattern])
    sql = f"""
        SELECT id, subject_title, topic_title, article_title, chapter_title,
               source_note_text, content_text, source_url, 999.0 AS score
        FROM articles
        WHERE {where}
        LIMIT ?
    """
    params.append(top_k)
    return list(conn.execute(sql, params).fetchall())


def search_articles(query: str, top_k: int = DEFAULT_TOP_K) -> List[SearchResult]:
    build_index(force=False)
    conn = connect_db()
    fts_query = build_fts_query(query)
    rows: Sequence[sqlite3.Row]
    try:
        rows = conn.execute(
            """
            SELECT a.id, a.subject_title, a.topic_title, a.article_title, a.chapter_title,
                   a.source_note_text, a.content_text, a.source_url,
                   bm25(articles_fts, 2.2, 1.3, 1.0, 1.0, 1.4, 3.0) AS score
            FROM articles_fts
            JOIN articles a ON a.id = articles_fts.rowid
            WHERE articles_fts MATCH ?
            ORDER BY score
            LIMIT ?
            """,
            (fts_query, top_k),
        ).fetchall()
    except sqlite3.OperationalError:
        rows = []

    if not rows:
        rows = like_fallback(conn, query, top_k)
    conn.close()

    results: List[SearchResult] = []
    for row in rows:
        results.append(
            SearchResult(
                id=int(row["id"]),
                subject_title=clean_text(row["subject_title"]),
                topic_title=clean_text(row["topic_title"]),
                article_title=clean_text(row["article_title"]),
                chapter_title=clean_text(row["chapter_title"]),
                source_note_text=clean_text(row["source_note_text"]),
                content_text=clean_text(row["content_text"]),
                source_url=clean_text(row["source_url"]),
                score=float(row["score"]),
            )
        )
    return results


def make_excerpt(text: str, query: str, max_chars: int = MAX_ANSWER_CONTEXT_CHARS) -> str:
    text = compact_text(text)
    if len(text) <= max_chars:
        return text
    tokens = tokenize_query(query, max_terms=8)
    lower = text.lower()
    positions = [lower.find(token.lower()) for token in tokens if lower.find(token.lower()) >= 0]
    if positions:
        center = min(positions)
        start = max(0, center - max_chars // 3)
    else:
        start = 0
    end = min(len(text), start + max_chars)
    excerpt = text[start:end]
    if start > 0:
        excerpt = "… " + excerpt
    if end < len(text):
        excerpt += " …"
    return excerpt


def render_sources(results: Sequence[SearchResult], query: str) -> str:
    cards = []
    for idx, result in enumerate(results, start=1):
        title = html.escape(result.article_title or "Không có tiêu đề điều")
        subject = html.escape(result.subject_title or "Không rõ đề mục")
        topic = html.escape(result.topic_title or "Không rõ chủ đề")
        chapter = html.escape(result.chapter_title or "")
        source_note = html.escape(result.source_note_text or "")
        excerpt = html.escape(make_excerpt(result.content_text, query))
        url = html.escape(result.source_url or "")
        link = f'<a href="{url}" target="_blank" rel="noopener noreferrer">Mở nguồn pháp điển</a>' if url else "Không có URL"
        cards.append(
            f"""
            <div class="source-card">
              <div class="source-title">{idx}. {title}</div>
              <div class="source-meta"><b>Chủ đề:</b> {topic} · <b>Đề mục:</b> {subject}</div>
              {f'<div class="source-meta"><b>Chương:</b> {chapter}</div>' if chapter else ''}
              {f'<div class="source-meta"><b>Nguồn:</b> {source_note}</div>' if source_note else ''}
              <div class="source-text">{excerpt}</div>
              <div class="source-meta">{link}</div>
            </div>
            """
        )
    return "\n".join(cards)


def synthesize_answer(question: str, results: Sequence[SearchResult]) -> str:
    if not results:
        return (
            "Tôi chưa tìm thấy điều/khoản đủ gần với câu hỏi trong chỉ mục Bộ pháp điển. "
            "Bạn nên thử hỏi bằng từ khóa cụ thể hơn, ví dụ: tên luật, lĩnh vực, hành vi, điều kiện, mức phạt, thủ tục hoặc đối tượng áp dụng."
        )

    top = results[0]
    bullets = []
    for result in results[:3]:
        excerpt = make_excerpt(result.content_text, question, max_chars=520)
        bullets.append(
            f"- **{result.article_title}** — {result.subject_title}. {excerpt}"
        )

    answer = f"""
<div class="notice">
<b>Lưu ý:</b> Đây là công cụ tra cứu và tổng hợp thông tin từ Bộ pháp điển, không thay thế ý kiến tư vấn của luật sư/cơ quan có thẩm quyền. Khi áp dụng vào vụ việc thật, cần kiểm tra văn bản gốc, hiệu lực và tình tiết cụ thể.
</div>

### Trả lời dựa trên kết quả gần nhất

Căn cứ gần nhất tôi tìm thấy là **{top.article_title}** trong đề mục **{top.subject_title}**, thuộc chủ đề **{top.topic_title}**.

Các đoạn liên quan:
{chr(10).join(bullets)}

### Nguồn tra cứu

{render_sources(results, question)}
"""
    return textwrap.dedent(answer).strip()


def chat_fn(message: str, history: List[Dict[str, str]], top_k: int) -> str:
    message = (message or "").strip()
    if not message:
        return "Bạn hãy nhập câu hỏi pháp luật cần tra cứu."
    results = search_articles(message, top_k=int(top_k))
    return synthesize_answer(message, results)


def status_fn() -> str:
    msg = build_index(force=False)
    return f"{msg}\nDataset: {DATASET_ID}/{DATASET_CONFIG}/{DATASET_SPLIT}\nIndex: {INDEX_PATH}"


def rebuild_fn() -> str:
    return build_index(force=True)


EXAMPLES = [
    "Người lao động đơn phương chấm dứt hợp đồng cần báo trước bao lâu?",
    "Điều kiện cấp giấy chứng nhận quyền sử dụng đất là gì?",
    "Xử phạt hành vi vi phạm an toàn thực phẩm như thế nào?",
    "Quyền và nghĩa vụ của công dân trong bảo vệ an ninh quốc gia",
    "Thủ tục đăng ký kết hôn với người nước ngoài",
]

with gr.Blocks(title="Chatbot Luật Việt Nam - Bộ pháp điển") as demo:
    gr.Markdown(
        """
        # ⚖️ Chatbot Luật Việt Nam — RAG từ Bộ pháp điển

        Dữ liệu: `tmquan/phapdien-moj-gov-vn` trên Hugging Face. Ứng dụng dùng SQLite FTS5 để tìm điều/khoản liên quan, sau đó trả lời kèm nguồn.  
        **Nguyên tắc:** không bịa quy định, không thay thế tư vấn pháp lý chuyên nghiệp.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.ChatInterface(
                fn=chat_fn,
                additional_inputs=[gr.Slider(3, 10, value=DEFAULT_TOP_K, step=1, label="Số nguồn trả về")],
                examples=[[example, DEFAULT_TOP_K] for example in EXAMPLES],
                cache_examples=False,
                chatbot=gr.Chatbot(height=560, buttons=["copy"]),
                textbox=gr.Textbox(
                    placeholder="Nhập câu hỏi, ví dụ: Điều kiện cấp sổ đỏ là gì?",
                    label="Câu hỏi pháp luật",
                ),
            )
        with gr.Column(scale=1):
            gr.Markdown("### Trạng thái dữ liệu")
            status = gr.Textbox(label="Index", lines=7, interactive=False)
            refresh_btn = gr.Button("Kiểm tra / tạo chỉ mục")
            rebuild_btn = gr.Button("Tạo lại chỉ mục", variant="secondary")
            refresh_btn.click(status_fn, outputs=status)
            rebuild_btn.click(rebuild_fn, outputs=status)
            gr.Markdown(
                """
                ### Gợi ý hỏi tốt
                - Nêu lĩnh vực: đất đai, lao động, hôn nhân, doanh nghiệp...
                - Nêu hành vi/thủ tục/mức phạt/đối tượng áp dụng.
                - Có thể hỏi theo tên luật hoặc điều khoản cụ thể.
                """
            )

    demo.load(status_fn, outputs=status)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=8).launch(theme=gr.themes.Soft(), css=APP_CSS)