"""Attach modification notices and stage validated platform release directories."""
import argparse
import json
import os
import shutil
import struct
from pathlib import Path

from scripts.integrity import sha256, verify_files
from scripts.provenance import model_identity, validate_model_structure

NOTICE = ('Modified by ixim / iximbox: eligible linear weights converted from '
          'Qwen-Image-2.1 to bitsandbytes LLM.int8 INT8. Built with Qwen. '
          'Non-commercial research/evaluation under the accompanying Qwen Research License.')
REVISION = 'b3179ad355be050328e483a9dfdd9e60cd62adfa'


def add_safetensors_notice(path):
    path = Path(path)
    temp = path.with_name(path.name + '.notice-tmp')
    try:
        with path.open('rb') as source:
            prefix = source.read(8)
            if len(prefix) != 8:
                raise ValueError(f'Invalid safetensors prefix: {path}')
            length = struct.unpack('<Q', prefix)[0]
            if length > 100_000_000:
                raise ValueError('Unexpected safetensors header size')
            header = json.loads(source.read(length))
            metadata = header.setdefault('__metadata__', {})
            if metadata.get('modification_notice') == NOTICE:
                return
            metadata.update(modification_notice=NOTICE, base_revision=REVISION)
            encoded = json.dumps(header, ensure_ascii=False, separators=(',', ':')).encode('utf-8')
            encoded += b' ' * ((8 - len(encoded) % 8) % 8)
            with temp.open('wb') as target:
                target.write(struct.pack('<Q', len(encoded)))
                target.write(encoded)
                shutil.copyfileobj(source, target, length=8 * 1024 * 1024)
        os.replace(temp, path)
    finally:
        if temp.exists():
            temp.unlink()


def annotate(root):
    root = Path(root)
    if validate_model_structure(root) != 'int8':
        raise ValueError('Only complete INT8 pipelines can be annotated')
    for component in ('transformer', 'text_encoder'):
        for path in (root / component).glob('*.safetensors'):
            add_safetensors_notice(path)
        for path in (root / component).glob('*.json'):
            data = json.loads(path.read_text(encoding='utf-8'))
            if 'weight_map' in data:
                data.setdefault('metadata', {})['modification_notice'] = NOTICE
            else:
                data['_modification_notice'] = NOTICE
            path.write_text(json.dumps(data, indent=2) + '\n', encoding='utf-8')
    (root / 'CHANGES.md').write_text(
        '# Modifications\n\n' + NOTICE + '\n\n'
        '- Converted eligible transformer and text encoder linear layers to LLM.int8, threshold 6.0.\n'
        '- Retained sensitive projections, normalization, embeddings, vision model and VAE in floating point.\n'
        '- Re-serialized component weights, shard indexes and configs; no fine-tuning.\n'
        '- Weight-file headers and modified JSON files contain modification notices.\n'
        '- Exact quantized module names and dtype counts are in the component reports.\n', encoding='utf-8')
    source = {'base_model': 'Qwen/Qwen-Image-2.1', 'base_revision': REVISION,
              'diffusers_commit': '80c7ed262aeffbeb43ef13ae04baeb9b84515a69',
              'method': 'bitsandbytes LLM.int8', 'threshold': 6.0,
              'weight_files': [{'path': p.relative_to(root).as_posix(),
                               'size': p.stat().st_size, 'sha256': sha256(p)}
                              for p in sorted(root.rglob('*.safetensors'))]}
    (root / 'conversion.json').write_text(json.dumps(source, indent=2), encoding='utf-8')


def evaluation_text(summary, language):
    s = summary
    if language == 'en':
        return (f'{s["pairs"]} paired outputs across {s["cases"]} cases at '
                f'{s["width"]}×{s["height"]} on an RTX 5090.\n\n'
                '| Metric | BF16 | INT8 |\n|---|---:|---:|\n'
                f'| Weight files (decimal GB) | {s["bf16_weight_bytes"]/1e9:.3f} | {s["int8_weight_bytes"]/1e9:.3f} |\n'
                f'| Mean call latency (s) | {s["bf16_mean_seconds"]:.2f} | {s["int8_mean_seconds"]:.2f} |\n'
                f'| Maximum CUDA allocated memory (GiB) | {s["bf16_max_allocated_gib"]:.2f} | '
                f'{s["int8_max_allocated_gib"]:.2f} |')
    return (f'在 RTX 5090 上完成 {s["cases"]} 类用例、{s["pairs"]} 对输出，'
            f'分辨率为 {s["width"]}×{s["height"]}。\n\n'
            '| 指标 | BF16 | INT8 |\n|---|---:|---:|\n'
            f'| 权重体积（十进制 GB） | {s["bf16_weight_bytes"]/1e9:.3f} | {s["int8_weight_bytes"]/1e9:.3f} |\n'
            f'| 平均调用耗时（秒） | {s["bf16_mean_seconds"]:.2f} | {s["int8_mean_seconds"]:.2f} |\n'
            f'| CUDA 已分配显存最高值（GiB） | {s["bf16_max_allocated_gib"]:.2f} | '
            f'{s["int8_max_allocated_gib"]:.2f} |')


def stage(model, evaluation, output, platform):
    model, evaluation, output = Path(model), Path(evaluation), Path(output)
    if output.exists():
        raise FileExistsError(f'Refusing to overwrite release directory: {output}')
    summary = json.loads((evaluation / 'summary.json').read_text())
    if summary['pairs'] < 14 or summary['cases'] < 7 or summary['steps'] != 40 or not summary['warmup']:
        raise ValueError('Release requires the complete paired 40-step evaluation with warmup')
    validate_evaluation(evaluation, model, summary)
    conversion = json.loads((model / 'conversion.json').read_text())
    verify_files(model, conversion['weight_files'])
    for component in ('transformer', 'text_encoder'):
        q = json.loads((model / component / 'config.json').read_text())['quantization_config']
        if not q.get('load_in_8bit'):
            raise ValueError(f'{component} is not saved as INT8')
    shutil.copytree(model, output, ignore=shutil.ignore_patterns('.cache', '__pycache__', '*.lock'))
    shutil.copytree(evaluation, output / 'evaluation')
    for name in ('scripts', 'benchmarks', 'cards', 'tests'):
        shutil.copytree(name, output / name, ignore=shutil.ignore_patterns('__pycache__'))
    shutil.copy2('README.md', output / 'REPRODUCE.md')
    for name in ('requirements.txt', 'PUBLISHING.md'):
        shutil.copy2(name, output / name)
    shutil.copy2('artifacts/download-verification.json', output / 'upstream-verification.json')
    card = Path(f'cards/{platform}.md').read_text(encoding='utf-8')
    card = card.replace('{{EVALUATION_EN}}', evaluation_text(summary, 'en'))
    card = card.replace('{{EVALUATION_ZH}}', evaluation_text(summary, 'zh'))
    from scripts.model_card import render_details, render_samples
    details = render_details(evaluation, platform)
    card = card.replace('{{EVALUATION_DETAILS_EN}}', details)
    card = card.replace('{{EVALUATION_DETAILS_ZH}}', details)
    samples = render_samples(evaluation, platform)
    card = card.replace('{{SAMPLES_EN}}', samples).replace('{{SAMPLES_ZH}}', samples)
    from scripts.audit_report import render_audit
    audit = render_audit(evaluation, platform)
    card = card.replace('{{AUDIT_EN}}', audit).replace('{{AUDIT_ZH}}', audit)
    if '{{' in card:
        raise ValueError('Unrendered model card template')
    (output / 'README.md').write_text(card, encoding='utf-8')
    if platform == 'modelscope':
        (output / 'configuration.json').write_text(json.dumps({'framework': 'pytorch',
                    'task': 'text-to-image-synthesis'}, indent=2), encoding='utf-8')
    manifest = [{'path': p.relative_to(output).as_posix(), 'size': p.stat().st_size, 'sha256': sha256(p)}
                for p in sorted(output.rglob('*')) if p.is_file()]
    (output / 'MANIFEST.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
    print(f'Staged {platform}: {len(manifest)} files, {sum(x["size"] for x in manifest)/1e9:.3f} GB', flush=True)


def validate_release(root):
    root = Path(root)
    rows = json.loads((root / 'MANIFEST.json').read_text())
    expected = {r['path'] for r in rows} | {'MANIFEST.json'}
    actual = {p.relative_to(root).as_posix() for p in root.rglob('*') if p.is_file()}
    if actual != expected:
        raise ValueError(f'Unexpected/missing release files: {actual ^ expected}')
    for name in expected:
        if any(part.startswith('.') for part in Path(name).parts) or name.endswith(('.log', '.incomplete')):
            raise ValueError(f'Non-release file: {name}')
    verify_files(root, rows)
    summary = json.loads((root / 'evaluation/summary.json').read_text())
    validate_evaluation(root / 'evaluation', root, summary)
    if (root / 'evaluation/editing').exists():
        from scripts.editing_validation import validate_editing
        validate_editing(root / 'evaluation')
    if (root / 'evaluation/editing-v2').exists():
        from scripts.editing_report import validate_suite
        runs = validate_suite(root / 'evaluation/editing-v2', root / 'scripts')
        for precision in ('bf16', 'int8'):
            core = json.loads((root / f'evaluation/{precision}/environment.json').read_text())
            if runs[precision][0]['model_identity']['fingerprint'] != core['model_identity']['fingerprint']:
                raise ValueError('Editing v2 model differs from retained main suite')
    return rows


def validate_evaluation(evaluation, model, summary, case_names=None, identity=None):
    from scripts.report import load_records
    from scripts.integrity import validate_pair
    from scripts.provenance import validate_roles
    import statistics
    evaluation = Path(evaluation)
    baseline = load_records(evaluation / 'bf16')
    candidate = load_records(evaluation / 'int8')
    is_core = case_names is None
    if is_core:
        case_names = ('portrait', 'english_text', 'chinese_text',
                      'composition', 'texture', 'rgba', 'edit')
    expected = {(name, seed) for name in case_names for seed in (42, 123)}
    if baseline.keys() != expected or candidate.keys() != expected:
        raise ValueError('Incomplete benchmark suite')
    ea = json.loads((evaluation / 'bf16/environment.json').read_text())
    eb = json.loads((evaluation / 'int8/environment.json').read_text())
    validate_roles(ea['model_identity'], eb['model_identity'])
    for key in ('offload_aux_fix', 'benchmark_sha256', 'runtime_helper_sha256'):
        if ea.get(key) != eb.get(key):
            raise ValueError(f'Benchmark implementation mismatch: {key}')
    identity = identity or model_identity(model)
    if identity['fingerprint'] != eb['model_identity']['fingerprint']:
        raise ValueError('Evaluated model differs from staged model')
    for key in ('gpu', 'cuda', 'packages', 'offload', 'warmup', 'generator_device', 'cases_sha256'):
        if ea[key] != eb[key]:
            raise ValueError(f'Runtime mismatch: {key}')
    if not ea['warmup'] or ea['offload'] != 'model':
        raise ValueError('Release evaluation requires warmup and model CPU offload')
    for key in expected:
        a, b = baseline[key], candidate[key]
        validate_pair(a, b)
        if (a['steps'], a['width'], a['height'], a['cfg'], a['kv_cache']) != (40, 1024, 1024, 1.0, True):
            raise ValueError(f'Unexpected release evaluation settings: {key}')
        if key[0] == 'edit' and a['input_sha256'] != baseline[('portrait', 42)]['image_sha256']:
            raise ValueError('Editing source image differs from documented input')
    computed = {'pairs': len(expected), 'cases': len(case_names), 'width': 1024, 'height': 1024, 'steps': 40, 'warmup': True,
                'bf16_mean_seconds': statistics.mean(r['seconds'] for r in baseline.values()),
                'int8_mean_seconds': statistics.mean(r['seconds'] for r in candidate.values()),
                'bf16_max_allocated_gib': max(r['peak_allocated_bytes'] / 2**30 for r in baseline.values()),
                'int8_max_allocated_gib': max(r['peak_allocated_bytes'] / 2**30 for r in candidate.values()),
                'bf16_weight_bytes': sum(f['size'] for f in ea['model_identity']['files'] if f['path'].endswith('.safetensors')),
                'int8_weight_bytes': sum(f['size'] for f in eb['model_identity']['files'] if f['path'].endswith('.safetensors'))}
    if computed != summary:
        raise ValueError('Summary does not match validated raw measurements')
    if is_core:
        supplement = evaluation / 'young_woman'
        extra_summary = json.loads((supplement / 'summary.json').read_text())
        validate_evaluation(supplement, model, extra_summary,
                            case_names=('young_chinese_woman',), identity=identity)
        for precision, original in (('bf16', ea), ('int8', eb)):
            extra = json.loads((supplement / precision / 'environment.json').read_text())
            if extra['model_identity']['fingerprint'] != original['model_identity']['fingerprint']:
                raise ValueError('Supplement uses a different model from the core evaluation')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('action', choices=['annotate', 'stage', 'check'])
    ap.add_argument('--model', default='models/int8')
    ap.add_argument('--evaluation', default='artifacts/eval')
    ap.add_argument('--output', default='release/huggingface')
    ap.add_argument('--platform', choices=['huggingface', 'modelscope'], default='huggingface')
    args = ap.parse_args()
    if args.action == 'annotate':
        annotate(args.model)
    elif args.action == 'stage':
        stage(args.model, args.evaluation, args.output, args.platform)
    else:
        print(f'Validated {len(validate_release(args.output))} files')


if __name__ == '__main__':
    main()