#!/usr/bin/env python3
"""
Export DeepSeek shared conversation to Markdown.

Usage (just run it — venv is created automatically):
    python deepseek_export.py <share_url_or_id> [-o output.md] [--json] [--txt] [--token TOKEN]

Examples:
    python deepseek_export.py https://chat.deepseek.com/share/nvy7v2ps1r6e2wyoyj
    python deepseek_export.py nvy7v2ps1r6e2wyoyj -o my_chat.md
    python deepseek_export.py nvy7v2ps1r6e2wyoyj --json -o chat.json
    python deepseek_export.py nvy7v2ps1r6e2wyoyj --txt

First run will create .venv/ and install requests via uv.
"""

# ── auto-venv bootstrap ─────────────────────────────────────────────────────
import os
import subprocess
import sys
from pathlib import Path

_SCRIPT_DIR = Path(__file__).resolve().parent
_VENV_DIR = _SCRIPT_DIR / ".venv"
_VENV_PYTHON = _VENV_DIR / ("Scripts/python.exe" if os.name == "nt" else "bin/python")


def _ensure_venv():
    """Create venv + install deps if missing, then re-exec inside it."""
    if sys.prefix != sys.base_prefix:
        return  # already inside a venv

    if not _VENV_PYTHON.exists():
        print("[boot] Creating .venv …", flush=True)
        subprocess.check_call(["uv", "venv", str(_VENV_DIR)], cwd=_SCRIPT_DIR)
        print("[boot] Installing dependencies …", flush=True)
        subprocess.check_call(
            ["uv", "pip", "install", "--python", str(_VENV_PYTHON), "requests"],
            cwd=_SCRIPT_DIR,
        )

    os.execv(str(_VENV_PYTHON), [str(_VENV_PYTHON), *sys.argv])


_ensure_venv()
# ── end bootstrap ────────────────────────────────────────────────────────────

import argparse
import json
import re
from datetime import datetime, timezone

import requests

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

SHARE_URL_RE = re.compile(
    r"(?:https?://chat\.deepseek\.com)?/share/(?P<id>[A-Za-z0-9_-]+)"
)


def extract_share_id(raw: str) -> str:
    m = SHARE_URL_RE.search(raw)
    if m:
        return m.group("id")
    if re.fullmatch(r"[A-Za-z0-9_-]{6,}", raw):
        return raw
    raise ValueError(f"Cannot extract share id from: {raw}")


HEADERS = {
    "accept": "application/json, text/plain, */*",
    "accept-language": "en-US,en;q=0.9",
    "referer": "https://chat.deepseek.com/",
    "user-agent": (
        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36"
    ),
}


# ---------------------------------------------------------------------------
# Strategy 1 — plain requests (works when WAF is not aggressive)
# ---------------------------------------------------------------------------

def fetch_via_requests(share_id: str, token: str | None = None) -> dict:
    s = requests.Session()
    s.headers.update(HEADERS)
    s.get("https://chat.deepseek.com/", timeout=15)

    if token:
        s.headers["authorization"] = f"Bearer {token}"

    url = f"https://chat.deepseek.com/api/v0/share/content?share_id={share_id}"
    r = s.get(url, timeout=15)
    r.raise_for_status()
    data = r.json()

    if data.get("code") != 0:
        raise RuntimeError(f"API error: {data.get('msg', data)}")
    return data


# ---------------------------------------------------------------------------
# Strategy 2 — Playwright headless browser (bypasses WAF reliably)
# ---------------------------------------------------------------------------

def fetch_via_playwright(share_id: str, token: str | None = None) -> dict:
    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        raise RuntimeError(
            "playwright not installed. Run:\n"
            "  uv pip install playwright && uv run playwright install chromium"
        )

    captured = {}

    def on_response(response):
        if "/api/v0/share/content" in response.url:
            try:
                captured["data"] = response.json()
            except Exception:
                pass

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        ctx = browser.new_context(
            user_agent=HEADERS["user-agent"],
            viewport={"width": 1280, "height": 800},
        )
        page = ctx.new_page()
        page.on("response", on_response)

        if token:
            ctx.add_cookies([{
                "name": "ds_session_id",
                "value": token,
                "domain": "chat.deepseek.com",
                "path": "/",
            }])

        page.goto(
            f"https://chat.deepseek.com/share/{share_id}",
            wait_until="networkidle",
            timeout=30000,
        )
        page.wait_for_timeout(3000)
        browser.close()

    if not captured.get("data"):
        raise RuntimeError("Playwright could not capture the share API response")
    data = captured["data"]
    if data.get("code") != 0:
        raise RuntimeError(f"API error: {data.get('msg', data)}")
    return data


# ---------------------------------------------------------------------------
# Main fetch (tries requests first, falls back to playwright)
# ---------------------------------------------------------------------------

def fetch_share(share_id: str, token: str | None = None) -> dict:
    try:
        print("[1/2] Trying direct API…", end=" ", flush=True)
        data = fetch_via_requests(share_id, token)
        print("OK")
        return data
    except Exception as e:
        print(f"failed ({e})")

    print("[2/2] Falling back to Playwright…", end=" ", flush=True)
    try:
        data = fetch_via_playwright(share_id, token)
        print("OK")
        return data
    except Exception as e:
        print(f"failed ({e})")
        raise RuntimeError(
            "Both strategies failed. If WAF is blocking you, install Playwright:\n"
            "  uv pip install playwright && uv run playwright install chromium\n"
            "Or pass your token with --token."
        )


# ---------------------------------------------------------------------------
# Formatting
# ---------------------------------------------------------------------------

def _ts(ts) -> str:
    if not ts:
        return ""
    try:
        return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
    except Exception:
        return str(ts)


def to_markdown(data: dict) -> str:
    biz = data.get("data", {}).get("biz_data", {})
    messages = biz.get("messages", []) or biz.get("chat_messages", [])

    title = biz.get("title") or "DeepSeek Conversation"
    model = biz.get("model_type") or biz.get("model") or ""

    parts = [f"# {title}\n"]
    if model:
        parts.append(f"**Model:** {model}")
    parts.append("")

    if not messages:
        parts.append("_No messages found._\n")
        return "\n".join(parts)

    for msg in messages:
        role = (msg.get("role") or "").upper()
        content = msg.get("content") or msg.get("message") or ""
        if not content and not role:
            continue

        thinking = msg.get("thinking_content") or msg.get("reasoning_content") or ""
        elapsed = msg.get("thinking_elapsed_secs")
        search = msg.get("search_results") or []

        if role == "USER":
            header = "You"
        elif role in ("ASSISTANT", "AI"):
            header = "DeepSeek"
        elif role == "SYSTEM":
            header = "System"
        else:
            header = role.title()

        parts.append(f"## {header}\n")

        if thinking:
            label = f"Thinking ({elapsed}s)" if elapsed else "Thinking"
            parts.append(f"<details><summary>{label}</summary>\n\n{thinking}\n\n</details>\n")

        if search:
            refs = []
            for item in search:
                t = item.get("title") or ""
                u = item.get("url") or ""
                if u:
                    refs.append(f"- [{t}]({u})" if t else f"- {u}")
            if refs:
                parts.append("**Sources:**\n" + "\n".join(refs) + "\n")

        if content:
            parts.append(f"{content}\n")
        parts.append("---\n")

    return "\n".join(parts)


def to_plain_text(data: dict) -> str:
    biz = data.get("data", {}).get("biz_data", {})
    messages = biz.get("messages", []) or biz.get("chat_messages", [])
    title = biz.get("title") or "DeepSeek Conversation"

    parts = [f"=== {title} ===\n"]
    for msg in messages:
        role = (msg.get("role") or "").upper()
        content = msg.get("content") or msg.get("message") or ""
        thinking = msg.get("thinking_content") or msg.get("reasoning_content") or ""
        if role == "USER":
            parts.append(f"[User]\n{content}\n")
        elif role in ("ASSISTANT", "AI"):
            if thinking:
                parts.append(f"[Thinking]\n{thinking}\n")
            parts.append(f"[DeepSeek]\n{content}\n")
        else:
            parts.append(f"[{role}]\n{content}\n")
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description="Export a DeepSeek shared conversation")
    ap.add_argument("share", help="Share URL or share ID")
    ap.add_argument("-o", "--output", help="Output file path (default: auto-named)")
    ap.add_argument("--json", action="store_true", help="Dump raw JSON instead of Markdown")
    ap.add_argument("--txt", action="store_true", help="Dump plain text")
    ap.add_argument("--token", help="Bearer token for auth (optional, needed for private shares)")
    args = ap.parse_args()

    share_id = extract_share_id(args.share)
    print(f"Share ID: {share_id}")

    data = fetch_share(share_id, args.token)

    if args.output:
        out = Path(args.output)
    elif args.json:
        out = Path(f"export_{share_id}.json")
    elif args.txt:
        out = Path(f"export_{share_id}.txt")
    else:
        out = Path(f"export_{share_id}.md")

    if args.json:
        content = json.dumps(data, indent=2, ensure_ascii=False)
    elif args.txt:
        content = to_plain_text(data)
    else:
        content = to_markdown(data)

    out.write_text(content, encoding="utf-8")
    print(f"Saved to {out}  ({len(content)} bytes)")


if __name__ == "__main__":
    main()