#!/usr/bin/env python3
"""Single-API consistent pipeline driver (full workflow, no historical cpp reuse).

Steps: 03/04/05/06/07 yaml -> tiling coverage -> fuzz config -> normal cpp
       -> compile -> normal precheck(1000, >=95%) -> deterministic check
       -> consistent cpp -> verify V1-V10 -> compile(TARGET=consistent)
       -> consistent precheck -> full consistent run -> parse report.
Status written into output/pipeline_status.json.
"""
import argparse
import json
import os
import re
import subprocess
import sys
import time

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(BASE, 'output')
FRAME = os.path.join(BASE, 'turing_sec', 'aclnnAPI_fuzz')
PLOG_DIR = os.environ.get('ASCEND_LOG_PATH', '/root/ascend/log/debug/plog')
STATUS_FILE = os.path.join(OUT, 'pipeline_status.json')

ENV = dict(os.environ)
ENV['ASCEND_HOME_PATH'] = ENV.get('ASCEND_HOME_PATH', '/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux')
ENV['ASCEND_OPP_PATH'] = ENV.get('ASCEND_OPP_PATH', '/usr/local/Ascend/ascend-toolkit/latest/opp')
ENV['LD_LIBRARY_PATH'] = ENV['ASCEND_HOME_PATH'] + '/lib64:' + ENV.get('LD_LIBRARY_PATH', '')


def run(cmd, timeout=600, cwd=BASE, log_file=None):
    p = subprocess.run(cmd, shell=True, cwd=cwd, capture_output=True, text=True, timeout=timeout, env=ENV)
    if log_file:
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f'\n===== CMD: {cmd}\n{p.stdout}\n{p.stderr}\n')
    return p.returncode, p.stdout + p.stderr


def load_status():
    if os.path.exists(STATUS_FILE):
        with open(STATUS_FILE, encoding='utf-8') as f:
            return json.load(f)
    return {}


def save_status(st):
    with open(STATUS_FILE, 'w', encoding='utf-8') as f:
        json.dump(st, f, ensure_ascii=False, indent=1)


def count_success(log_path):
    """统计 FUZZ_RESULT success/faild 数量"""
    s, d = 0, 0
    if not os.path.exists(log_path):
        return 0, 0
    with open(log_path, encoding='utf-8', errors='ignore') as f:
        for line in f:
            if '[FUZZ_RESULT]' in line:
                if ' success' in line:
                    s += 1
                elif ' faild' in line:
                    d += 1
    return s, d


def consistent_stats(log_path):
    s, m, e = 0, 0, 0
    if not os.path.exists(log_path):
        return 0, 0, 0
    with open(log_path, encoding='utf-8', errors='ignore') as f:
        for line in f:
            if '[FUZZ_CONSISTENT]' in line:
                if 'type=consistent' in line and 'success' in line:
                    s += 1
                elif 'hash_mismatch' in line:
                    m += 1
                elif 'execute_error' in line:
                    e += 1
    return s, m, e


def clear_plog():
    subprocess.run(f'rm -f {PLOG_DIR}/plog-*.log', shell=True)


def pipeline(api, status):
    t0 = time.time()
    rec = {'status': 'RUNNING', 'start': time.strftime('%F %T')}
    status[api] = rec
    save_status(status)

    tdir = os.path.join(OUT, api, 'test-design')
    fdir = os.path.join(OUT, api, 'fuzz_logs')
    os.makedirs(tdir, exist_ok=True)
    os.makedirs(fdir, exist_ok=True)
    drv_log = os.path.join(fdir, 'pipeline_driver.log')

    def finish(st, detail=''):
        rec['status'] = st
        rec['detail'] = detail
        rec['end'] = time.strftime('%F %T')
        rec['minutes'] = round((time.time() - t0) / 60, 1)
        save_status(status)
        print(f'[{api}] {st} {detail} ({rec["minutes"]}min)')

    # Step 1: 03 参数定义
    rc, out = run(f'python scripts/auto_gen_param_yaml.py {api}', log_file=drv_log)
    if rc != 0:
        return finish('NO_PARAM_YAML', out.strip().splitlines()[-1] if out.strip() else '')

    # Step 2: 05 约束定义
    rc, out = run(f'python scripts/auto_gen_constraint_yaml.py --param_yaml {tdir}/03_参数定义.yaml --output {tdir}/05_约束定义.yaml', log_file=drv_log)
    if rc != 0:
        return finish('NO_CONSTRAINT_YAML', '')

    # Step 3-6: 04/06/07
    run(f'python skills/ascendc-st-design/scripts/generate_test_factors.py {tdir}/03_参数定义.yaml {tdir}/04_测试因子.yaml', log_file=drv_log)
    run(f'python skills/ascendc-st-design/scripts/generate_implicit_constraints.py {tdir}/04_测试因子.yaml {tdir}/05_约束定义.yaml', log_file=drv_log)
    run(f'python skills/ascendc-st-design/scripts/generate_solver_config.py {tdir}/05_约束定义.yaml {tdir}/06_求解配置.yaml', log_file=drv_log)
    run(f'python skills/ascendc-st-design/scripts/generate_factor_values.py {tdir}/06_求解配置.yaml {tdir}/05_约束定义.yaml {tdir}/04_测试因子.yaml {tdir}/07_因子值.xlsx --max-cases 1000', log_file=drv_log)

    # Step 7: tiling 覆盖（失败允许, N=0）
    rc, out = run(f'python skills/ascendc-fuzz-cpp-gen/scripts/gen_tiling_key_coverage.py --op_name {api} --param_yaml {tdir}/03_参数定义.yaml --output {tdir}/08_tiling_key_coverage.yaml', log_file=drv_log)
    n_tk = 0
    m = re.search(r'api_reachable=(\d+)', out)
    if m:
        n_tk = int(m.group(1))
    rec['n_tk'] = n_tk

    # Step 8: fuzz config（md 路径从索引取）
    idx = json.load(open('/tmp/opencode/op_index.json')) if os.path.exists('/tmp/opencode/op_index.json') else {}
    md = idx.get(api, {}).get('md', '')
    rc, out = run(f'python skills/ascendc-fuzz-cpp-gen/scripts/st_to_fuzz_config.py --param_def {tdir}/03_参数定义.yaml --constraint_def {tdir}/05_约束定义.yaml --api_doc "{md}" --op_name {api} --output {tdir}/fuzz_op_config.yaml', log_file=drv_log)
    if rc != 0 or not os.path.exists(f'{tdir}/fuzz_op_config.yaml'):
        return finish('CONFIG_FAIL', '')

    # Step 9: normal CPP
    tiling_arg = f'--tiling_yaml {tdir}/08_tiling_key_coverage.yaml' if n_tk > 0 else ''
    rc, out = run(f'python skills/ascendc-fuzz-cpp-gen/scripts/gen_fuzz_cpp.py --op_name {api} --st_config {tdir}/fuzz_op_config.yaml {tiling_arg} --output {OUT}/{api}/', log_file=drv_log)
    normal_cpp = os.path.join(OUT, api, f'{api}.cpp')
    if rc != 0 or not os.path.exists(normal_cpp):
        return finish('GEN_CPP_FAIL', out.strip().splitlines()[-1] if out.strip() else '')

    # Step 10: 编译 normal
    run('rm -f ' + os.path.join(FRAME, 'src/op/*.cpp'))
    run(f'cp {normal_cpp} {FRAME}/src/op/')
    rc, out = run('make op_fuzz', cwd=FRAME, timeout=600, log_file=drv_log)
    if rc != 0:
        return finish('COMPILE_FAIL', '')

    # Step 11: normal 预检 1000 条（最多3轮, malloc失败多时缩减size）
    precheck_log = os.path.join(fdir, f'{api}_fuzz_precheck.log')
    final_succ = 0
    for attempt in range(3):
        clear_plog()
        with open(precheck_log, 'w', encoding='utf-8') as lf:
            prc = subprocess.run(f'./aclnnAPI_fuzz {api}_fuzz 1000', shell=True, cwd=os.path.join(FRAME, 'out'), stdout=lf, stderr=subprocess.STDOUT, timeout=900, env=ENV).returncode
        succ, fail = count_success(precheck_log)
        final_succ = succ
        rate = succ / 1000.0
        rec['precheck'] = f'{succ}/1000={rate:.1%}'
        save_status(status)
        if rate >= 0.95:
            break
        # malloc 失败主导时缩减 sizeTargets 重试
        n_malloc = 0
        try:
            n_malloc = open(precheck_log, encoding='utf-8', errors='ignore').read().count('aclrtMalloc failed')
        except Exception:
            pass
        if n_malloc >= 20 and attempt < 2:
            cpp = open(normal_cpp, encoding='utf-8').read()
            cpp2 = re.sub(r'const long long sizeTargets\[4\] = \{(\d+)LL, (\d+)LL, (\d+)LL, (\d+)LL\};',
                          lambda mm: 'const long long sizeTargets[4] = {%dLL, %dLL, %dLL, %dLL};' % tuple(max(1000, int(g) // 4) for g in mm.groups()),
                          cpp)
            if cpp2 != cpp:
                open(normal_cpp, 'w', encoding='utf-8').write(cpp2)
                run('rm -f ' + os.path.join(FRAME, 'src/op/*.cpp'))
                run(f'cp {normal_cpp} {FRAME}/src/op/')
                run('make op_fuzz', cwd=FRAME, timeout=600, log_file=drv_log)
            continue
        break
    if final_succ < 950:
        return finish('PRECHECK_FAIL', f'succ={final_succ}/1000')

    # Step 12: 确定性判定
    rc, out = run(f'python skills/gen-consistent-from-normal/scripts/parse_md_deterministic.py --api {api} --output {tdir}/deterministic_check.json', log_file=drv_log)
    need_det = 0
    try:
        dj = json.load(open(f'{tdir}/deterministic_check.json', encoding='utf-8'))
        need_det = 1 if dj.get('need_deterministic') else 0
    except Exception:
        pass
    rec['need_deterministic'] = need_det

    # Step 13: consistent CPP
    rc, out = run(f'python skills/gen-consistent-from-normal/scripts/gen_consistent_from_normal.py --input {normal_cpp} --output {OUT}/{api}/{api}_consistent.cpp --api {api} --deterministic_json {tdir}/deterministic_check.json', log_file=drv_log)
    cons_cpp = os.path.join(OUT, api, f'{api}_consistent.cpp')
    if rc != 0 or not os.path.exists(cons_cpp):
        return finish('GEN_CONSISTENT_FAIL', '')

    # Step 14: verify V1-V10
    rc, out = run(f'python skills/gen-consistent-from-normal/scripts/verify_consistent.py --cpp {cons_cpp} --api {api} --need_deterministic {need_det}', log_file=drv_log)
    if 'RESULT: PASS' not in out:
        return finish('VERIFY_FAIL', out.strip().splitlines()[-1] if out.strip() else '')

    # Step 15: 编译 consistent（清空 src_consistent 只放当前）
    run('rm -f ' + os.path.join(FRAME, 'src/src_consistent/*.cpp'))
    run(f'cp {cons_cpp} {FRAME}/src/src_consistent/')
    rc, out = run('make op_fuzz TARGET=consistent', cwd=FRAME, timeout=600, log_file=drv_log)
    if rc != 0:
        return finish('COMPILE_CONSISTENT_FAIL', '')

    # Step 16: consistent 预调试（N×2 或 100 基线）
    pre_n = max(n_tk * 2, 100) if n_tk > 0 else 100
    clear_plog()
    pre_log = os.path.join(fdir, f'{api}_consistent_precheck.log')
    try:
        with open(pre_log, 'w', encoding='utf-8') as lf:
            prc = subprocess.run(f'./aclnnAPI_fuzz_consistent {api}_consistent_fuzz {pre_n}', shell=True, cwd=os.path.join(FRAME, 'out'), stdout=lf, stderr=subprocess.STDOUT, timeout=1200, env=ENV).returncode
        cs, cm, ce = consistent_stats(pre_log)
    except subprocess.TimeoutExpired:
        return finish('CONSISTENT_PRECHECK_TIMEOUT', '')
    rec['consistent_precheck'] = f'consistent={cs} mismatch={cm} err={ce}/{pre_n}'
    save_status(status)
    if cs == 0 and cm == 0:
        return finish('CONSISTENT_PRECHECK_CRASH', f'rc={prc}')

    # Step 17: 全量压测（N×50 封顶 10000, 无 tiling 基线 1000）
    full_n = min(n_tk * 50, 10000) if n_tk > 0 else 1000
    clear_plog()
    full_log = os.path.join(fdir, f'{api}_consistent_fuzz_output.log')
    with open(full_log, 'w', encoding='utf-8') as lf:
        p = subprocess.run(f'./aclnnAPI_fuzz_consistent {api}_consistent_fuzz {full_n}', shell=True, cwd=os.path.join(FRAME, 'out'), stdout=lf, stderr=subprocess.STDOUT, timeout=1800, env=ENV)
    if p.returncode != 0:
        rec['full_run_rc'] = p.returncode
    # 回传 plog
    run(f'cp {PLOG_DIR}/plog-*.log {fdir}/{api}_consistent_fuzz_plog.log 2>/dev/null', log_file=drv_log)

    # Step 18: 解析报告
    rc, out = run(f'python skills/gen-consistent-from-normal/scripts/parse_consistent_log.py --log {full_log} --api {api} --plog {fdir}/{api}_consistent_fuzz_plog.log --run_num 2', log_file=drv_log)
    m = re.search(r'total=(\d+) tks=\d+ consistent=(\d+) hash_mismatch=(\d+) execute_error=(\d+) exception=(\d+) aicore_error=(\w+) need_issue=(\w+)', out)
    if m:
        rec.update({
            'total': int(m.group(1)), 'consistent': int(m.group(2)),
            'hash_mismatch': int(m.group(3)), 'execute_error': int(m.group(4)),
            'aicore_error': m.group(6), 'need_issue': m.group(7),
        })
        finish('OK' if int(m.group(3)) == 0 and m.group(6) == 'False' else 'OK_WITH_ISSUE',
               f'mismatch={m.group(3)} aicore={m.group(6)}')
    else:
        finish('PARSE_FAIL', out.strip().splitlines()[-1] if out.strip() else '')


def count_success_from_run(out_text):
    s, d = 0, 0
    for line in out_text.splitlines():
        if '[FUZZ_RESULT]' in line:
            if ' success' in line:
                s += 1
            elif ' faild' in line:
                d += 1
    return s, d


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('api')
    args = ap.parse_args()
    status = load_status()
    pipeline(args.api, status)


if __name__ == '__main__':
    main()