#!/usr/bin/env python3
import sys, json, re, html
from pathlib import Path
from datetime import datetime

S = re.sub
D = re.DOTALL

def h2m(h):
  for tag, md in [("br\\s*/?", "\n"), ("h1", "\n# "), ("h2", "\n## "), ("h3", "\n### "), ("h4", "\n#### ")]:
    h = S(rf"<{tag}[^>]*>(.*?)</{tag}>" if "/" not in tag else rf"<{tag}>", md if "/" in tag else rf"{md}\1\n", h, flags=D)
  for tag, wrap in [("strong", "**"), ("b", "**"), ("em", "*"), ("i", "*"), ("code", "`")]:
    h = S(rf"<{tag}>(.*?)</{tag}>", rf"{wrap}\1{wrap}", h, flags=D)
  h = S(r"<a[^>]+href=[\"']([^\"']*)[\"'][^>]*>(.*?)</a>", r"[\2](\1)", h, flags=D)
  h = S(r"<li>(.*?)</li>", r"- \1", h, flags=D)
  h = S(r"</?(?:ul|ol|p|div|span|table|thead|tbody|tr)[^>]*>", "\n", h)
  for tag in ("th", "td"): h = S(rf"<{tag}[^>]*>(.*?)</{tag}>", r"| \1 ", h, flags=D)
  return S(r"\n{3,}", "\n\n", html.unescape(S(r"<[^>]+>", "", h))).strip()

def safename(t): return (S(r"[^a-zA-Z0-9\u4e00-\u9fff]+", "-", t).strip("-") or "untitled")[:80]

def title_of(item):
  raw = item.get("title", "untitled")
  if raw.startswith("Created Gemini Canvas titled "): return raw[29:]
  if raw.startswith("Prompted "): return raw[9:].split("\n")[0].strip()[:60]
  return raw

def user_prompt(item):
  raw = item.get("title", "")
  if raw.startswith("Prompted "): return raw[9:].replace("\xa0", " ").strip()
  return ""

def content_of(item):
  if item.get("title", "").startswith("Created Gemini Canvas"):
    parts = [s["name"] for s in item.get("subtitles", []) if s.get("name")]
  else:
    parts = []
    prompt = user_prompt(item)
    if prompt: parts.append(f"## User\n\n{prompt}")
    replies = [m for s in item.get("safeHtmlItem", []) if s.get("html") for m in [h2m(s["html"])] if m]
    if replies: parts.append("## Model\n\n" + "\n\n---\n\n".join(replies))
  return "\n\n".join(parts)

def find_json(src):
  for name in ["我的活动/Gemini Apps/我的活动记录.json", "Gemini Apps/我的活动记录.json", "我的活动记录.json",
                "My Activity/Gemini Apps/My Activity.json", "Gemini Apps/My Activity.json", "My Activity.json"]:
    if (p := src / name).exists(): return p
  if src.is_file() and src.suffix == ".json": return src
  sys.exit(f"Error: no activity JSON in {src}")

def main():
  if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
    sys.exit("gemini-takeout - Convert Google Gemini chats to Markdown\n\nUsage: gemini-takeout <dir> [outdir]")

  src, out = Path(sys.argv[1]).resolve(), Path(sys.argv[2] if len(sys.argv) > 2 else "gemini").resolve()
  assert src.exists(), f"{src} not found"

  jf = find_json(src)
  print(f"Reading {jf.name}...")
  data = json.loads(jf.read_text(encoding="utf-8"))
  assert isinstance(data, list), "expected JSON array"

  out.mkdir(exist_ok=True)
  n = att = 0

  for item in data:
    title, content = title_of(item), content_of(item)
    if not content: continue

    t = item.get("time", "")
    dt = datetime.fromisoformat(t.replace("Z", "+00:00")) if t else datetime.now()
    dp, typ = dt.strftime("%Y-%m-%d"), "canvas" if item.get("title", "").startswith("Created Gemini Canvas") else "prompt"

    md = f"---\nprovider: gemini\ncreated: {dt.strftime('%Y-%m-%dT%H:%M:%S.000Z')}\ntype: {typ}\n---\n\n# {title}\n\n{content}\n"

    for af in item.get("attachedFiles", []):
      if (af_src := jf.parent / af).exists():
        (af_out := out / af).exists() or af_out.write_bytes(af_src.read_bytes())
        md += f"\n- [{af}]({af})"
        att += 1

    safe = safename(title)
    fname = f"{dp}-{safe}.chat.md"
    idx = 2
    while (out / fname).exists(): fname, idx = f"{dp}-{safe}-{idx}.chat.md", idx + 1
    (out / fname).write_text(md, encoding="utf-8")
    n += 1
    print(f"✓ {fname}")

  print(f"\nDone! {n} chats → {out}/" + (f" ({att} attached files copied)" if att else ""))

if __name__ == "__main__": main()
