import argparse
import hashlib
import json
import math
from pathlib import Path
import shutil
import struct
import time
import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from common import ROOT, inside, save_json, append, permit_compute

torch.set_grad_enabled(False)
FROZEN = ('embed_tokens', 'lm_head', 'model.language_model.norm.weight', 'model.norm.weight')


def headers(model):
    result = {}
    for file in sorted(inside(model).glob('*.safetensors')):
        with file.open('rb') as f:
            length = struct.unpack('<Q', f.read(8))[0]
            head = json.loads(f.read(length))
        for name, info in head.items():
            if name != '__metadata__':
                if name in result:
                    raise ValueError(f'Duplicate tensor: {name}')
                result[name] = dict(info, file=str(file))
    return result


def load(model, name):
    info = headers(model)[name]
    with safe_open(info['file'], framework='pt', device='cpu') as f:
        return f.get_tensor(name)


def validate(name, original, value):
    if not name.startswith('model.language_model.layers.') or any(x in name for x in FROZEN):
        raise ValueError('Разрешены только внутренние слои языковой модели')
    if value.shape != original.shape or not torch.isfinite(value).all():
        raise ValueError('Неверная форма или нечисловые значения')
    w, v = original.float(), value.float()
    if torch.equal(w, v):
        raise ValueError('Правка не меняет веса')
    ratio = (w * v).sum() / w.square().sum().clamp_min(1e-30)
    if (v - ratio * w).norm() <= 1e-6 * v.norm().clamp_min(1e-30):
        raise ValueError('Масштабирование всего тензора запрещено')
    cast = value.to(original.dtype).contiguous()
    if not torch.isfinite(cast).all() or torch.equal(cast, original):
        raise ValueError('Правка исчезла или стала нечисловой после приведения dtype')
    return cast


def digest(tensor):
    return hashlib.sha256(tensor.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()


def save_edit(model, directory, name, value, hypothesis):
    permit_compute()
    directory = inside(directory)
    if not directory.is_relative_to(ROOT / 'edits'):
        raise ValueError('Правки сохраняются только в edits/')
    if not hypothesis.strip():
        raise ValueError('Нужна техническая гипотеза')
    original = load(model, name)
    cast = validate(name, original, value)
    directory.mkdir(parents=True, exist_ok=True)
    file = directory / (name + '.safetensors')
    save_file({name: cast}, str(file))
    record = {'time': time.time(), 'source_model': str(inside(model)), 'tensor': name, 'hypothesis': hypothesis, 'status': 'unconfirmed',
              'changed_entries': int((cast != original).sum()), 'tensor_entries': original.numel(),
              'base_sha256': digest(original), 'edited_sha256': digest(cast),
              'relative_frobenius': float((cast.float()-original.float()).norm()/original.float().norm())}
    save_json(directory / (name + '.json'), record)
    append(ROOT / 'logs/experiments.jsonl', record)
    return record


def project_writers(w, direction, columns, alpha):
    w = w.float()
    u = torch.as_tensor(direction, dtype=torch.float32)
    if not torch.isfinite(u).all() or u.norm() == 0:
        raise ValueError('Некорректное направление')
    u = u / u.norm()
    coeff = u @ w
    mask = torch.zeros_like(coeff)
    mask[columns] = coeff[columns]
    return w - float(alpha) * torch.outer(u, mask)


def lineage(model):
    current, ancestors, inherited = inside(model), [], set()
    while True:
        if str(current) in ancestors:
            raise ValueError('Цикл в происхождении модели')
        ancestors.append(str(current))
        manifest = current / 'wsp_manifest.json'
        if not manifest.exists():
            return current, ancestors, inherited
        record = json.loads(manifest.read_text())
        inherited.update(record['tensors'])
        current = inside(record['source'])


def weight_fingerprint(model):
    result = []
    for path in sorted(inside(model).glob('*.safetensors')):
        with path.open('rb') as stream:
            result.append([path.name, hashlib.file_digest(stream, 'sha256').hexdigest()])
    return result


def export(model, edits, out):
    permit_compute()
    model, edits, out = inside(model), inside(edits), inside(out)
    if out.exists():
        raise ValueError('Каталог назначения должен быть новым')
    if not out.is_relative_to(ROOT / 'models') or out == ROOT / 'models/original':
        raise ValueError('Недопустимый каталог модели')
    metadata = headers(model)
    changes = {}
    for file in edits.glob('*.safetensors'):
        with safe_open(str(file), framework='pt', device='cpu') as f:
            for name in f.keys():
                if name in changes:
                    raise ValueError('Повторная правка тензора')
                base = load(model, name)
                changes[name] = validate(name, base, f.get_tensor(name))
                manifest = json.loads((edits / (name + '.json')).read_text())
                if manifest['base_sha256'] != digest(base) or manifest['edited_sha256'] != digest(changes[name]):
                    raise ValueError('Хеш правки не совпал')
    if not changes:
        raise ValueError('Нет правок для экспорта')
    original, ancestors, inherited = lineage(model)
    original_metadata = headers(original)
    if metadata.keys() != original_metadata.keys() or any(
            metadata[name]['shape'] != original_metadata[name]['shape'] or
            metadata[name]['dtype'] != original_metadata[name]['dtype'] for name in metadata):
        raise ValueError('Архитектура родителя отличается от исходной модели')
    cumulative = {}
    for name in sorted(inherited | changes.keys()):
        value = changes[name] if name in changes else load(model, name)
        cumulative[name] = int((value != load(original, name)).sum())
    parent_fingerprint = weight_fingerprint(model)
    original_fingerprint = parent_fingerprint if original == model else weight_fingerprint(original)
    staging = out.with_name(out.name + '.building')
    if staging.exists():
        raise ValueError('Обнаружен незавершённый экспорт')
    staging.mkdir(parents=True)
    for file in model.iterdir():
        if file.is_file() and file.suffix != '.safetensors':
            shutil.copyfile(file, staging / file.name)
    for shard in sorted({x['file'] for x in metadata.values()}):
        with safe_open(shard, framework='pt', device='cpu') as f:
            tensors = {name: changes[name] if name in changes else f.get_tensor(name) for name in f.keys()}
        save_file(tensors, str(staging / Path(shard).name), metadata={'format': 'pt'})
        del tensors
    total = sum(math.prod(x['shape']) for x in metadata.values() if x['dtype'].startswith(('BF', 'F')))
    language_total = sum(math.prod(x['shape']) for name, x in metadata.items()
                         if name.startswith('model.language_model.') and x['dtype'].startswith(('BF', 'F')))
    changed = sum(int((value != load(model, name)).sum()) for name, value in changes.items())
    cumulative_changed = sum(cumulative.values())
    save_json(staging / 'wsp_manifest.json', {'schema_version': 2,
                                           'source': str(model), 'edits': str(edits), 'tensors': list(changes),
                                           'changed_entries': changed, 'stored_parameter_entries': total,
                                           'changed_percent': 100 * changed / total,
                                           'language_parameter_entries': language_total,
                                           'changed_percent_language': 100 * changed / language_total,
                                           'changed_entries_reference': 'immediate_parent',
                                           'parent_weight_fingerprint': parent_fingerprint,
                                           'original_source': str(original), 'ancestor_chain_parent_first': ancestors,
                                           'original_weight_fingerprint': original_fingerprint,
                                           'changed_entries_vs_original': cumulative_changed,
                                           'changed_percent_vs_original': 100 * cumulative_changed / total,
                                           'changed_percent_language_vs_original': 100 * cumulative_changed / language_total,
                                           'changed_entries_vs_original_by_tensor': cumulative,
                                           'status': 'unconfirmed'})
    staging.rename(out)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('command', choices=['inventory', 'export'])
    parser.add_argument('--model', default=str(ROOT / 'models/original'))
    parser.add_argument('--edits')
    parser.add_argument('--out')
    args = parser.parse_args()
    if args.command == 'inventory':
        save_json(ROOT / 'reports/tensor_inventory.json', headers(args.model))
    else:
        export(args.model, args.edits, args.out)