#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse, subprocess, json, os, re, unicodedata, datetime, sys, time, selectors
from pathlib import Path

# 强制使用 markdownify
try:
    from markdownify import markdownify as md_convert
except Exception:
    print("[错误] 请先安装依赖：pip install -U markdownify tqdm", file=sys.stderr)
    sys.exit(3)

# 可选进度条
try:
    from tqdm import tqdm
except Exception:
    tqdm = None

# --- JXA：使用 NSFileHandle 写 stdout，避免 console.log 兼容性问题 ---
JXA_STREAM_TEMPLATE = r"""
ObjC.import('Foundation');

function writeln(obj){
  var s = (typeof obj === 'string') ? obj : JSON.stringify(obj);
  s = s + "\n";
  var data = $.NSString.stringWithString(s).dataUsingEncoding($.NSUTF8StringEncoding);
  $.NSFileHandle.fileHandleWithStandardOutput.writeData(data);
}

function exportDiariesStream(rootFolderPath, accountFilter, pattern) {
  // 立即心跳，证明脚本已启动
  writeln({type:'hello', ts:(new Date()).toISOString()});

  var Notes = Application('Notes');
  try { Notes.includeStandardAdditions = true; } catch(e) {}
  try { Notes.activate(); } catch(e) {}

  var rx = pattern ? new RegExp(pattern) : null;

  function findFolderByPath(account, pathStr){
    var parts = (pathStr || '').split('/').filter(function(x){ return x.length>0;});
    if (parts.length === 0) return null;
    var current = null;
    var level = account.folders();
    for (var i=0;i<parts.length;i++){
      var name = parts[i];
      var next = null;
      try {
        for (var j=0;j<level.length;j++){
          var f = level[j];
          if (f.name() === name){ next = f; break; }
        }
      } catch(e) {}
      if (!next) return null;
      current = next;
      try { level = next.folders(); } catch(e) { level = []; }
    }
    return current;
  }

  function bfs(root){
    var out = [];
    var q = [{folder: root, rel: [root.name()]}];
    while (q.length){
      var cur = q.shift();
      out.push(cur);
      var subs = [];
      try { subs = cur.folder.folders(); } catch(e) { subs = []; }
      for (var i=0;i<subs.length;i++){
        var sub = subs[i];
        q.push({folder: sub, rel: cur.rel.concat([sub.name()])});
      }
    }
    return out;
  }

  var selected = [];
  try {
    var accs = Notes.accounts();
    for (var ai=0; ai<accs.length; ai++){
      var acc = accs[ai];
      if (accountFilter && acc.name() !== accountFilter) continue;
      var root = findFolderByPath(acc, rootFolderPath);
      if (root) selected.push({acc: acc, root: root});
    }
  } catch(e) {
    writeln({type:'error', stage:'accounts', message:String(e)});
    writeln({type:'done'});
    return;
  }

  writeln({type:'meta', accounts: selected.map(function(r){return r.acc.name();}), rootPath: rootFolderPath});

  if (selected.length === 0){
    writeln({type:'warn', message:'root not found'});
    writeln({type:'done'});
    return;
  }

  for (var si=0; si<selected.length; si++){
    var acc = selected[si].acc;
    var root = selected[si].root;
    var folders = bfs(root);

    for (var fi=0; fi<folders.length; fi++){
      var item = folders[fi];
      var folder = item.folder;
      var relPath = item.rel;

      var fname = '';
      try { fname = folder.name(); } catch(e) { fname=''; }
      if (rx && !rx.test(fname)) continue;

      var notesArr = [];
      try { notesArr = folder.notes(); } catch(e) { notesArr = []; }
      var cnt = 0;
      try { cnt = notesArr.length; } catch(e) {
        try { for (var t=0; t<notesArr.length; t++){ cnt++; } } catch(e2) { cnt = 0; }
      }
      writeln({type:'folder', path: relPath, count: cnt});

      for (var ni=0; ni<cnt; ni++){
        var n = null;
        try { n = notesArr[ni]; } catch(e) { continue; }

        var created=null, modified=null, idVal=null, title='', bodyHtml='';
        var atts=[];
        try { created = n.creationDate(); } catch(e){}
        try { modified = n.modificationDate(); } catch(e){}
        try { idVal = (typeof n.id === 'function') ? n.id() : null; } catch(e){}
        try { title = n.name(); } catch(e){ title=''; }

        try {
          var as = n.attachments();
          for (var ak=0; ak<as.length; ak++){
            var a = as[ak];
            var t=null; try { t = a.contentType ? a.contentType() : null; } catch(e){}
            atts.push({name: a.name(), type: t});
          }
        } catch(e) {}

        try { bodyHtml = n.body(); } catch(e) { bodyHtml=''; }

        writeln({
          type:'note',
          account: acc.name(),
          folderPath: relPath,
          title: title,
          createdISO: (created && created.toISOString) ? created.toISOString() : null,
          modifiedISO: (modified && modified.toISOString) ? modified.toISOString() : null,
          id: idVal,
          bodyHtml: bodyHtml,
          attachments: atts
        });
      }
    }
  }

  writeln({type:'done'});
}
exportDiariesStream(%s, %s, %s);
"""

def build_jxa(root_folder_path: str, account_filter: str, pattern: str) -> str:
    return JXA_STREAM_TEMPLATE % (
        json.dumps(root_folder_path),
        json.dumps(account_filter) if account_filter else 'null',
        json.dumps(pattern if pattern is not None else '.*')
    )

# 带“握手超时”的流式读取
def stream_jxa(root_folder_path: str, account_filter: str, pattern: str, handshake_timeout: int = 30):
    script = build_jxa(root_folder_path, account_filter, pattern)
    p = subprocess.Popen(
        ["osascript", "-l", "JavaScript", "-e", script],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, bufsize=1
    )
    sel = selectors.DefaultSelector()
    sel.register(p.stdout, selectors.EVENT_READ)
    sel.register(p.stderr, selectors.EVENT_READ)
    start = time.perf_counter()
    any_output = False
    stderr_buf = ""

    try:
        done = False
        while not done:
            events = sel.select(timeout=0.5)
            if not events:
                if not any_output and (time.perf_counter() - start) > handshake_timeout:
                    try: p.kill()
                    except Exception: pass
                    raise RuntimeError(f"JXA handshake 超时（>{handshake_timeout}s）。stderr: {stderr_buf.strip()}")
                continue

            for key, _ in events:
                if key.fileobj is p.stdout:
                    line = p.stdout.readline()
                    if line == "":
                        done = True
                        break
                    line = line.strip()
                    if not line:
                        continue
                    any_output = True
                    if line.startswith("{") and line.endswith("}"):
                        try:
                            obj = json.loads(line)
                            yield obj
                        except Exception:
                            pass
                else:
                    chunk = p.stderr.read()
                    if chunk:
                        stderr_buf += chunk

        ret = p.wait()
        if ret != 0:
            raise RuntimeError(f"osascript 退出码 {ret}。stderr: {stderr_buf.strip()}")
    finally:
        try:
            sel.unregister(p.stdout); sel.unregister(p.stderr)
        except Exception:
            pass
        try:
            p.stdout.close(); p.stderr.close()
        except Exception:
            pass

def slugify(s: str, maxlen=160) -> str:
    s = '' if s is None else str(s)
    s = unicodedata.normalize('NFKC', s).strip()
    s = s.replace('\u200b', '').replace('\u200c', '').replace('\u200d', '')
    s = re.sub(r'[\\/:*?"<>|]+', '-', s)
    s = re.sub(r'\s+', ' ', s)
    s = (s or 'untitled')[:maxlen]
    return s

def ensure_unique_path(base_path: Path) -> Path:
    if not base_path.exists():
        return base_path
    stem, suffix = base_path.stem, base_path.suffix
    parent = base_path.parent
    i = 1
    while True:
        p = parent / f"{stem}-{i}{suffix}"
        if not p.exists():
            return p
        i += 1

def html_to_md(html: str) -> str:
    return md_convert(html or '', heading_style="ATX")

def write_note_md(note: dict, out_root: Path, front_matter=True, max_title_len=160) -> Path:
    rel_parts_raw = note.get('folderPath') or []
    if not isinstance(rel_parts_raw, list):
        rel_parts_raw = [str(rel_parts_raw)]
    rel_parts = rel_parts_raw[1:]  # 跳过 root
    rel_parts = [slugify(p, 80) for p in rel_parts if p]
    target_dir = out_root.joinpath(*rel_parts)
    target_dir.mkdir(parents=True, exist_ok=True)

    title = note.get('title') or 'Untitled'
    created_iso = note.get('createdISO') or ''
    date_prefix = ''
    if created_iso:
        try:
            dt = datetime.datetime.fromisoformat(created_iso.replace('Z', '+00:00'))
            date_prefix = dt.strftime('%Y-%m-%d')
        except Exception:
            pass

    fname_core = slugify(title, max_title_len)
    fname = f"{date_prefix}-{fname_core}" if date_prefix else fname_core
    out_path = ensure_unique_path(target_dir / f"{fname}.md")

    body_md = html_to_md(note.get('bodyHtml') or '')
    if front_matter:
        fm = [
            "---",
            f"title: {title}",
            f"account: {note.get('account')}",
            f"folder: {'/'.join(note.get('folderPath') or [])}",
            f"created: {note.get('createdISO')}",
            f"modified: {note.get('modifiedISO')}",
            f"id: {note.get('id') or ''}",
            f"attachments: {len(note.get('attachments') or [])}",
            "source: Apple Notes",
            "---",
            ""
        ]
        content = "\n".join(fm) + body_md
    else:
        content = body_md

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

def quick_accounts_applescript():
    script = 'tell application "Notes" to get name of accounts'
    p = subprocess.run(['osascript', '-e', script], capture_output=True, text=True)
    if p.returncode != 0:
        return False, f"osascript error: {p.stderr.strip()}"
    out = p.stdout.strip()
    if not out:
        return True, {'accounts': []}
    parts = [x.strip() for x in out.split(',') if x.strip()]
    return True, {'accounts': parts}

def prewarm_notes():
    try:
        subprocess.run(["open", "-a", "Notes"], check=False)
        time.sleep(2)
    except Exception:
        pass

def main():
    ap = argparse.ArgumentParser(description="Export Apple Notes with realtime progress (markdownify)")
    ap.add_argument('-r','--root-folder', required=True, help="起始目录路径，支持多级，如 'Journal/2025/Before'")
    ap.add_argument('-a','--account', default=None, help="账户过滤，如 'iCloud' 或 'On My Mac'")
    ap.add_argument('-o','--output', default=os.path.expanduser('~/Desktop/NotesExport'), help='输出目录')
    ap.add_argument('-p','--pattern', default=r'.*', help="匹配文件夹名的正则；默认导出所有（含 root 自身）")
    ap.add_argument('--no-front', action='store_true', help='不写 Front Matter')
    ap.add_argument('--max-title-len', type=int, default=160, help='文件名标题最大长度')
    ap.add_argument('--handshake-timeout', type=int, default=30, help='JXA 首次输出超时时间（秒）')
    args = ap.parse_args()

    out_root = Path(args.output).expanduser().resolve()
    out_root.mkdir(parents=True, exist_ok=True)

    ok, info = quick_accounts_applescript()
    if not ok:
        print("[诊断] 无法读取 Notes 账户：", info, file=sys.stderr)
        print("建议：系统设置 > 隐私与安全性 > 自动化，允许“终端/你的运行器”控制“备忘录”；手动打开备忘录预热。")
        prewarm_notes()
    else:
        print(f"[诊断] 账户列表：{', '.join(info.get('accounts', [])) or '(空)'}")

    print("正在扫描并导出（流式）…")
    t0 = time.perf_counter()
    total_known = 0
    exported = 0
    pbar = None

    try:
        for obj in stream_jxa(args.root_folder, args.account, args.pattern, handshake_timeout=args.handshake_timeout):
            typ = obj.get('type')
            if typ == 'hello':
                print(f"[JXA] 启动成功 {obj.get('ts')}")
            elif typ == 'meta':
                accs = obj.get('accounts') or []
                root_path = obj.get('rootPath') or args.root_folder
                print(f"[JXA] 账户: {', '.join(accs) if accs else '(无)'} | 根路径: {root_path}")
                if tqdm:
                    pbar = tqdm(total=0, desc="导出", unit="篇", smoothing=0.1)
                else:
                    print("[提示] 未检测到 tqdm，将以简易计数显示。安装：pip install tqdm")
            elif typ == 'folder':
                count = int(obj.get('count') or 0)
                total_known += count
                if pbar:
                    pbar.total = total_known
                    pbar.refresh()
                print(f"[JXA] 文件夹 {'/'.join(obj.get('path') or [])} | 预计 {count} 篇")
            elif typ == 'note':
                write_note_md(obj, out_root, front_matter=not args.no_front, max_title_len=args.max_title_len)
                exported += 1
                if pbar:
                    pbar.update(1)
                elif exported % 20 == 0 or (total_known and exported == total_known):
                    print(f"进度：{exported}/{total_known}")
            elif typ == 'warn':
                print(f"[警告] {obj.get('message')}")
            elif typ == 'error':
                print(f"[错误@{obj.get('stage')}] {obj.get('message')}", file=sys.stderr)
            elif typ == 'done':
                break
    except RuntimeError as e:
        print(f"[致命] {e}", file=sys.stderr)
        print("操作建议：确认自动化权限已允许；手动打开备忘录等待同步完成；或缩小匹配范围重试。")
        sys.exit(2)
    finally:
        if pbar:
            pbar.close()

    dt = time.perf_counter() - t0
    print(f"完成：导出 {exported}/{total_known if total_known else exported} 条，输出目录：{out_root}（总耗时 {dt:.1f} 秒）")

if __name__ == '__main__':
    main()