#!/usr/bin/env python3
"""
sync_notebook.py — Đồng bộ hóa các nguồn tài liệu mới từ NotebookLM vào master file.
Chạy trực tiếp từ terminal hoặc thông qua bot Telegram.

Cách dùng:
  python3 scripts/sync_notebook.py --notebook "Báo cáo Du Lịch" --master-file "content/research/domains/tourism/master_tourism.md"
"""
import os
import sys
import argparse
import subprocess
import json
import re
import requests

# Cấu hình 9router cục bộ
NINE_ROUTER_URL = "http://127.0.0.1:20128/v1/chat/completions"
NINE_ROUTER_KEY = os.getenv("LOCAL9R_KEY", "LOCAL9R_KEY")

NLM_PATH = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"

def get_existing_ids(master_path):
    """Đọc master file và trích xuất tất cả ID của nguồn đã tồn tại."""
    if not os.path.exists(master_path):
        return set()
    with open(master_path, 'r', encoding='utf-8') as f:
        content = f.read()
    # Tìm kiếm các mẫu - **ID:** `[uuid-hoặc-hash]`
    ids = re.findall(r'-\s+\*\*ID:\*\*\s+`([^`]+)`', content)
    return set(ids)

def get_existing_categories(master_path):
    """Đọc master file và lấy danh sách các nhóm hiện có (dòng bắt đầu bằng ## Nhóm: )."""
    if not os.path.exists(master_path):
        return []
    categories = []
    with open(master_path, 'r', encoding='utf-8') as f:
        for line in f:
            match = re.match(r'^##\s+Nhóm:\s*(.+)$', line.strip())
            if match:
                categories.append(match.group(1).strip())
    return categories

def resolve_notebook_id(notebook_ref):
    """Resolve notebook title or ID to actual notebook ID."""
    if re.fullmatch(r"[0-9a-fA-F-]{36}", notebook_ref):
        return notebook_ref
    try:
        result = subprocess.run([NLM_PATH, "notebook", "list", "--json"], capture_output=True, text=True, check=True)
        notebooks = json.loads(result.stdout)
        for nb in notebooks:
            if nb.get("id") == notebook_ref or nb.get("title") == notebook_ref:
                return nb.get("id")
        for nb in notebooks:
            if notebook_ref.lower() in str(nb.get("title", "")).lower():
                return nb.get("id")
    except Exception as e:
        print(f"❌ Không thể resolve notebook: {e}", file=sys.stderr)
    return None


def fetch_notebooklm_sources(notebook_ref):
    """Gọi nlm CLI để lấy danh sách sources dạng JSON."""
    notebook_id = resolve_notebook_id(notebook_ref)
    if not notebook_id:
        print(f"❌ Không tìm thấy notebook phù hợp: {notebook_ref}", file=sys.stderr)
        sys.exit(1)
    try:
        cmd = [NLM_PATH, "source", "list", notebook_id, "--json"]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return notebook_id, json.loads(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"❌ Lỗi khi gọi nlm source list: {e.stderr}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"❌ Lỗi xử lý dữ liệu: {e}", file=sys.stderr)
        sys.exit(1)

def fetch_source_content(source_id):
    """Gọi nlm CLI để lấy nội dung thô của source."""
    try:
        cmd = [NLM_PATH, "source", "content", source_id]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"❌ Lỗi khi lấy nội dung source {source_id}: {e.stderr}", file=sys.stderr)
        return None

def call_llm_for_summary(content, categories):
    """Gọi 9router để tóm tắt (EN + VN) và phân nhóm."""
    categories_str = ", ".join([f"'{c}'" for c in categories])
    prompt = f"""Bạn là một chuyên gia phân tích chính sách và nghiên cứu dữ liệu. 
Hãy đọc tài liệu dưới đây và thực hiện các nhiệm vụ sau:
1. Tóm tắt tài liệu bằng tiếng Anh chuyên nghiệp (khoảng 3-5 câu).
2. Dịch/viết tóm tắt bằng tiếng Việt với văn phong trang trọng, học thuật. 
   ⚠️ QUY TẮC BẮT BUỘC: Tuyệt đối KHÔNG sử dụng bất kỳ dấu gạch nối (-) nào trong văn xuôi hoặc danh sách của bản tóm tắt tiếng Việt. Thay vào đó, hãy sử dụng từ ngữ liên kết tự nhiên, dấu phẩy hoặc viết liền mạch.
3. Trích xuất chính xác 5 từ khóa cốt lõi (bằng tiếng Anh).
4. Phân loại tài liệu này vào một trong các nhóm hiện có sau đây: [{categories_str}]. Nếu không phù hợp với nhóm nào, hãy gợi ý một tên nhóm mới (bắt đầu bằng danh từ chung, ví dụ: 'Báo cáo Kinh tế', 'Số liệu thống kê').

Nội dung tài liệu:
\"\"\"
{content[:8000]}
\"\"\"

Hãy trả về kết quả dưới dạng JSON có cấu trúc như sau:
{{
  "category": "Tên nhóm được chọn",
  "keywords": ["kw1", "kw2", "kw3", "kw4", "kw5"],
  "summary_en": "Bản tóm tắt tiếng Anh",
  "summary_vi": "Bản tóm tắt tiếng Việt (TUYỆT ĐỐI KHÔNG CÓ DẤU GẠCH NỐI)"
}}
"""
    headers = {
        "Authorization": f"Bearer {NINE_ROUTER_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "anthropic/claude-3-5-sonnet", # Default model chất lượng cao
        "messages": [
            {"role": "system", "content": "You are a helpful assistant that returns only valid JSON matching the requested schema."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.2,
        "response_format": {"type": "json_object"}
    }
    
    try:
        r = requests.post(NINE_ROUTER_URL, headers=headers, json=payload, timeout=60)
        r.raise_for_status()
        res_json = r.json()
        raw_text = res_json['choices'][0]['message']['content']
        return json.loads(raw_text)
    except Exception as e:
        print(f"⚠️ Lỗi gọi LLM: {e}. Thử fallback sang local model...", file=sys.stderr)
        # Fallback thử dùng mô hình local khác nếu có
        payload["model"] = "openai/gpt-4o-mini"
        try:
            r = requests.post(NINE_ROUTER_URL, headers=headers, json=payload, timeout=60)
            res_json = r.json()
            return json.loads(res_json['choices'][0]['message']['content'])
        except Exception as ex:
            print(f"❌ Cả 2 model đều thất bại: {ex}", file=sys.stderr)
            return None

def format_entry(title, source_id, category, keywords, summary_en, summary_vi):
    """Định dạng block markdown theo chuẩn master file."""
    # Đảm bảo summary_vi thực sự không có dấu gạch nối
    summary_vi_clean = summary_vi.replace(" - ", " ").replace("-", " ")
    kw_str = ", ".join([f"`{k}`" for k in keywords])
    return f"""### {title}
- **ID:** `{source_id}`
- **Nhóm:** {category}
- **Từ khóa:** ** {kw_str}
- **Tóm tắt:** ** {summary_en}
- **Tóm tắt (VN):** ** {summary_vi_clean}

"""

def patch_master_file(master_path, new_entries_by_cat):
    """Chèn các block mới vào đúng vị trí dưới các tiêu đề nhóm tương ứng."""
    with open(master_path, 'r', encoding='utf-8') as f:
        content = f.read()

    for category, entries in new_entries_by_cat.items():
        cat_header = f"## Nhóm: {category}"
        # Tìm tiêu đề nhóm trong file
        if cat_header in content:
            # Chèn ngay bên dưới tiêu đề nhóm
            pattern = re.escape(cat_header) + r'\n'
            replacement = f"{cat_header}\n\n" + "".join(entries)
            content = re.sub(pattern, replacement, content, count=1)
        else:
            # Nếu nhóm này chưa có trong file, append vào cuối file
            content += f"\n\n## Nhóm: {category}\n\n" + "".join(entries)

    with open(master_path, 'w', encoding='utf-8') as f:
        f.write(content)

def main():
    parser = argparse.ArgumentParser(description="Sync NotebookLM sources to Master Markdown file")
    parser.add_argument("--notebook", required=True, help="Notebook name or ID")
    parser.add_argument("--master-file", required=True, help="Path to the master markdown file")
    parser.add_argument("--limit", type=int, default=0, help="Optional limit for testing (0 = all)")
    args = parser.parse_args()

    master_path = args.master_file
    if not os.path.exists(master_path):
        print(f"❌ Không tìm thấy master file tại {master_path}", file=sys.stderr)
        sys.exit(1)

    print(f"📖 Đọc master file: {master_path}...")
    existing_ids = get_existing_ids(master_path)
    categories = get_existing_categories(master_path)
    print(f"✅ Đã tìm thấy {len(existing_ids)} nguồn hiện có và {len(categories)} nhóm.")

    print(f"🌐 Đang lấy danh sách nguồn từ NotebookLM: '{args.notebook}'...")
    notebook_id, all_sources = fetch_notebooklm_sources(args.notebook)
    print(f"✅ NotebookLM notebook id: {notebook_id}")
    print(f"✅ NotebookLM có {len(all_sources)} nguồn.")

    # Lọc nguồn mới
    new_sources = [s for s in all_sources if s['id'] not in existing_ids]
    if args.limit and args.limit > 0:
        new_sources = new_sources[:args.limit]
    if not new_sources:
        print("🎉 Không có nguồn mới nào cần đồng bộ.")
        sys.exit(0)

    print(f"🔔 Phát hiện {len(new_sources)} nguồn mới cần đồng bộ!")

    new_entries_by_cat = {}
    for idx, src in enumerate(new_sources, 1):
        src_id = src['id']
        title = src.get('title', f"Nguồn_{src_id[:8]}")
        print(f"[{idx}/{len(new_sources)}] Đang xử lý: '{title}' (ID: {src_id})...")
        
        # 1. Lấy nội dung
        raw_content = fetch_source_content(src_id)
        if not raw_content:
            print(f"⚠️ Bỏ qua nguồn do lỗi fetch content: {title}")
            continue

        # 2. Gọi LLM tóm tắt & phân nhóm
        analysis = call_llm_for_summary(raw_content, categories)
        if not analysis:
            print(f"⚠️ Bỏ qua nguồn do lỗi LLM: {title}")
            continue

        # 3. Format entry
        category = analysis.get('category', 'Tin tức & Báo cáo thị trường')
        keywords = analysis.get('keywords', ['news', 'report'])
        summary_en = analysis.get('summary_en', '')
        summary_vi = analysis.get('summary_vi', '')

        entry_md = format_entry(title, src_id, category, keywords, summary_en, summary_vi)
        
        # Lưu vào danh mục tương ứng
        if category not in new_entries_by_cat:
            new_entries_by_cat[category] = []
        new_entries_by_cat[category].append(entry_md)

        # Cập nhật danh sách categories nếu xuất hiện category mới
        if category not in categories:
            categories.append(category)

    if new_entries_by_cat:
        print("✍️ Đang ghi các nguồn mới vào file master...")
        patch_master_file(master_path, new_entries_by_cat)
        print("🚀 Đồng bộ hoàn tất!")
        for cat, entries in new_entries_by_cat.items():
            print(f"  • Nhóm '{cat}': +{len(entries)} nguồn mới.")
    else:
        print("⚠️ Không có nguồn nào được ghi thành công.")

if __name__ == "__main__":
    main()