ViGent2/backend/app/services/uploader/xiaohongshu_uploader.py

"""
Xiaohongshu (小红书) uploader using Playwright
Based on social-auto-upload implementation
"""
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any
import asyncio

from playwright.async_api import Playwright, async_playwright
from loguru import logger

from .base_uploader import BaseUploader
from .cookie_utils import set_init_script


class XiaohongshuUploader(BaseUploader):
    """Xiaohongshu video uploader using Playwright"""
    
    # 超时配置 (秒)
    UPLOAD_TIMEOUT = 300  # 视频上传超时
    PUBLISH_TIMEOUT = 120  # 发布检测超时
    POLL_INTERVAL = 1  # 轮询间隔
    
    def __init__(
        self,
        title: str,
        file_path: str,
        tags: List[str],
        publish_date: Optional[datetime] = None,
        account_file: Optional[str] = None,
        description: str = ""
    ):
        super().__init__(title, file_path, tags, publish_date, account_file, description)
        self.upload_url = "https://creator.xiaohongshu.com/publish/publish?from=homepage&target=video"
    
    async def set_schedule_time(self, page, publish_date):
        """Set scheduled publish time"""
        try:
            logger.info("[小红书] 正在设置定时发布时间...")
            
            # Click "定时发布" label
            label_element = page.locator("label:has-text('定时发布')")
            await label_element.click()
            await asyncio.sleep(1)
            
            # Format time
            publish_date_hour = publish_date.strftime("%Y-%m-%d %H:%M")
            
            # Fill datetime input
            await page.locator('.el-input__inner[placeholder="选择日期和时间"]').click()
            await page.keyboard.press("Control+KeyA")
            await page.keyboard.type(str(publish_date_hour))
            await page.keyboard.press("Enter")
            
            await asyncio.sleep(1)
            logger.info(f"[小红书] 已设置定时发布: {publish_date_hour}")
            
        except Exception as e:
            logger.error(f"[小红书] 设置定时发布失败: {e}")
    
    async def upload(self, playwright: Playwright) -> dict:
        """Main upload logic with guaranteed resource cleanup"""
        browser = None
        context = None
        try:
            # Launch browser (headless for server deployment)
            browser = await playwright.chromium.launch(headless=True)
            context = await browser.new_context(
                viewport={"width": 1600, "height": 900},
                storage_state=self.account_file
            )
            context = await set_init_script(context)
            
            page = await context.new_page()
            
            # Go to upload page
            await page.goto(self.upload_url)
            logger.info(f"[小红书] 正在上传: {self.file_path.name}")
            
            # Upload video file
            await page.locator("div[class^='upload-content'] input[class='upload-input']").set_input_files(str(self.file_path))
            
            # Wait for upload to complete (with timeout)
            import time
            upload_start = time.time()
            while time.time() - upload_start < self.UPLOAD_TIMEOUT:
                try:
                    upload_input = await page.wait_for_selector('input.upload-input', timeout=3000)
                    preview_new = await upload_input.query_selector(
                        'xpath=following-sibling::div[contains(@class, "preview-new")]'
                    )
                    
                    if preview_new:
                        stage_elements = await preview_new.query_selector_all('div.stage')
                        upload_success = False
                        
                        for stage in stage_elements:
                            text_content = await page.evaluate('(element) => element.textContent', stage)
                            if '上传成功' in text_content:
                                upload_success = True
                                break
                        
                        if upload_success:
                            logger.info("[小红书] 检测到上传成功标识")
                            break
                        else:
                            logger.info("[小红书] 未找到上传成功标识，继续等待...")
                    else:
                        logger.info("[小红书] 未找到预览元素，继续等待...")
                    
                    await asyncio.sleep(self.POLL_INTERVAL)
                    
                except Exception as e:
                    logger.info(f"[小红书] 检测过程: {str(e)}，重新尝试...")
                    await asyncio.sleep(0.5)
            else:
                logger.error("[小红书] 视频上传超时")
                return {
                    "success": False,
                    "message": "视频上传超时",
                    "url": None
                }
            
            # Fill title and tags
            await asyncio.sleep(1)
            logger.info("[小红书] 正在填充标题和话题...")
            
            title_container = page.locator('div.plugin.title-container').locator('input.d-text')
            if await title_container.count():
                await title_container.fill(self.title[:30])
            
            # Add tags
            css_selector = ".tiptap"
            for tag in self.tags:
                await page.type(css_selector, "#" + tag)
                await page.press(css_selector, "Space")
            
            logger.info(f"[小红书] 总共添加 {len(self.tags)} 个话题")
            
            # Set scheduled publish time if needed
            if self.publish_date != 0:
                await self.set_schedule_time(page, self.publish_date)
            
            # Click publish button (with timeout)
            publish_start = time.time()
            while time.time() - publish_start < self.PUBLISH_TIMEOUT:
                try:
                    if self.publish_date != 0:
                        await page.locator('button:has-text("定时发布")').click()
                    else:
                        await page.locator('button:has-text("发布")').click()
                    
                    await page.wait_for_url(
                        "https://creator.xiaohongshu.com/publish/success?**",
                        timeout=3000
                    )
                    logger.success("[小红书] 视频发布成功")
                    break
                except Exception:
                    logger.info("[小红书] 视频正在发布中...")
                    await asyncio.sleep(0.5)
            else:
                logger.warning("[小红书] 发布检测超时，请手动确认")
            
            # Save updated cookies
            await context.storage_state(path=self.account_file)
            logger.success("[小红书] Cookie 更新完毕")
            
            await asyncio.sleep(2)
            
            return {
                "success": True,
                "message": "发布成功，待审核" if self.publish_date == 0 else "已设置定时发布",
                "url": None
            }
            
        except Exception as e:
            logger.exception(f"[小红书] 上传失败: {e}")
            return {
                "success": False,
                "message": f"上传失败: {str(e)}",
                "url": None
            }
        finally:
            # 确保资源释放
            if context:
                try:
                    await context.close()
                except Exception:
                    pass
            if browser:
                try:
                    await browser.close()
                except Exception:
                    pass
    
    async def main(self) -> Dict[str, Any]:
        """Execute upload"""
        async with async_playwright() as playwright:
            return await self.upload(playwright)