#!/usr/bin/env python3
import glob
import json
import os
import re
import subprocess
import sys
import time
import tempfile
import urllib.request
from pathlib import Path
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload

NLM_BIN = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"
NLM_PYTHON = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/python3"
STATE_FILE = "/opt/ai-os/products/ceo/projects/learning/data/synced_drive_files.json"
GOOGLE_TOKEN_FILE = os.path.expanduser("~/.hermes/google_token.json")
ROOT_FOLDER_ID = "1BUtCsz-SfdDCYspmwQlvSiyffFP2TFfY" # Class Materials folder
TELEGRAM_CHAT_ID = "-1003707758328"
TELEGRAM_THREAD_ID = 9896

NOTEBOOK_MAPPING = {
    "public_policy": "825ad1eb-2ca9-4576-8dd3-aadf468e5578", # Nhập môn Chính sách công
    "microeconomics": "e8a30630-0df2-4932-8740-9ffe7d8882bc", # Kinh tế học vi mô 1
    "macroeconomics": "99090327-f66b-4d8d-8b9e-a3c00acb599f", # Kinh tế học vĩ mô
    "default": "3e44792c-0dfb-45f4-a918-da80390f051c"       # FSPPM - Courses
}

def get_telegram_token():
    token = os.environ.get("TELEGRAM_BOT_TOKEN")
    if token:
        return token
    for p in glob.glob("/proc/[0-9]*/environ"):
        try:
            with open(p, "rb") as f:
                env_vars = f.read().decode("utf-8", errors="ignore").split("\0")
                for v in env_vars:
                    if v.startswith("TELEGRAM_BOT_TOKEN="):
                        return v.split("=", 1)[1]
        except Exception:
            continue
    return None

def send_telegram_notification(text):
    token = get_telegram_token()
    if not token:
        return False
    data = {
        "chat_id": TELEGRAM_CHAT_ID,
        "message_thread_id": TELEGRAM_THREAD_ID,
        "text": text,
        "parse_mode": "Markdown",
        "disable_web_page_preview": False
    }
    req = urllib.request.Request(
        f"https://api.telegram.org/bot{token}/sendMessage",
        data=json.dumps(data).encode("utf-8"),
        headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(req) as resp:
            res = json.loads(resp.read().decode("utf-8"))
            return res.get("ok", False)
    except Exception as e:
        return False

def refresh_auth():
    refresh_script = "/opt/ai-os/products/ceo/scripts/notebooklm_refresh.py"
    res = subprocess.run([NLM_PYTHON, refresh_script], capture_output=True, text=True)
    if res.returncode == 0:
        storage = Path.home() / ".notebooklm-mcp-cli"
        try:
            with open(storage / "profiles" / "default" / "cookies.json") as f:
                cookies = json.load(f)
            with open(storage / "profiles" / "default" / "metadata.json") as f:
                meta = json.load(f)
            payload = {
                "cookies": cookies,
                "csrf_token": meta.get("csrf_token", ""),
                "session_id": meta.get("session_id", ""),
                "build_label": meta.get("build_label", ""),
                "extracted_at": time.time()
            }
            with open(storage / "auth.json", "w") as f:
                json.dump(payload, f, indent=2)
            return True
        except Exception:
            pass
    return False

def load_state():
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, "r") as f:
            return json.load(f)
    return {}

def save_state(state):
    os.makedirs(os.path.dirname(STATE_FILE), exist_ok=True)
    with open(STATE_FILE, "w") as f:
        json.dump(state, f, indent=2, ensure_ascii=False)

def get_drive_service():
    with open(GOOGLE_TOKEN_FILE) as f:
        token_data = json.load(f)
    creds = Credentials.from_authorized_user_info(token_data)
    return build("drive", "v3", credentials=creds)

def list_drive_files(service, folder_id, path_prefix=""):
    results = service.files().list(
        q=f"'{folder_id}' in parents and trashed = false",
        fields="files(id, name, mimeType, modifiedTime, webViewLink, parents)",
        pageSize=100
    ).execute()
    files = results.get("files", [])
    valid_items = []
    for item in files:
        cur_path = f"{path_prefix}/{item['name']}" if path_prefix else item['name']
        if item["mimeType"] == "application/vnd.google-apps.folder":
            valid_items.extend(list_drive_files(service, item["id"], cur_path))
        else:
            m = item["mimeType"]
            if "pdf" in m or "word" in m or "document" in m or "text" in m or "presentation" in m or "spreadsheet" in m or "excel" in m:
                valid_items.append({
                    "id": item["id"],
                    "name": item["name"],
                    "mimeType": item["mimeType"],
                    "modifiedTime": item["modifiedTime"],
                    "path": cur_path,
                    "parents": item.get("parents", []),
                    "webViewLink": item.get("webViewLink", f"https://drive.google.com/file/d/{item['id']}/view")
                })
    return valid_items

def convert_ms_to_google_format(service, file_id, file_name, mime_type, parent_id=None):
    if mime_type.startswith("application/vnd.google-apps."):
        target_t = "sheets" if "spreadsheet" in mime_type else "doc" if "document" in mime_type else "slides" if "presentation" in mime_type else "pdf"
        return file_id, file_name, mime_type, target_t

    target_mime = None
    target_type = None
    name_lower = file_name.lower()
    
    if "word" in mime_type or name_lower.endswith((".docx", ".doc")):
        target_mime = "application/vnd.google-apps.document"
        target_type = "doc"
    elif "spreadsheet" in mime_type or "excel" in mime_type or name_lower.endswith((".xlsx", ".xls")):
        target_mime = "application/vnd.google-apps.spreadsheet"
        target_type = "sheets"
    elif "presentation" in mime_type or "powerpoint" in mime_type or name_lower.endswith((".pptx", ".ppt")):
        target_mime = "application/vnd.google-apps.presentation"
        target_type = "slides"

    if not target_mime:
        return file_id, file_name, mime_type, target_type

    clean_title = file_name
    for ext in [".docx", ".doc", ".xlsx", ".xls", ".pptx", ".ppt"]:
        if clean_title.lower().endswith(ext):
            clean_title = clean_title[:-len(ext)]
            break
            
    # Check if already converted in parent to prevent duplicate copies
    try:
        pid = parent_id[0] if isinstance(parent_id, list) else parent_id
        p_query = f"'{pid}' in parents and " if pid else ""
        safe_title = clean_title.replace("'", "\'")
        existing = service.files().list(
            q=f"{p_query}name = '{safe_title}' and mimeType = '{target_mime}' and trashed = false",
            fields="files(id)"
        ).execute().get("files", [])
        if existing:
            return existing[0]["id"], clean_title, target_mime, target_type
    except Exception:
        pass

    body = {"name": clean_title, "mimeType": target_mime}
    if parent_id:
        body["parents"] = parent_id

    try:
        converted = service.files().copy(fileId=file_id, body=body).execute()
        return converted["id"], clean_title, target_mime, target_type
    except Exception as e:
        return file_id, file_name, mime_type, target_type

def extract_title_with_prefix(path, original_name):
    parts = path.split('/')
    date_code = ''
    instructor_code = ''
    filename = original_name
    
    # 1. Orientation / General files
    if "So Tay Hoc Vien" in filename or "Library_Rules" in filename or "Online_Databases" in filename or "HuongDanTrichDanTLTK" in filename:
        return f"[G]_{filename}"
    
    # 2. Books
    if "Birkland_2016" in filename or "Kingdon_2013" in filename or "Moore_1995" in filename or "Gomez-Barroso_2023" in filename:
        return f"[B]_{filename}"

    # 3. Path analysis for Date and Instructor
    for part in parts:
        m_date = re.match(r'(?:20)?26\.(\d{2})\.(\d{2})', part)
        if m_date:
            mm, dd = m_date.groups()
            date_code = f"{dd}{mm}"
            
        p_lower = part.lower()
        if 'do thien anh tuan' in p_lower and 'nguyen xuan thanh' in p_lower:
            instructor_code = 'NXT_DTAT'
        elif 'do thien anh tuan' in p_lower:
            instructor_code = 'DTAT'
        elif 'nguyen xuan thanh' in p_lower:
            instructor_code = 'NXT'
        elif 'vu thanh tu anh' in p_lower:
            instructor_code = 'VTTA'
        elif 'huynh nhat nam' in p_lower:
            instructor_code = 'HNN'
            
    prefix = ""
    if date_code and instructor_code:
        prefix = f"[{date_code}-{instructor_code}]"
    elif date_code:
        prefix = f"[{date_code}]"
    elif instructor_code:
        prefix = f"[{instructor_code}]"
        
    if prefix:
        return f"{prefix} {filename}"
    return filename

def determine_notebook(path, name):
    p_lower = path.lower()
    if "nhập môn chính sách công" in p_lower or "nhap mon chinh sach cong" in p_lower:
        return NOTEBOOK_MAPPING["public_policy"], "Nhập môn Chính sách công"
    if "macroeconomics" in p_lower or "vĩ mô" in p_lower :
        return NOTEBOOK_MAPPING["macroeconomics"], "Kinh tế học vĩ mô"
    if "microecnomics" in p_lower or "microeconomics" in p_lower or "vi mô" in p_lower:
        return NOTEBOOK_MAPPING["microeconomics"], "Kinh tế học vi mô 1"
    return NOTEBOOK_MAPPING["default"], "FSPPM - Courses"

def add_source_to_nlm(service, notebook_id, file_id, file_name, mime_type):
    doc_type = "pdf"
    if "presentation" in mime_type:
        doc_type = "slides"
    elif "spreadsheet" in mime_type or "excel" in mime_type:
        doc_type = "sheets"
    elif "word" in mime_type or "document" in mime_type or "text" in mime_type:
        doc_type = "doc"

    cmd = [
        NLM_BIN, "source", "add", notebook_id,
        "--drive", file_id,
        "--type", doc_type,
        "--title", file_name
    ]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    except subprocess.TimeoutExpired:
        print(f"Timeout adding {file_name} via Drive ID {file_id}")
        return False, "Timeout adding source via Drive"

    if res.returncode != 0 and ("Authentication" in res.stdout or "expired" in res.stdout or "auth" in res.stdout.lower()):
        refresh_auth()
        try:
            res = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        except subprocess.TimeoutExpired:
            return False, "Timeout adding source via Drive after auth refresh"

    if res.returncode == 0:
        source_id = None
        for line in res.stdout.splitlines():
            if "Source ID:" in line:
                source_id = line.split("Source ID:")[1].strip()
        return True, source_id

    return False, res.stderr or res.stdout

def main():
    service = get_drive_service()
    state = load_state()
    files = list_drive_files(service, ROOT_FOLDER_ID)

    new_synced = []
    failed_items = []

    for f in files:
        f_id = f["id"]
        if f_id in state:
            continue

        nb_id, course_name = determine_notebook(f["path"], f["name"])
        formatted_title = extract_title_with_prefix(f["path"], f["name"])
        
        use_id, use_name, use_mime, target_type = convert_ms_to_google_format(service, f_id, formatted_title, f["mimeType"], f.get("parents"))
        
        ok, res_or_id = add_source_to_nlm(service, nb_id, use_id, use_name, use_mime)
        
        if ok:
            state[f_id] = {
                "name": formatted_title,
                "original_name": f["name"],
                "path": f["path"],
                "notebook_id": nb_id,
                "course_name": course_name,
                "source_id": res_or_id,
                "synced_at": f["modifiedTime"],
                "webViewLink": f["webViewLink"]
            }
            new_synced.append({
                "name": formatted_title,
                "course": course_name,
                "path": f["path"],
                "link": f["webViewLink"],
                "notebook_url": f"https://notebooklm.google.com/notebook/{nb_id}"
            })
            save_state(state)
        else:
            failed_items.append((formatted_title, res_or_id))

    if new_synced:
        lines = []
        lines.append("🔔 *BẢN TIN ĐỒNG BỘ TÀI LIỆU FSPPM MỚI*\n")
        lines.append(f"Tổng số tài liệu mới vừa nạp thành công: *{len(new_synced)}*")
        lines.append(f"Sổ tay tiếp nhận: https://notebooklm.google.com/notebook/{NOTEBOOK_MAPPING['public_policy']}\n")
        
        groups = {}
        for item in new_synced:
            parts = item["path"].split("/")
            group_key = parts[1] if len(parts) > 2 else "Tài liệu chung"
            if len(parts) > 3 and "Tài liệu đọc" in parts[1]:
                group_key = f"{parts[1]} / {parts[2]}"
            groups.setdefault(group_key, []).append(item)
            
        for group_name, items in groups.items():
            lines.append(f"📁 *Phân mục:* {group_name} ({len(items)} tài liệu)")
            for it in items[:5]:
                lines.append(f"  • {it['name']}")
                lines.append(f"    Xem: {it['link']}")
            if len(items) > 5:
                lines.append(f"  ... và thêm {len(items) - 5} tài liệu khác.")
            lines.append("")
        
        full_msg = "\n".join(lines)
        print(full_msg)
        send_telegram_notification(full_msg)
    else:
        pass

if __name__ == "__main__":
    main()