#!/usr/bin/env python3
"""
Tipitaka Query Helper
=====================
Helper functions สำหรับ query ฐานข้อมูลพระไตรปิฎก

Usage:
    from tipitaka_query import TipitakaDB
    
    db = TipitakaDB("tipitaka_mcu.db")
    
    # ค้นหาคำ
    results = db.search("อริยสัจ", limit=10)
    
    # ดึงหน้า
    page = db.get_page(volume=10, page=50)
    
    # ดึงตามเลขข้อ
    page = db.get_by_item(volume=10, item=123)
    
    # ดึงสารบัญ
    toc = db.get_toc(volume=10)
"""

import re
import sqlite3
from pathlib import Path
from typing import Optional, List, Dict, Any
from dataclasses import dataclass


def clean_text(
    text: str, 
    remove_line_numbers: bool = True, 
    remove_footnotes: bool = True,
    remove_item_numbers: bool = False,
    for_tts: bool = False
) -> str:
    """
    ทำความสะอาดข้อความพระไตรปิฎก (รองรับ Clean Dataset)
    
    Args:
        text: ข้อความดิบ
        remove_line_numbers: ลบเลขบรรทัด (001, 002, ...)
        remove_footnotes: ลบเชิงอรรถ (บรรทัดที่ขึ้นต้นด้วย @ หรือ [เชิงอรรถ])
        remove_item_numbers: ลบเลขข้อ [๑], [๒๓], ๑- เป็นต้น
        for_tts: เตรียมสำหรับ TTS (ลบทุกอย่างที่ไม่ควรอ่าน)
    
    Returns:
        ข้อความที่ clean แล้ว
    """
    if for_tts:
        remove_line_numbers = True
        remove_footnotes = True
        remove_item_numbers = True
    
    lines = text.split('\n')
    cleaned = []
    
    for line in lines:
        # 1. ลบเลขบรรทัด (ถ้ามีหลุดมา)
        if remove_line_numbers:
            line = re.sub(r'^\s*\d{3}\s+', '', line)
        
        # 2. ข้ามบรรทัดเชิงอรรถ
        # รองรับทั้ง @ (แบบเก่า) และ [เชิงอรรถ] (แบบคลีน)
        s_line = line.strip()
        if remove_footnotes:
            if s_line.startswith('@') or s_line.startswith('[เชิงอรรถ]'):
                continue
        
        # 3. ลบเลขข้อ และ footnote markers
        if remove_item_numbers:
            # ลบ [๑], [1]
            line = re.sub(r'\[[\u0E50-\u0E59]+\]', '', line)
            line = re.sub(r'\[\d+\]', '', line)
            # ลบ ๑-, ๑-๒, 1-, 1-2
            line = re.sub(r'[\u0E50-\u0E59]+-[\u0E50-\u0E59]*', '', line)
            line = re.sub(r'\d+-\d*', '', line)
        else:
            # เก็บ [๑] ไว้ แต่ลบ ๑- (footnote reference) ทิ้ง
            # ต้องระวังไม่ให้ลบ ๑- ที่เป็นส่วนหนึ่งของเลขข้อ เช่น "๑-๕. เรื่อง..."
            # ปกติ footnote reference จะอยู่หลังคำ/ประโยคทันที โดยไม่มีช่องว่าง
            line = re.sub(r'(?<=[^\s])[\u0E50-\u0E59]+-', '', line)
        
        cleaned.append(line)
    
    # รวมบรรทัดและลบช่องว่างซ้ำ
    result = '\n'.join(cleaned)
    result = re.sub(r'\n{3,}', '\n\n', result)
    result = re.sub(r'  +', ' ', result)
    
    return result.strip()


@dataclass
class SearchResult:
    """ผลลัพธ์การค้นหา"""
    volume_number: int
    volume_title: str
    page_number: int
    snippet: str
    rank: float


@dataclass
class PageContent:
    """เนื้อหาหน้า"""
    volume_number: int
    volume_title: str
    page_number: int
    content_text: str
    content_html: Optional[str] = None


@dataclass
class TocItem:
    """รายการสารบัญ"""
    title: str
    page_number: int
    level: int


class TipitakaDB:
    """Interface สำหรับ query ฐานข้อมูลพระไตรปิฎก"""
    
    def __init__(self, db_path: str | Path):
        self.db_path = Path(db_path)
        if not self.db_path.exists():
            raise FileNotFoundError(f"Database not found: {db_path}")
        self._conn = None
    
    @property
    def conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self._conn = sqlite3.connect(self.db_path)
            self._conn.row_factory = sqlite3.Row
        return self._conn
    
    def close(self):
        if self._conn:
            self._conn.close()
            self._conn = None
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
    
    # ========== Search Functions ==========
    
    def search(self, query: str, limit: int = 20, volume: Optional[int] = None) -> List[SearchResult]:
        """
        ค้นหาคำในพระไตรปิฎก (Full-Text Search)
        
        Args:
            query: คำที่ต้องการค้นหา
            limit: จำนวนผลลัพธ์สูงสุด
            volume: จำกัดการค้นหาในเล่มที่ระบุ (optional)
        
        Returns:
            List of SearchResult
        """
        sql = """
            SELECT 
                v.volume_number,
                v.title as volume_title,
                p.page_number,
                snippet(pages_fts, 0, '>>>', '<<<', '...', 50) as snippet,
                rank
            FROM pages_fts
            JOIN pages p ON pages_fts.rowid = p.id
            JOIN volumes v ON p.volume_id = v.id
            WHERE pages_fts MATCH ?
        """
        params = [query]
        
        if volume:
            sql += " AND v.volume_number = ?"
            params.append(volume)
        
        sql += " ORDER BY rank LIMIT ?"
        params.append(limit)
        
        cursor = self.conn.execute(sql, params)
        
        return [
            SearchResult(
                volume_number=row['volume_number'],
                volume_title=row['volume_title'],
                page_number=row['page_number'],
                snippet=row['snippet'],
                rank=row['rank']
            )
            for row in cursor.fetchall()
        ]
    
    def search_simple(self, keyword: str, limit: int = 20) -> List[Dict[str, Any]]:
        """
        ค้นหาแบบง่าย (LIKE) - ใช้เมื่อ FTS ไม่ทำงาน
        
        Returns:
            List of dicts with volume, page, excerpt
        """
        sql = """
            SELECT 
                v.volume_number,
                v.title as volume_title,
                p.page_number,
                substr(p.content_text, max(1, instr(p.content_text, ?) - 50), 150) as excerpt
            FROM pages p
            JOIN volumes v ON p.volume_id = v.id
            WHERE p.content_text LIKE ?
            LIMIT ?
        """
        cursor = self.conn.execute(sql, [keyword, f"%{keyword}%", limit])
        
        return [dict(row) for row in cursor.fetchall()]
    
    # ========== Get Functions ==========
    
    def get_page(self, volume: int, page: int, include_html: bool = False) -> Optional[PageContent]:
        """
        ดึงเนื้อหาหน้าที่ระบุ
        
        Args:
            volume: เลขเล่ม (1-45)
            page: เลขหน้า
            include_html: รวม HTML content ด้วยหรือไม่
        """
        sql = """
            SELECT 
                v.volume_number,
                v.title as volume_title,
                p.page_number,
                p.content_text,
                p.content_html
            FROM pages p
            JOIN volumes v ON p.volume_id = v.id
            WHERE v.volume_number = ? AND p.page_number = ?
        """
        cursor = self.conn.execute(sql, [volume, page])
        row = cursor.fetchone()
        
        if not row:
            return None
        
        return PageContent(
            volume_number=row['volume_number'],
            volume_title=row['volume_title'],
            page_number=row['page_number'],
            content_text=row['content_text'],
            content_html=row['content_html'] if include_html else None
        )
    
    def get_pages(self, volume: int, start_page: int, end_page: int) -> List[PageContent]:
        """ดึงหลายหน้าพร้อมกัน"""
        sql = """
            SELECT 
                v.volume_number,
                v.title as volume_title,
                p.page_number,
                p.content_text
            FROM pages p
            JOIN volumes v ON p.volume_id = v.id
            WHERE v.volume_number = ? AND p.page_number BETWEEN ? AND ?
            ORDER BY p.page_number
        """
        cursor = self.conn.execute(sql, [volume, start_page, end_page])
        
        return [
            PageContent(
                volume_number=row['volume_number'],
                volume_title=row['volume_title'],
                page_number=row['page_number'],
                content_text=row['content_text']
            )
            for row in cursor.fetchall()
        ]
    
    def get_section_text(
        self, 
        volume: int, 
        start_page: int, 
        end_page: int,
        clean: bool = True,
        remove_line_numbers: bool = True,
        remove_footnotes: bool = True,
        remove_item_numbers: bool = False,
        for_tts: bool = False
    ) -> str:
        """
        ดึงเนื้อหาหลายหน้าและรวมเป็น string เดียว
        
        Args:
            volume: เลขเล่ม (1-45)
            start_page: หน้าเริ่มต้น
            end_page: หน้าสิ้นสุด
            clean: ทำความสะอาดข้อความ
            remove_line_numbers: ลบเลขบรรทัด
            remove_footnotes: ลบเชิงอรรถ
            remove_item_numbers: ลบเลขข้อ [๑], [๒๓]
            for_tts: เตรียมสำหรับ TTS (ลบทุกอย่าง)
        
        Returns:
            เนื้อหาทั้งหมดรวมกันเป็น string
        """
        pages = self.get_pages(volume, start_page, end_page)
        
        texts = []
        for p in pages:
            text = p.content_text
            if clean:
                text = clean_text(
                    text, 
                    remove_line_numbers, 
                    remove_footnotes,
                    remove_item_numbers,
                    for_tts
                )
            texts.append(text)
        
        return "\n\n".join(texts)
    
    def get_by_item(self, volume: int, item_number: int) -> Optional[PageContent]:
        """
        ดึงเนื้อหาตามเลขข้อ/เลขย่อหน้า
        
        Args:
            volume: เลขเล่ม (1-45)
            item_number: เลขข้อ
        """
        sql = """
            SELECT page_number
            FROM item_numbers i
            JOIN volumes v ON i.volume_id = v.id
            WHERE v.volume_number = ? AND i.item_number = ?
        """
        cursor = self.conn.execute(sql, [volume, item_number])
        row = cursor.fetchone()
        
        if not row:
            return None
        
        return self.get_page(volume, row['page_number'])
    
    # ========== Section Search ==========
    
    def find_section(self, title: str, volume: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        ค้นหา section/พระสูตรจากชื่อในสารบัญ
        
        Args:
            title: ชื่อที่ต้องการค้นหา (บางส่วนก็ได้)
            volume: จำกัดการค้นหาในเล่มที่ระบุ (optional)
        
        Returns:
            List of dict with volume, title, start_page, end_page
        """
        sql = """
            SELECT 
                v.volume_number,
                v.title as volume_title,
                c.id,
                c.title,
                c.page_number,
                c.level
            FROM contents c
            JOIN volumes v ON c.volume_id = v.id
            WHERE c.title LIKE ?
        """
        params = [f"%{title}%"]
        
        if volume:
            sql += " AND v.volume_number = ?"
            params.append(volume)
        
        sql += " ORDER BY v.volume_number, c.id"
        
        cursor = self.conn.execute(sql, params)
        matches = cursor.fetchall()
        
        results = []
        for match in matches:
            vol = match['volume_number']
            start_page = match['page_number']
            section_id = match['id']
            section_level = match['level']
            
            # หา end_page โดยดู section ถัดไปที่มี level เท่ากันหรือน้อยกว่า
            sql_next = """
                SELECT page_number
                FROM contents c
                JOIN volumes v ON c.volume_id = v.id
                WHERE v.volume_number = ? AND c.id > ? AND c.level <= ?
                ORDER BY c.id
                LIMIT 1
            """
            cursor = self.conn.execute(sql_next, [vol, section_id, section_level])
            next_row = cursor.fetchone()
            
            if next_row:
                end_page = next_row['page_number'] - 1
            else:
                # ถ้าไม่มี section ถัดไป ให้ใช้หน้าสุดท้ายของเล่ม
                sql_max = """
                    SELECT MAX(page_number) as max_page
                    FROM pages p
                    JOIN volumes v ON p.volume_id = v.id
                    WHERE v.volume_number = ?
                """
                cursor = self.conn.execute(sql_max, [vol])
                end_page = cursor.fetchone()['max_page']
            
            results.append({
                'volume': vol,
                'volume_title': match['volume_title'],
                'title': match['title'],
                'start_page': start_page,
                'end_page': end_page,
                'level': section_level
            })
        
        return results
    
    def get_section(
        self, 
        title: str, 
        volume: Optional[int] = None,
        for_tts: bool = False
    ) -> Optional[str]:
        """
        ดึงเนื้อหาพระสูตร/section ทั้งหมดจากชื่อ
        
        Args:
            title: ชื่อพระสูตร (บางส่วนก็ได้)
            volume: เล่มที่ต้องการ (optional, ถ้าไม่ระบุจะใช้ผลลัพธ์แรก)
            for_tts: เตรียมสำหรับ TTS
        
        Returns:
            เนื้อหาทั้งหมด หรือ None ถ้าไม่พบ
        """
        sections = self.find_section(title, volume)
        
        if not sections:
            return None
        
        # ใช้ผลลัพธ์แรก
        section = sections[0]
        
        return self.get_section_text(
            volume=section['volume'],
            start_page=section['start_page'],
            end_page=section['end_page'],
            for_tts=for_tts
        )
    
    # ========== Table of Contents ==========
    
    def get_toc(self, volume: int) -> List[TocItem]:
        """ดึงสารบัญของเล่มที่ระบุ"""
        sql = """
            SELECT c.title, c.page_number, c.level
            FROM contents c
            JOIN volumes v ON c.volume_id = v.id
            WHERE v.volume_number = ?
            ORDER BY c.id
        """
        cursor = self.conn.execute(sql, [volume])
        
        return [
            TocItem(
                title=row['title'],
                page_number=row['page_number'],
                level=row['level']
            )
            for row in cursor.fetchall()
        ]
    
    # ========== Volume Info ==========
    
    def get_volumes(self) -> List[Dict[str, Any]]:
        """ดึงรายการเล่มทั้งหมด"""
        sql = """
            SELECT volume_number, title, pitaka, total_pages
            FROM volumes
            ORDER BY volume_number
        """
        cursor = self.conn.execute(sql)
        return [dict(row) for row in cursor.fetchall()]
    
    def get_volume_info(self, volume: int) -> Optional[Dict[str, Any]]:
        """ดึงข้อมูลเล่มที่ระบุ"""
        sql = """
            SELECT volume_number, title, pitaka, total_pages
            FROM volumes
            WHERE volume_number = ?
        """
        cursor = self.conn.execute(sql, [volume])
        row = cursor.fetchone()
        return dict(row) if row else None
    
    # ========== Stats ==========
    
    def get_stats(self) -> Dict[str, Any]:
        """ดึงสถิติของฐานข้อมูล"""
        stats = {}
        
        # Volume count
        cursor = self.conn.execute("SELECT COUNT(*) FROM volumes")
        stats['total_volumes'] = cursor.fetchone()[0]
        
        # Page count
        cursor = self.conn.execute("SELECT COUNT(*) FROM pages")
        stats['total_pages'] = cursor.fetchone()[0]
        
        # Content items
        cursor = self.conn.execute("SELECT COUNT(*) FROM contents")
        stats['total_toc_items'] = cursor.fetchone()[0]
        
        # Item numbers
        cursor = self.conn.execute("SELECT COUNT(*) FROM item_numbers")
        stats['total_item_numbers'] = cursor.fetchone()[0]
        
        # By pitaka
        cursor = self.conn.execute("""
            SELECT pitaka, COUNT(*) as count, SUM(total_pages) as pages
            FROM volumes
            GROUP BY pitaka
        """)
        stats['by_pitaka'] = {row['pitaka']: {'volumes': row['count'], 'pages': row['pages']} 
                             for row in cursor.fetchall()}
        
        return stats


# ========== Convenience Functions ==========

def search_text(keyword: str, db_path: str = "tipitaka_mcu.db", limit: int = 10) -> List[Dict]:
    """Quick search function"""
    with TipitakaDB(db_path) as db:
        results = db.search(keyword, limit=limit)
        return [
            {
                'volume': r.volume_number,
                'title': r.volume_title,
                'page': r.page_number,
                'snippet': r.snippet
            }
            for r in results
        ]


def get_page(volume: int, page: int, db_path: str = "tipitaka_mcu.db") -> Optional[str]:
    """Quick get page function"""
    with TipitakaDB(db_path) as db:
        result = db.get_page(volume, page)
        return result.content_text if result else None


def get_toc(volume: int, db_path: str = "tipitaka_mcu.db") -> List[Dict]:
    """Quick get TOC function"""
    with TipitakaDB(db_path) as db:
        items = db.get_toc(volume)
        return [{'title': i.title, 'page': i.page_number, 'level': i.level} for i in items]


# ========== CLI ==========

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Query Tipitaka Database")
    parser.add_argument("--db", default="tipitaka_mcu.db", help="Database path")
    
    subparsers = parser.add_subparsers(dest="command")
    
    # Search command
    search_parser = subparsers.add_parser("search", help="Search text")
    search_parser.add_argument("query", help="Search query")
    search_parser.add_argument("-n", "--limit", type=int, default=10, help="Max results")
    search_parser.add_argument("-v", "--volume", type=int, help="Limit to volume")
    
    # Get page command
    page_parser = subparsers.add_parser("page", help="Get page content")
    page_parser.add_argument("volume", type=int, help="Volume number")
    page_parser.add_argument("page", type=int, help="Page number")
    
    # TOC command
    toc_parser = subparsers.add_parser("toc", help="Get table of contents")
    toc_parser.add_argument("volume", type=int, help="Volume number")
    
    # Stats command
    subparsers.add_parser("stats", help="Show database statistics")
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return
    
    db = TipitakaDB(args.db)
    
    if args.command == "search":
        results = db.search(args.query, limit=args.limit, volume=args.volume)
        for r in results:
            print(f"\n[เล่ม {r.volume_number} หน้า {r.page_number}] {r.volume_title}")
            print(f"  {r.snippet}")
    
    elif args.command == "page":
        page = db.get_page(args.volume, args.page)
        if page:
            print(f"=== {page.volume_title} - หน้า {page.page_number} ===\n")
            print(page.content_text)
        else:
            print("ไม่พบหน้าที่ระบุ")
    
    elif args.command == "toc":
        toc = db.get_toc(args.volume)
        for item in toc:
            indent = "  " * item.level
            print(f"{indent}{item.title} (หน้า {item.page_number})")
    
    elif args.command == "stats":
        stats = db.get_stats()
        print(f"Total volumes: {stats['total_volumes']}")
        print(f"Total pages: {stats['total_pages']}")
        print(f"Total TOC items: {stats['total_toc_items']}")
        print(f"Total item numbers: {stats['total_item_numbers']}")
        print("\nBy Pitaka:")
        for pitaka, data in stats['by_pitaka'].items():
            print(f"  {pitaka}: {data['volumes']} volumes, {data['pages']} pages")
    
    db.close()


if __name__ == "__main__":
    main()