import os
import subprocess
from pathlib import Path

import torch

from .tokenizer import TokenizerInterface

SIGMA_DATA = 1.0
T5_EMBEDDING_MAX_LENGTH = 512


def _data_root():
    default_root = "/opt/ml" if os.environ.get("SAGEMAKER", "") == "enabled" else \
        os.environ.get("ANYDATA_LOCAL_ROOT", "/data")
    return Path(os.environ.get("ANYDATA_CACHE_ROOT", default_root)).resolve()


def _resolve_s3_path(path, is_folder):
    if not isinstance(path, str) or not path.startswith("s3://"):
        return path
    stripped = path.replace("s3://", "")
    parts = stripped.split("/")
    if len(parts) < 2:
        return path
    rel = "/".join(parts[1:])
    local = _data_root() / rel
    local.parent.mkdir(parents=True, exist_ok=True)
    if is_folder:
        if not local.exists():
            subprocess.run(["aws", "s3", "sync", path, str(local), "--quiet"], check=True)
    else:
        if not local.exists():
            subprocess.run(["aws", "s3", "cp", path, str(local), "--quiet"], check=True)
    return str(local)


class VAE(torch.nn.Module):
    def __init__(self, device, vae_path, text_encoder_path, downsample=(4, 8, 8), load_text_encoder=True):
        super().__init__()
        self.tensor_kwargs = dict(device=device, dtype=torch.bfloat16)
        self.downsample = downsample
        self.text_encoder = None

        vae_local = _resolve_s3_path(vae_path, is_folder=False)
        self.video_tokenizer = TokenizerInterface(
            chunk_duration=25,
            load_mean_std=False,
            name="tokenizer",
            vae_pth=vae_local,
        )

        if load_text_encoder:
            from .text_encoder import CosmosT5TextEncoder
            text_local = _resolve_s3_path(text_encoder_path, is_folder=True)
            self.text_encoder = CosmosT5TextEncoder(device=device, cache_dir=text_local)

        self.cached_lat_language = {}

    @torch.no_grad()
    def encode_text(self, prompt):
        if self.text_encoder is None:
            raise RuntimeError("Text encoder is disabled for this VAE instance")
        prompts = prompt if isinstance(prompt, list) else [prompt]
        text_embedding, text_mask = self.text_encoder.encode_prompts(
            prompts, max_length=T5_EMBEDDING_MAX_LENGTH, return_mask=True
        )
        return {
            "t5_text_embeddings": text_embedding.to(**self.tensor_kwargs),
            "t5_text_mask": text_mask.to(**self.tensor_kwargs),
        }

    @torch.no_grad()
    def encode_single_video(self, raw_video):
        _, cp, tp, _, _ = raw_video.shape
        if cp != 3:
            raise ValueError(f"Raw video must have 3 channels, got {cp}")
        tcf = self.video_tokenizer.temporal_compression_factor
        if (tp - 1) % tcf != 0:
            raise ValueError(f"raw Tp-1 must be divisible by temporal compression factor {tcf}, got {tp}")
        raw_video = raw_video.to(**self.tensor_kwargs, non_blocking=True)
        latent_video = self.video_tokenizer.encode(raw_video)
        return latent_video * SIGMA_DATA

    @torch.no_grad()
    def encode_single_video_rgb(self, raw_video):
        if raw_video.shape[1] != 3:
            raise ValueError("RGB video must have 3 channels")
        return self.encode_single_video(raw_video * 2.0 - 1.0)

    @torch.no_grad()
    def encode_single_video_depth(self, raw_video, key=None):
        del key
        if raw_video.shape[1] != 1:
            raise ValueError("Depth video must have 1 channel")
        raw_video = raw_video * 2.0 - 1.0
        raw_video = torch.cat([raw_video] * 3, dim=1)
        return self.encode_single_video(raw_video)