"""Persistent, unsliced YOR embeddings shared across training recipes."""
import fcntl
import hashlib
import inspect
import json
import logging
import os
from pathlib import Path
import uuid

import cv2
import hydra
from omegaconf import OmegaConf
import torch
from tqdm import tqdm

from datasets.core import TrajectoryDataset
from utils.inference import _embed_in_chunks, eval_mode

ROOT = Path(__file__).resolve().parents[1]
LOG = logging.getLogger(__name__)
SCHEMA = 1
LOADERS = {'datasets.yor.YorDataset', 'datasets.yor_jittered.YorJitteredDataset'}


def source_manifest(dataset_config):
    """Metadata fingerprint: paths, sizes and nanosecond mtimes, not video reads."""
    config = OmegaConf.to_container(dataset_config, resolve=True)
    target = config['_target_']
    if target not in LOADERS:
        raise ValueError(f'Embedding cache does not support {target}')
    root = Path(config['data_directory']).expanduser().resolve()
    if target == 'datasets.yor.YorDataset':
        files = [root/'actions.pth', root/'seq_lengths.pkl', root/'meta.json']
        files.extend(sorted((root/'obses').glob('episode_*.pth')))
    else:
        files = []
        for episode in sorted(root.glob('ep_*')):
            if not episode.is_dir():
                continue
            files.extend(episode/name for name in ['robot_data.csv', 'provenance.json'])
            for view in ['odin', 'wrist_left', 'wrist_right'][:int(config.get('views', 3))]:
                files.extend([episode/(view+'.mp4'), episode/(view+'_index.csv')])
    if not files:
        raise ValueError(f'No dataset files under {root}')
    entries = []
    for path in files:
        stat = path.stat()
        entries.append([str(path.relative_to(root)), stat.st_size, stat.st_mtime_ns])
    # RAM prefetch changes I/O only, not the resulting training examples.
    config.pop('prefetch', None)
    config['data_directory'] = str(root)
    return dict(dataset=config, files=entries)


def cache_identity(dataset_config, encoder_config, encoder):
    model = getattr(encoder, 'module', encoder)
    weight_hash = hashlib.sha256()
    for name, value in sorted(model.state_dict().items()):
        value = value.detach().cpu().contiguous()
        weight_hash.update(name.encode())
        weight_hash.update(str((tuple(value.shape), value.dtype)).encode())
        weight_hash.update(value.reshape(-1).view(torch.uint8).numpy().tobytes())
    source_files = [Path(__file__), ROOT/'datasets/yor.py', ROOT/'datasets/yor_jittered.py',
                    ROOT/'utils/convert_yor.py', ROOT/'utils/inference.py',
                    Path(inspect.getfile(type(model)))]
    identity = dict(schema=SCHEMA, source=source_manifest(dataset_config),
                    encoder=OmegaConf.to_container(encoder_config, resolve=True),
                    weights_sha256=weight_hash.hexdigest(), torch_version=str(torch.__version__),
                    opencv_version=cv2.__version__,
                    code_sha256={str(p.relative_to(ROOT)) if p.is_relative_to(ROOT) else str(p):
                                 hashlib.sha256(p.read_bytes()).hexdigest() for p in source_files})
    key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()
    return key, identity


class CachedYorDataset(TrajectoryDataset):
    """Whole episodes in CPU RAM; splitting and temporal sampling happen later."""
    def __init__(self, payload, cache_path):
        self.data = payload['episodes']
        self.seq_lengths = payload['seq_lengths']
        self.episode_names = payload['episode_names']
        self.source_episodes = payload['source_episodes']
        self.cache_path = str(cache_path)
        if not self.data or len(self.data) != len(self.seq_lengths) or len(self.data) != len(self.episode_names):
            raise ValueError('Invalid cached episode metadata')
        if self.source_episodes is not None and len(self.source_episodes) != len(self.data):
            raise ValueError('Invalid cached source groups')
        for length, (obs, actions, goal) in zip(self.seq_lengths, self.data):
            if length <= 0 or obs.ndim != 4 or actions.shape != (length, 58) or obs.shape[0] != length or goal.shape != (length, 1, 1, 1):
                raise ValueError('Invalid cached trajectory shapes')

    def __len__(self):
        return len(self.data)

    def get_seq_length(self, idx):
        return self.seq_lengths[idx]

    def get_frames(self, idx, frames):
        indices = list(frames)
        return tuple(value[indices] for value in self.data[idx])

    def __getitem__(self, idx):
        return self.data[idx]

    def get_all_actions(self):
        return torch.cat([episode[1] for episode in self.data])


def load_or_build_yor_cache(dataset_config, encoder_config, encoder, cache_root, force=False):
    """Lock builders, publish atomically, and skip raw-loader construction on hits."""
    key, identity = cache_identity(dataset_config, encoder_config, encoder)
    folder = Path(cache_root).expanduser()/Path(identity['source']['dataset']['data_directory']).name
    folder.mkdir(parents=True, exist_ok=True)
    path = folder/(key+'.pt')
    LOG.info('YOR embedding cache: %s', path)
    with (folder/(key+'.lock')).open('a') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if path.exists() and not force:
            try:
                payload = torch.load(path, map_location='cpu', weights_only=True)
                if payload['identity'] != identity:
                    raise ValueError('Cache fingerprint mismatch')
                cached = CachedYorDataset(payload, path)
            except Exception as exc:
                raise RuntimeError(f'Cannot load {path}; rebuild it with --force') from exc
            LOG.info('YOR embedding cache HIT: %s episodes, %.1f MiB; loaded into RAM',
                     len(cached), path.stat().st_size/2**20)
            return cached
        LOG.info('YOR embedding cache MISS: building all frames before splitting/sampling')
        dataset = hydra.utils.instantiate(dataset_config)
        model = getattr(encoder, 'module', encoder)
        device = next(model.parameters()).device
        episodes = []
        with eval_mode(model, no_grad=True):
            for i in tqdm(range(len(dataset)), desc='Caching YOR embeddings'):
                obs, actions, goal = dataset[i]
                embeddings = _embed_in_chunks(model, obs.to(device)).detach().cpu().contiguous()
                if not torch.isfinite(embeddings).all():
                    raise ValueError(f'Nonfinite embeddings in episode {i}')
                episodes.append((embeddings, actions.detach().cpu().clone(), goal.detach().cpu().clone()))
        # Do not publish a cache assembled from changing source files.
        if source_manifest(dataset_config) != identity['source']:
            raise RuntimeError('Dataset changed during cache construction; retry with a stable dataset')
        if hasattr(dataset, 'episodes'):
            names = [p.name for p in dataset.episodes]
        else:
            metadata = json.loads((Path(dataset.data_directory)/'meta.json').read_text())
            names = metadata.get('episodes', [f'episode_{i:05d}' for i in range(len(dataset))])
        payload = dict(identity=identity, episodes=episodes,
                       seq_lengths=[dataset.get_seq_length(i) for i in range(len(dataset))],
                       episode_names=names, source_episodes=getattr(dataset, 'source_episodes', names))
        cached = CachedYorDataset(payload, path)
        temporary = folder/(key+'.'+uuid.uuid4().hex+'.tmp')
        try:
            with temporary.open('wb') as stream:
                torch.save(payload, stream)
                stream.flush()
                os.fsync(stream.fileno())
            os.replace(temporary, path)
        finally:
            temporary.unlink(missing_ok=True)
        (folder/(key+'.json')).write_text(json.dumps(dict(identity=identity,
            episode_names=names, source_episodes=payload['source_episodes'],
            seq_lengths=payload['seq_lengths'], bytes=path.stat().st_size), indent=2)+'\n')
        LOG.info('YOR embedding cache SAVED: %s (%.1f MiB)', path, path.stat().st_size/2**20)
        return cached