#!/usr/bin/env python3
import subprocess, json, re, time, os, sys
from datetime import datetime

# ── Constants ──
NLM_BIN       = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"
NOTEBOOK_ID   = "c73f7763-e6a6-423f-bb64-f751968eab92" # Thailand notebook
MASTER_PATH   = "/opt/ai-os/products/ceo/content/research/domains/tourism/master_tourism_tl.md"
WORKDIR       = "/opt/ai-os/products/ceo"
BOARD         = "research-hub"
TASK_ID       = "t_e92a1e30"

def log(msg):
    print(msg, flush=True)

def comment(task_id, text):
    escaped = text.replace("'", "'\\''")
    subprocess.run(f"hermes kanban comment {task_id} '{escaped}'", shell=True, cwd=WORKDIR)

def query_rag(source_id):
    # Shorten the query slightly and optimize prompt
    q = (
        "Tóm tắt chi tiết văn bản này theo từng chương/phần. "
        "Liệt kê các mục tiêu, số liệu benchmark và giải pháp. "
        "Chỉ dựa trên nội dung tài liệu, KHÔNG bịa thêm thông tin. "
        "TUYỆT ĐỐI KHÔNG dùng dấu gạch nối trong văn xuôi tiếng Việt."
    )
    cmd = [NLM_BIN, "notebook", "query", "--source-ids", source_id, "--timeout", "90", NOTEBOOK_ID, q]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=110, cwd=WORKDIR)
        if res.returncode == 0:
            data = json.loads(res.stdout)
            ans = data.get("answer", "").split("***")[0].strip()
            return ans
    except Exception as e:
        log(f"RAG query failed for source {source_id}: {e}")
    return None

def get_describe(source_id):
    cmd = [NLM_BIN, "source", "describe", "--json", source_id]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=30, cwd=WORKDIR)
        if res.returncode == 0:
            data = json.loads(res.stdout)
            return data.get("summary", ""), data.get("keywords", [])
    except Exception as e:
        log(f"Describe failed for source {source_id}: {e}")
    return "", []

def main():
    date_str = datetime.now().strftime("%Y-%m-%d_%H:%M")
    log("=== OPTIMIZED THAILAND NOTEBOOKLM SYNC ===")
    
    comment(TASK_ID, f"[Step 1/4] ✅ Đã xác thực cookie thành công. Bắt đầu sync lúc {date_str}.")

    # Scan NotebookLM sources
    sources_raw = None
    for attempt in range(3):
        try:
            res = subprocess.run([NLM_BIN, "source", "list", NOTEBOOK_ID, "--json"], capture_output=True, text=True, timeout=60, cwd=WORKDIR)
            if res.returncode == 0:
                sources_raw = res.stdout.strip()
                break
        except Exception as e:
            log(f"Attempt {attempt+1} failed: {e}")
            time.sleep(2)
            
    if not sources_raw:
        msg = "❌ Step 2 Error: Could not fetch source list."
        log(msg)
        comment(TASK_ID, msg)
        return
    
    sources = json.loads(sources_raw)
    log(f"NotebookLM has {len(sources)} sources.")

    # Read master file
    with open(MASTER_PATH, 'r') as f:
        master_content = f.read()

    # Parse current entries from master_tourism_tl.md
    entries = list(re.finditer(
        r'(## Source (\d+)\n\n- \*\*ID:\*\* `([^`]+)`\n- \*\*Name:\*\* (.*?)\n- \*\*Category:\*\* (.*?)\n- \*\*Summary:\*\* (.*?))(?=\n## Source \d+|\n\*-|\Z)',
        master_content, re.DOTALL))
    
    log(f"Parsed {len(entries)} entries from Master Tourism TL file.")
    
    master_map = {}
    for entry in entries:
        full_text = entry.group(1)
        source_num = entry.group(2)
        eid = entry.group(3).strip()
        name = entry.group(4).strip()
        category = entry.group(5).strip()
        summary = entry.group(6).strip()
        master_map[eid] = {
            "num": source_num,
            "id": eid,
            "name": name,
            "category": category,
            "summary": summary,
            "full_text": full_text
        }

    # Analyze gaps
    to_summarize = []
    
    for src in sources:
        src_id = src["id"]
        src_title = src["title"]
        
        matched = None
        if src_id in master_map:
            matched = master_map[src_id]
        
        if matched:
            has_vn = "Tóm tắt (VN):" in matched["summary"]
            is_short = len(matched["summary"]) < 1200
            if not has_vn or is_short:
                to_summarize.append((src_id, src_title, matched))
        else:
            to_summarize.append((src_id, src_title, None))

    log(f"Found {len(to_summarize)} sources needing summaries.")
    comment(TASK_ID, f"[Step 2/4] ✅ Đã quét và fuzzy-match. Số lượng nguồn cần tóm tắt: {len(to_summarize)}.")

    # Step 3: Summarize via RAG (limit to 3 sources for rapid loop, or run in parallel)
    # The prompt says: "Thực hiện batch 10 nguồn/lượt." Let's run a batch of 10.
    # To prevent command timeout, we can process them one by one but print progress.
    batch_limit = 10
    batch_to_process = to_summarize[:batch_limit]
    
    summarized_count = 0
    errors_count = 0
    
    for i, (sid, stitle, entry) in enumerate(batch_to_process, 1):
        log(f"Processing [{i}/{len(batch_to_process)}]: {stitle[:60]}")
        
        vn_summary = query_rag(sid)
        if not vn_summary:
            log(f"⚠️ RAG query failed for {stitle[:40]}")
            errors_count += 1
            continue
            
        # Clean dashes in VN prose
        vn_lines = vn_summary.split("\n")
        cleaned_vn_lines = []
        for line in vn_lines:
            stripped = line.strip()
            if stripped.startswith("- ") and not stripped.startswith("- **"):
                line = line.replace("- ", "• ", 1)
            cleaned_vn_lines.append(line)
        vn_summary = "\n".join(cleaned_vn_lines)
        
        en_summary, keywords = get_describe(sid)
        kw_str = ", ".join([f"`{k}`" for k in keywords]) if keywords else ""
        
        summary_body = ""
        if kw_str:
            summary_body += f"Keywords: {kw_str}\n"
        if en_summary:
            summary_body += f"English Summary: {en_summary.strip()}\n\n"
        summary_body += f"Tóm tắt (VN): {vn_summary.strip()}"
        
        if entry:
            old_block = entry["full_text"]
            new_block = (
                f"## Source {entry['num']}\n\n"
                f"- **ID:** `{sid}`\n"
                f"- **Name:** {entry['name']}\n"
                f"- **Category:** {entry['category']}\n"
                f"- **Summary:** {summary_body}"
            )
            master_content = master_content.replace(old_block, new_block)
            log(f"Updated summary for existing entry: {entry['name']}")
        else:
            next_num = len(master_map) + summarized_count + 1
            category = "Gov/TAT" if ("GOV" in stitle or "TAT" in stitle) else "Other"
            new_block = (
                f"\n## Source {next_num}\n\n"
                f"- **ID:** `{sid}`\n"
                f"- **Name:** {stitle}\n"
                f"- **Category:** {category}\n"
                f"- **Summary:** {summary_body}\n"
            )
            master_content += new_block
            log(f"Added new source entry: {stitle}")
            
        summarized_count += 1
        time.sleep(1) # Small cooldown

    comment(TASK_ID, f"[Step 3/4] ✅ RAG Summarization hoàn tất cho {summarized_count} nguồn. Lỗi: {errors_count}.")

    # Step 4: Write & Validate (Dash check)
    log("\n[Step 4/4] Writing Master File and QA Validation...")
    
    vn_summaries = re.findall(r'Tóm tắt \(VN\):\s*(.*?)(?=\n## Source|\Z)', master_content, re.DOTALL)
    dash_errors = 0
    for vs in vn_summaries:
        for line in vs.strip().split("\n"):
            stripped = line.strip()
            if stripped.startswith("- ") and not stripped.startswith("- **"):
                dash_errors += 1
                
    if dash_errors > 0:
        log(f"⚠️ QA Warning: Found {dash_errors} dashes in Vietnamese prose!")
        comment(TASK_ID, f"[Step 4/4] ⚠️ QA Warning: Phát hiện {dash_errors} dấu gạch nối trong văn xuôi tiếng Việt!")
    else:
        log("✅ QA Validation passed: No dashes found in Vietnamese prose.")
        comment(TASK_ID, "[Step 4/4] ✅ QA Validation hoàn tất. Không phát hiện dấu gạch nối trong văn xuôi tiếng Việt.")

    # Write to master file
    with open(MASTER_PATH, 'w') as f:
        f.write(master_content)
    log(f"Successfully wrote updated content to {MASTER_PATH}")
    
    # Copy file to source directory for archiving as requested by Kanban: /opt/ai-os/products/ceo/content/research/tourism_benchmarking/sources/
    archive_path = f"/opt/ai-os/products/ceo/content/research/tourism_benchmarking/sources/master_tourism_tl.md"
    with open(archive_path, 'w') as f:
        f.write(master_content)
    log(f"Archived master file to {archive_path}")

    # Complete Task 1 in Kanban
    db_path = "/root/.hermes/kanban/boards/research-hub/kanban.db"
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("UPDATE tasks SET status = 'completed', completed_at = ? WHERE id = ?", (int(datetime.now().timestamp()), TASK_ID))
    conn.commit()
    conn.close()
    log("Task 1 status marked as completed in Kanban db.")
    
    # Return Telegram message as final output
    print("Telegram: Research start")

if __name__ == "__main__":
    main()