import os
import time
import subprocess
import sys
import argparse
import time
from openai import OpenAI

def fetch_generator(statement, model, maxtokens):
    APIKEY = os.environ.get('DEEPSEEK_API_KEY')
    if not APIKEY:
        raise Exception('需要设置 DEEPSEEK_API_KEY 环境变量')

    client = OpenAI(
        api_key=APIKEY,
        base_url='https://api.deepseek.com')

    start_time = time.time()
    msgs = [{ 'role': 'user', 'content': 
             """
请根据以下信息学竞赛题目生成数据生成器。只输出 gen.py 代码。
要求：
1. 解析题目中的输入格式、数据范围和子任务划分
2. 创建 data/ 目录并生成所有 .in 文件
3. 数据严格符合题目约束
4. 使用合理的数据梯度，包含边界情况
5. 数据文件名一定为 'data{编号}.in'，编号不含前导零
6. 没有子任务的，数据文件直接放在 data/ 下；有子任务的，为每个子任务创建 subtask{编号}/ 目录，把数据文件移动到对应的子任务目录中
7. 数据点在 10 到 20 个内
8. 务必确保你给出的数据是正确的！
输出只包含可以直接运行的 gen.py 的代码，无其他内容。
题目描述：""" + statement}]
    resp = client.chat.completions.create(
            model=model,
            messages=msgs,
            stream=True,
            temperature=0.0,
            max_tokens=maxtokens)

    response = ''
    for c in resp:
        d = c.choices[0].delta
        if hasattr(d, 'reasoning_content') and d.reasoning_content is not None:
            delta = d.reasoning_content
        elif hasattr(d, 'content') and d.content is not None:
            delta = d.content
            if response == '':
                delta = '\nThinking done\n' + delta
            response += d.content
        else:
            continue
        
        with open('stream', 'a') as f:
            f.write(delta)
        if hasattr(c, 'usage') and c.usage:
            tokens = c.usage

    if response.startswith('```python'):
        response = response[9:-3]
    return (response, tokens, time.time() - start_time)

def execute(s):
    print('正在执行命令：', s)

    try:
        subprocess.run(['sh', '-c', s], check=True, timeout=20)
    except subprocess.CalledProcessError as e:
        raise Exception('命令执行错误')
    except subprocess.TimeoutExpired:
        raise Exception('命令执行超时（20s）')

def generate_in():
    # 其实应该用个沙盒比较安全，但是我懒得写了。可以最后扔到容器里头跑。
    execute('python gen.py')

def flatten_data():
    execute('find -type f | xargs -I {} mv -n {} . || true')
    execute('rmdir subtask* || true')

def run_std():
    if not os.path.exists('../std.cpp'):
        raise Exception('未找到 std.cpp')
    execute('g++ ../std.cpp -o std -Wall -Wextra -O2 -std=c++14')

    for i in os.listdir():
        if not i.endswith('.in'):
            print(f'警告：未知文件 {i}')
            continue
        out = i[:-2] + 'ans'
        execute(f'./std < {i} > {out}')
    execute('rm std')

def gen_problem_conf(time, mem):
    execute(f'python ../autoconf.py {time} {mem}')

def compress_data():
    execute('zip ../data.zip -r .')

def cleanup():
    os.chdir('..')
    execute('rm -r data/')

def pretty_size(d):
    if d < 1024:
        return f'{d}'
    if d < 1048576:
        return f'{d/1024:.2f} KB'
    return f'{d/1048576:.2f} MB'

if __name__ == '__main__':
    start_time = time.time()

    parser = argparse.ArgumentParser()
    parser.add_argument('statement', help='题面文件')
    parser.add_argument('--model', type=str, default='deepseek-v4-pro', help='使用的模型（默认为 deepseek-v4-pro，可以改为 deepseek-v4-flash）')
    parser.add_argument('--maxtokens', type=int, default=20000, help='token 上限（默认为 20000）')
    parser.add_argument('--time', type=int, default=1, help='时间限制（秒，默认为 1）')
    parser.add_argument('--mem', type=int, default=256, help='空间限制（MB，默认为 256）')
    parser.add_argument('--offline', help='不生成新的 gen.py', action='store_true')
    args = parser.parse_args()

    if os.path.exists('data'):
        print('正在清空已有的 data 目录')
        execute('rm -r data')
    if os.path.exists('data.zip'):
        print('删除已存在的 data.zip')
        execute('rm data.zip')

    tokens = 0

    with open(args.statement) as f:
        if not args.offline:
            print(f'正在请求 {args.model}，token 上限为 {args.maxtokens} 个')
            resp, tokens, duration = fetch_generator(f.read(), args.model, args.maxtokens)
            with open('gen.py', 'w') as fw:
                fw.write(resp)
            print(f'--- tokens: {tokens}, time cost: {duration:.6f} seconds')
        else:
            if not os.path.exists('gen.py'):
                raise Exception('离线模式下必须提供已有的 gen.py')
            print('（离线模式）')

        generate_in()
        os.chdir('data')
        gen_problem_conf(args.time, args.mem)

        # run_std 需要一个扁平的目录，也就是所有文件放在 data/ 下，而 autoconf.py 不要求这么做
        flatten_data()
        run_std()
        compress_data()
        cleanup()
        prettysize = pretty_size(os.path.getsize('data.zip'))
        print(f'已生成 data.zip（{prettysize}），耗时 {(time.time() - start_time):.2f} 秒，消耗 {tokens} 个 token')