import gradio as gr
import asyncio
import sys
import json
import random
from pathlib import Path
import zipfile
from datetime import datetime

TOKEN_DIR = Path("/app/tokens")
TOKEN_DIR.mkdir(exist_ok=True)
ACCOUNTS_FILE = TOKEN_DIR / "accounts.txt"

# 全局状态
current_process = None

async def run_registration(proxy: str, loop_mode: bool = False, progress=gr.Progress()):
    """
    执行注册，loop_mode=False 表示单次（带 --once），loop_mode=True 表示无限循环（不带 --once）
    返回实时日志的异步生成器
    """
    global current_process
    cmd = [sys.executable, "/app/register.py"]
    if not loop_mode:
        cmd.append("--once")
    if proxy and proxy.strip():
        cmd.extend(["--proxy", proxy.strip()])

    process = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    current_process = process

    lines = []
    async for line in process.stdout:
        decoded = line.decode('utf-8', errors='replace').rstrip()
        lines.append(decoded)
        if len(lines) % 10 == 0:
            progress(0.5, desc="运行中...")
        yield "\n".join(lines)

    await process.wait()
    current_process = None
    progress(1.0, desc="进程结束")
    yield "\n".join(lines) + "\n\n✅ 进程已结束。"

def stop_registration():
    global current_process
    if current_process and current_process.returncode is None:
        current_process.terminate()
        return "⏹️ 已发送停止信号"
    return "没有正在运行的进程"

def list_token_files():
    files = list(TOKEN_DIR.glob("token_*.json"))
    return [str(f) for f in files]

def get_accounts_content():
    if ACCOUNTS_FILE.exists():
        return ACCOUNTS_FILE.read_text(encoding='utf-8')
    return "暂无账号信息"

def download_accounts():
    if ACCOUNTS_FILE.exists():
        return str(ACCOUNTS_FILE)
    return None

def download_all_files():
    token_files = list(TOKEN_DIR.glob("token_*.json"))
    if not token_files and not ACCOUNTS_FILE.exists():
        return None

    zip_path = TOKEN_DIR / f"all_files_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for f in token_files:
            zipf.write(f, arcname=f.name)
        if ACCOUNTS_FILE.exists():
            zipf.write(ACCOUNTS_FILE, arcname="accounts.txt")
    return str(zip_path)

def download_latest():
    files = list(TOKEN_DIR.glob("token_*.json"))
    if files:
        latest = max(files, key=lambda p: p.stat().st_mtime)
        return latest
    return None

def update_file_list():
    files = list_token_files()
    count = len(files)
    return gr.update(value=files), gr.update(value=count)

# 构建 Gradio 界面
with gr.Blocks(title="OpenAI 自动注册 (双模式版)") as demo:
    gr.Markdown("""
    # 🤖 OpenAI 自动注册 (双模式版)
    - **单次注册**：执行一次注册，完成后停止。
    - **无限循环**：启动后 `register.py` 自己无限循环注册，可随时停止。
    - 所有生成的 token 文件及 `accounts.txt` 均可单独下载或打包下载全部。
    """)

    with gr.Row():
        proxy_input = gr.Textbox(
            label="代理地址 (可选)",
            placeholder="例如 http://127.0.0.1:7890"
        )

    with gr.Row():
        single_btn = gr.Button("▶️ 单次注册", variant="primary")
        loop_btn = gr.Button("🔄 无限循环", variant="primary")
        stop_btn = gr.Button("⏹️ 停止", variant="secondary")

    output = gr.Textbox(label="实时日志", lines=20, interactive=False)
    progress_bar = gr.Progress()

    gr.Markdown("---\n### 📁 生成的 Token 文件")
    with gr.Row():
        file_count = gr.Number(value=0, label="当前 Token 数量", interactive=False)
        file_list = gr.Files(label="所有 Token 文件", file_count="multiple")
    with gr.Row():
        refresh_btn = gr.Button("🔄 刷新文件列表")
        download_latest_btn = gr.Button("📥 下载最新 Token")
        download_all_btn = gr.Button("📦 下载全部文件")

    gr.Markdown("---\n### 📋 账号信息 (accounts.txt)")
    with gr.Row():
        accounts_display = gr.Textbox(label="accounts.txt 内容", lines=10, interactive=False)
        refresh_accounts_btn = gr.Button("🔄 刷新账号列表")
        download_accounts_btn = gr.Button("📥 下载 accounts.txt")

    # 单次注册
    async def single_run(proxy):
        yield gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True), None
        logs = ""
        async for log in run_registration(proxy, loop_mode=False):
            logs = log
            yield gr.update(), gr.update(), gr.update(), logs
        yield gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=False), logs

    single_btn.click(
        fn=single_run,
        inputs=proxy_input,
        outputs=[single_btn, loop_btn, stop_btn, output],
        queue=True
    ).then(
        fn=update_file_list,
        outputs=[file_list, file_count]
    ).then(
        fn=get_accounts_content,
        outputs=accounts_display
    )

    # 无限循环
    async def loop_run(proxy):
        yield gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=True), None
        logs = ""
        async for log in run_registration(proxy, loop_mode=True):
            logs = log
            yield gr.update(), gr.update(), gr.update(), logs
        yield gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=False), logs

    loop_btn.click(
        fn=loop_run,
        inputs=proxy_input,
        outputs=[single_btn, loop_btn, stop_btn, output],
        queue=True
    ).then(
        fn=update_file_list,
        outputs=[file_list, file_count]
    ).then(
        fn=get_accounts_content,
        outputs=accounts_display
    )

    # 停止
    def stop_fn():
        msg = stop_registration()
        return gr.update(interactive=False), msg

    stop_btn.click(
        fn=stop_fn,
        outputs=[stop_btn, output]
    ).then(
        fn=update_file_list,
        outputs=[file_list, file_count]
    ).then(
        fn=get_accounts_content,
        outputs=accounts_display
    )

    # 刷新文件列表
    refresh_btn.click(
        fn=update_file_list,
        outputs=[file_list, file_count]
    )

    download_latest_btn.click(
        fn=download_latest,
        outputs=gr.File(label="下载最新 Token")
    )

    download_all_btn.click(
        fn=download_all_files,
        outputs=gr.File(label="下载全部文件")
    )

    # 账号相关
    refresh_accounts_btn.click(
        fn=get_accounts_content,
        outputs=accounts_display
    )

    download_accounts_btn.click(
        fn=download_accounts,
        outputs=gr.File(label="下载 accounts.txt")
    )

demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860)