#!/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

BASE = Path('/Users/bot1/Volumes/root_for_ai/AI工作区/通用_产品宣传视频_古钱币杜邦纸钱袋包_20260530_1702/outputs/dreamina-seedance2-four-segments-20260530')
BIN = '/Users/bot1/.hermes/profiles/video/home/.local/bin/dreamina'
SEGMENTS = [
    (1, 'cf32490b-2902-45c3-a995-d73a2b06ec8e'),
    (2, '473e9d3a-004d-41ce-af70-fd2b86cff779'),
    (3, '80113c72-683d-40ea-b28b-c73016d6f4b5'),
    (4, '71d46c4a-16ae-4fc4-9c89-bcc6e2ad190a'),
]

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

def status_from(text):
    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)
    low = text.lower()
    if 'success' in low:
        return 'success'
    if 'fail' in low:
        return 'fail'
    if 'querying' in low or 'generating' in low:
        return 'querying'
    return None

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

max_seconds = int(sys.argv[1]) if len(sys.argv) > 1 else 900
interval = int(sys.argv[2]) if len(sys.argv) > 2 else 20
start = time.time()
done = {n: False for n, _ in SEGMENTS}
last = {n: None for n, _ in SEGMENTS}

while True:
    summary = []
    for n, submit_id in SEGMENTS:
        seg_dir = BASE / f'segment_{n:02d}'
        download_dir = seg_dir / 'downloads'
        download_dir.mkdir(parents=True, exist_ok=True)
        if done[n]:
            files = media_files(download_dir)
            summary.append({'segment': n, 'submit_id': submit_id, 'gen_status': 'success', 'media_files': [str(p) for p in files]})
            continue
        res = run_query(submit_id, download_dir)
        out = res.stdout
        with (seg_dir / 'query_poll.log').open('a', encoding='utf-8') as f:
            f.write(f"\n--- {datetime.now().astimezone().isoformat()} rc={res.returncode} ---\n{out}\n")
        st = status_from(out) or last[n]
        last[n] = st
        files = media_files(download_dir)
        obj = {
            'updated_at': datetime.now().astimezone().isoformat(),
            'segment': n,
            'submit_id': submit_id,
            'returncode': res.returncode,
            'gen_status': st,
            'media_files': [str(p) for p in files],
        }
        (seg_dir / 'query_status.json').write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding='utf-8')
        summary.append(obj)
        if st == 'success' and files:
            done[n] = True
        elif st in {'fail', 'failed'}:
            (BASE / 'poll_summary.json').write_text(json.dumps({'updated_at': datetime.now().astimezone().isoformat(), 'summary': summary}, ensure_ascii=False, indent=2), encoding='utf-8')
            print(json.dumps(summary, ensure_ascii=False, indent=2))
            sys.exit(2)
    (BASE / 'poll_summary.json').write_text(json.dumps({'updated_at': datetime.now().astimezone().isoformat(), 'summary': summary}, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    if all(done.values()):
        sys.exit(0)
    if time.time() - start >= max_seconds:
        sys.exit(124)
    time.sleep(interval)