import os
import time
import re
import argparse
import traceback
from playwright.sync_api import sync_playwright

URL = "https://creator.xiaohongshu.com/publish/publish"


# =========================
# 等待编辑 iframe 出现
# =========================
def find_editor_frame(page, timeout_sec=30):
    end = time.time() + timeout_sec
    while time.time() < end:
        for f in page.frames:
            try:
                if f.locator("input.upload-input[type='file']").count() > 0:
                    return f
            except:
                pass
        time.sleep(0.5)
    return None


# =========================
# 切换到“上传图文”
# =========================
def click_tab_js(page, tab_text: str) -> bool:
    script = """
(t) => {
  const titles = Array.from(document.querySelectorAll("span.title"));
  const hit = titles.find(x => (x.textContent || "").trim() === t);
  if (!hit) return false;
  const tab = hit.closest(".creator-tab") || hit.parentElement;
  if (!tab) return false;
  tab.click();
  return true;
}
"""
    return bool(page.evaluate(script, tab_text))


# =========================
# 上传图片
# =========================
def upload_image(frame, image_path: str):
    loc = frame.locator("input.upload-input[type='file']")
    for i in range(min(loc.count(), 10)):
        el = loc.nth(i)
        acc = (el.get_attribute("accept") or "").lower()
        if any(x in acc for x in [".jpg", ".jpeg", ".png", ".webp"]):
            el.set_input_files(image_path)
            return True
    return False


# =========================
# 获取图片数量（判断是否上传完成）
# =========================
def get_image_count(page) -> int:
    try:
        body = page.inner_text("body")
        m = re.search(r"图片编辑\s*(\d+)\s*/\s*\d+", body)
        if m:
            return int(m.group(1))
    except:
        pass
    return -1


def wait_image_uploaded(page, start_cnt, timeout_sec=240):
    end = time.time() + timeout_sec
    while time.time() < end:
        cnt = get_image_count(page)
        if cnt >= 1 and cnt != start_cnt:
            return True
        if "上传失败" in page.inner_text("body"):
            return False
        time.sleep(2)
    return False


# =========================
# 填写标题
# =========================
def fill_title(frame, title: str):
    candidates = [
        "input[placeholder*='标题']",
        "textarea[placeholder*='标题']",
        "input[placeholder*='填写标题']",
        "textarea[placeholder*='填写标题']",
    ]
    for sel in candidates:
        try:
            el = frame.locator(sel).first
            if el.count() == 0:
                continue
            el.click(timeout=2000)
            el.fill("")
            el.type(title, delay=10)
            val = el.input_value()
            if val and title[:2] in val:
                return True
        except:
            pass
    return False


# =========================
# 填写正文
# =========================
def fill_body(frame, page, body: str):
    try:
        box = frame.locator("[contenteditable='true']").first
        if box.count() > 0:
            box.click()
            page.keyboard.press("Meta+A")
            page.keyboard.type(body, delay=5)
            return True
    except:
        pass

    try:
        ta = frame.locator("textarea").first
        ta.click()
        ta.fill(body)
        return True
    except:
        return False


# =========================
# 点击发布
# =========================
def click_publish(page):
    try:
        page.locator("button:has-text('发布')").last.click(timeout=5000)
        return True
    except:
        pass

    for t in ["发布", "发布笔记", "立即发布"]:
        try:
            page.locator(f"text={t}").first.click(timeout=3000)
            return True
        except:
            pass

    return False


# =========================
# 主发布函数
# =========================
def publish(image_path: str, title: str, body: str, state_path="xhs_state.json", headless=False):

    image_path = os.path.abspath(image_path)

    if not os.path.exists(image_path):
        raise FileNotFoundError(image_path)

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=headless)
        context = browser.new_context(storage_state=state_path)
        page = context.new_page()

        page.goto(URL, wait_until="domcontentloaded")
        page.wait_for_timeout(4000)

        click_tab_js(page, "上传图文")
        page.wait_for_timeout(1500)

        start_cnt = get_image_count(page)

        frame = find_editor_frame(page)
        if not frame:
            raise RuntimeError("No editor frame")

        if not upload_image(frame, image_path):
            raise RuntimeError("Upload input not found")

        if not wait_image_uploaded(page, start_cnt, timeout_sec=240):
            raise RuntimeError("Image upload not completed")

        if not fill_title(frame, title):
            raise RuntimeError("Title not filled")

        if not fill_body(frame, page, body):
            raise RuntimeError("Body not filled")

        page.wait_for_timeout(1000)

        if not click_publish(page):
            raise RuntimeError("Publish button not found")

        page.wait_for_timeout(2500)

        browser.close()


# =========================
# CLI入口
# =========================
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True)
    parser.add_argument("--title", required=True)
    parser.add_argument("--body", required=True)
    parser.add_argument("--state", default="xhs_state.json")
    parser.add_argument("--headless", action="store_true")

    args = parser.parse_args()

    try:
        publish(args.image, args.title, args.body, state_path=args.state, headless=args.headless)
        print("OK")
    except Exception as e:
        print("ERR:", e)
        traceback.print_exc()
        raise