#!/usr/bin/env python3
import subprocess, json, re, time, os, sys, sqlite3
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" # The task we are executing

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

def shell(cmd, timeout=120):
    try:
        res = subprocess.run(cmd, shell=True, capture_output=True, text=True,
                             timeout=timeout, cwd=WORKDIR)
        if res.returncode == 0:
            return res.stdout.strip()
        return None
    except Exception as e:
        log(f"  ⚠️ shell error: {e}")
        return None

def refresh_auth():
    log("🔄 Refreshing NotebookLM auth...")
    r = subprocess.run(["python3", "scripts/notebooklm_refresh.py"],
                       capture_output=True, text=True, cwd=WORKDIR, timeout=60)
    return r.returncode == 0

def kanban_call(action, *args):
    cmd = f"hermes kanban {action} " + " ".join(args)
    return shell(cmd, timeout=30)

def comment(task_id, text):
    # Escape single quotes for shell command formatting
    escaped = text.replace("'", "'\\''")
    return shell(f"hermes kanban comment {task_id} '{escaped}'", timeout=30)

def normalize_title(title):
    t = title.lower()
    t = re.sub(r'\.md$|\.pdf$|\.docx$|\.txt$', '', t)
    t = re.sub(r'[^a-z0-9]', '', t)
    return t

def query_rag(source_id):
    q = (
        "Hãy 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 cụ thể, số liệu benchmark và giải pháp nếu có. "
        "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", "180", NOTEBOOK_ID, q]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=200, cwd=WORKDIR)
        if res.returncode == 0:
            data = json.loads(res.stdout)
            # Remove markdown citations, asterisks, and format cleanly
            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=60, 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("=== THAILAND NOTEBOOKLM SYNC ===")
    
    # Check Auth
    log("\n[Step 1/4] Checking Auth & Doctor...")
    doctor = shell(f"{NLM_BIN} doctor 2>&1", timeout=30)
    if doctor and "Authentication" in doctor and "expired" in doctor:
        if not refresh_auth():
            msg = "❌ Step 1 Error: Auth expired and refresh failed."
            log(msg)
            comment(TASK_ID, msg)
            return
        log("✅ Auth refreshed.")
    elif not doctor:
        log("⚠️ Doctor returned nothing. Attempting refresh...")
        refresh_auth()
    else:
        log("✅ Auth OK.")
    
    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
    log("\n[Step 2/4] Fetching NotebookLM Source List...")
    sources_raw = shell(f"{NLM_BIN} source list {NOTEBOOK_ID} --json 2>&1", timeout=60)
    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
    if not os.path.exists(MASTER_PATH):
        msg = f"❌ Error: Master file not found at {MASTER_PATH}"
        log(msg)
        comment(TASK_ID, msg)
        return
        
    with open(MASTER_PATH, 'r') as f:
        master_content = f.read()

    # Parse current entries from master_tourism_tl.md
    # Master file format:
    # ## Source 1
    # - **ID:** `...`
    # - **Name:** ...
    # - **Category:** ...
    # - **Summary:** ...
    #
    # We will search for entries using regex:
    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()
        norm = normalize_title(name)
        master_map[eid] = {
            "num": source_num,
            "id": eid,
            "name": name,
            "category": category,
            "summary": summary,
            "full_text": full_text
        }
        master_map[norm] = master_map[eid] # Allow lookup by norm title as well

    # Analyze gaps
    to_summarize = []
    ids_updated = 0
    
    for src in sources:
        src_id = src["id"]
        src_title = src["title"]
        src_norm = normalize_title(src_title)
        
        matched = None
        if src_id in master_map:
            matched = master_map[src_id]
        elif src_norm in master_map:
            matched = master_map[src_norm]
            # If matched by title but ID is different, we update the ID
            if matched["id"] != src_id:
                ids_updated += 1
                master_content = master_content.replace(f"- **ID:** `{matched['id']}`", f"- **ID:** `{src_id}`")
                matched["id"] = src_id
                log(f"✏️ Updated ID for {matched['name']} -> {src_id}")
        
        # Check if the summary is missing or needs expansion (e.g. doesn't have "Tóm tắt (VN)" or is too short)
        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 (new or stale/missing VN summary).")
    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 (up to 10 sources per run)
    log("\n[Step 3/4] Running RAG Summarization...")
    # Limit to 10 sources as requested by Task 1 prompt
    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]}")
        
        # RAG query for VN
        vn_summary = query_rag(sid)
        if not vn_summary:
            log(f"⚠️ RAG query failed for {stitle[:40]}")
            errors_count += 1
            continue
            
        # Check if there are dashes in VN prose and clean them
        # Let's do a basic strip/replace if it contains list items
        # Clean up any dashes
        vn_lines = vn_summary.split("\n")
        cleaned_vn_lines = []
        for line in vn_lines:
            stripped = line.strip()
            if stripped.startswith("- ") and not stripped.startswith("- **"):
                # Replace bullet dash with standard Vietnamese prose bullet or just text
                line = line.replace("- ", "• ", 1)
            cleaned_vn_lines.append(line)
        vn_summary = "\n".join(cleaned_vn_lines)
        
        # Get EN summary + keywords
        en_summary, keywords = get_describe(sid)
        kw_str = ", ".join([f"`{k}`" for k in keywords]) if keywords else ""
        
        # Format the summary body
        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:
            # We are updating an existing entry
            # Let's locate the full block and replace its summary
            old_block = entry["full_text"]
            # Reconstruct the block
            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:
            # We are adding a new entry
            next_num = len(master_map) // 2 + summarized_count + 1 # rough estimate
            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(3) # Delay to avoid rate limit

    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...")
    
    # Dash check in Tóm tắt (VN)
    # Search for all "Tóm tắt (VN):" and scan for starting "-" lines
    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
    # Move status of t_e92a1e30 to completed
    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()