import hashlib
import json
from pathlib import Path
import shutil
import urllib.request
from datasets import load_dataset
from common import ROOT, save_json


def dump(path, items):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(''.join(json.dumps(x, ensure_ascii=False) + '\n' for x in items))


def main():
    practice = ROOT / 'data/practice'
    practice.mkdir(exist_ok=True)
    for f in (ROOT / 'sources/lm_understand-logs/exp8_math/practice').glob('*.jsonl'):
        shutil.copyfile(f, practice / f.name)
    manifest = []
    specs = [('HuggingFaceH4/MATH-500', None, 'math500'), ('openai/gsm8k', 'main', 'gsm8k')]
    for repo, config, name in specs:
        with urllib.request.urlopen(f'https://huggingface.co/api/datasets/{repo}') as response:
            revision = json.load(response)['sha']
        ds = load_dataset(repo, config, split='test', revision=revision)
        if name == 'math500':
            items = [{'id': r['unique_id'], 'problem': r['problem'], 'answer': r['answer']} for r in ds]
        else:
            items = [{'id': f'gsm8k_test_{i}', 'problem': r['question'], 'answer': r['answer'].split('####')[-1].strip().replace(',', '')} for i, r in enumerate(ds)]
        path = ROOT / f'data/locked_test/{name}.jsonl'
        dump(path, items)
        manifest.append({'repo': repo, 'revision': revision, 'rows': len(items), 'sha256': hashlib.sha256(path.read_bytes()).hexdigest()})
    with urllib.request.urlopen('https://api.github.com/repos/google-research/google-research/commits/master') as response:
        revision = json.load(response)['sha']
    target = ROOT / 'sources/instruction_following_eval'
    target.mkdir(exist_ok=True)
    for name in ('instructions.py', 'instructions_registry.py', 'instructions_util.py', 'evaluation_lib.py', 'data/input_data.jsonl'):
        with urllib.request.urlopen(f'https://raw.githubusercontent.com/google-research/google-research/{revision}/instruction_following_eval/{name}') as response:
            content = response.read()
        out = ROOT / 'data/locked_test/ifeval.jsonl' if name.startswith('data/') else target / name
        out.write_bytes(content)
    (target / '__init__.py').touch()
    manifest.append({'repo': 'google-research/google-research', 'revision': revision, 'path': 'instruction_following_eval'})
    save_json(ROOT / 'data/provenance.json', manifest)
    from split_data import main as split_data
    split_data()
    print('Данные скачаны, исследовательские и отложенные выборки разделены.')


if __name__ == '__main__':
    main()