#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

SUBMIT_ID = '815fc095-2cf3-415d-ba13-2b6d31050593'
BIN = '/Users/bot1/.hermes/profiles/video/home/.local/bin/dreamina'
OUT_DIR = Path('/Users/bot1/Volumes/root_for_ai/AI工作区/通用_产品宣传视频_古钱币杜邦纸钱袋包_20260530_1702/outputs/dreamina-seedance2-full-20260530_194640')
DOWNLOAD_DIR = OUT_DIR / 'downloads'
DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
LOG = OUT_DIR / 'query_poll.log'
STATUS = OUT_DIR / 'query_status.json'

def run_query():
    env = os.environ.copy()
    env['HOME'] = '/Users/bot1'
    cmd = [BIN, 'query_result', '--submit_id', SUBMIT_ID, '--download_dir', str(DOWNLOAD_DIR)]
    res = subprocess.run(cmd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    return res.returncode, res.stdout

def extract_status(text: str) -> str | None:
    m = re.search(r'"gen_status"\s*:\s*"([^"]+)"|gen_status[=:]\s*([A-Za-z_\-]+)', text)
    if m:
        return next(g for g in m.groups() if g)
    # Some outputs only say success/fail in prose.
    low = text.lower()
    if 'success' in low:
        return 'success'
    if 'fail' in low or 'failed' in low:
        return 'fail'
    if 'querying' in low or 'generating' in low:
        return 'querying'
    return None

def media_files():
    exts = {'.mp4', '.mov', '.m4v', '.webm'}
    return sorted([p for p in DOWNLOAD_DIR.rglob('*') if p.suffix.lower() in exts], key=lambda p: p.stat().st_mtime, reverse=True)

max_seconds = int(sys.argv[1]) if len(sys.argv) > 1 else 480
interval = int(sys.argv[2]) if len(sys.argv) > 2 else 20
start = time.time()
last_status = None
attempt = 0
with LOG.open('a', encoding='utf-8') as log:
    log.write(f"\n===== poll start {datetime.now().astimezone().isoformat()} max={max_seconds}s interval={interval}s =====\n")
    while True:
        attempt += 1
        rc, out = run_query()
        now = datetime.now().astimezone().isoformat()
        log.write(f"\n--- attempt {attempt} {now} rc={rc} ---\n{out}\n")
        log.flush()
        last_status = extract_status(out) or last_status
        files = media_files()
        status_obj = {
            'updated_at': now,
            'submit_id': SUBMIT_ID,
            'returncode': rc,
            'gen_status': last_status,
            'download_dir': str(DOWNLOAD_DIR),
            'media_files': [str(p) for p in files],
            'attempt': attempt,
        }
        STATUS.write_text(json.dumps(status_obj, ensure_ascii=False, indent=2), encoding='utf-8')
        print(json.dumps(status_obj, ensure_ascii=False, indent=2))
        if last_status == 'success' and files:
            sys.exit(0)
        if last_status in {'fail', 'failed'}:
            sys.exit(2)
        if time.time() - start >= max_seconds:
            sys.exit(124)
        time.sleep(interval)