import io
import modal
import os
import subprocess
import sys
import time
import shutil
from datetime import datetime
from typing import Dict, List, Tuple

# --- Modal Setup ---
app = modal.App("hunyuan-training")
volume = modal.Volume.from_name(
    "hunyuan-volume", create_if_missing=True
)
image = modal.Image.debian_slim().apt_install(
    "git", "wget", "curl"
).pip_install(
    "tomli",
    "toml",
    "diffusers",
    "transformers",
    "accelerate",
    "safetensors",
    "tensorboard",
    "ftfy",
    "protobuf==3.20.1",
    "deepspeed",
    "huggingface_hub[hf_transfer]",
).run_commands(
    "pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118"
).env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})

# --- Constants ---
GIT_REPO = "https://github.com/tdrussell/diffusion-pipe.git"
REMOTE_ROOT = "/root/volume"

# --- Helper Functions ---
@app.function(image=image, volumes={REMOTE_ROOT: volume})
def hf_download_repo(repo_id: str, model_type: str, ignore_patterns: List[str] = None):
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id=repo_id,
        local_dir=f"/root/volume/models/{model_type}",
        ignore_patterns=ignore_patterns or [],
    )
    print(f"Repository {repo_id} downloaded to /root/volume/models/{model_type}")

@app.function(image=image, volumes={REMOTE_ROOT: volume})
def hf_download_file(repo_id: str, filename: str, model_type: str):
    from huggingface_hub import hf_hub_download

    hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=f"/root/volume/models/{model_type}",
    )
    print(f"File {filename} from {repo_id} downloaded to /root/volume/models/{model_type}")

# --- Modal Functions ---
@app.function(image=image, volumes={REMOTE_ROOT: volume})
def download_models():
    models_to_download_from_hf_repos: List[Tuple[str, str, List[str]]] = [
        ("Kijai/llava-llama-3-8b-text-encoder-tokenizer", "llm"),
        ("openai/clip-vit-large-patch14", "clip", ["flax_model.msgpack", "pytorch_model.bin", "tf_model.h5"]),  # Exclude these files
    ]

    models_to_download_from_hf_files: List[Tuple[str, str, str]] = [
        ("Kijai/HunyuanVideo_comfy", "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors", "hunyuan"),
        ("Kijai/HunyuanVideo_comfy", "hunyuan_video_vae_bf16.safetensors", "hunyuan"),
    ]

    # Download entire HF repos
    print("Downloading entire Hugging Face repositories...")
    for repo in models_to_download_from_hf_repos:
        if len(repo) == 3:
            repo_id, model_type, ignore_patterns = repo
        elif len(repo) == 2:
            repo_id, model_type = repo
            ignore_patterns = []
        else:
            raise ValueError(f"Unexpected tuple length: {len(repo)}")
    
        hf_download_repo.remote(repo_id, model_type, ignore_patterns)

    # Download individual files from Hugging Face
    print("Downloading individual files from Hugging Face...")
    for repo_id, filename, model_type in models_to_download_from_hf_files:
        hf_download_file.remote(repo_id, filename, model_type)

@app.function(
    image=image,
    gpu="T4",
    volumes={REMOTE_ROOT: volume},
    timeout=86400,
)
def train():
    import os
    import subprocess

    # Clone the diffusion-pipe repository
    if not os.path.exists('/root/diffusion-pipe'):
        subprocess.run(['git', 'clone', '--recurse-submodules', GIT_REPO, '/root/diffusion-pipe'], check=True)
    else:
        print("diffusion-pipe repository already cloned.")

    # Change directory to diffusion-pipe
    os.chdir('/root/diffusion-pipe')

    # Install requirements
    subprocess.run(['pip', 'install', '-r', 'requirements.txt'], check=True)

    # Prepare 'config.toml' and 'dataset.toml' with adjusted paths

    # Prepare 'config.toml'
    config_toml_content = """
# Dataset config file.
output_dir = '/root/volume/training_output'
dataset = 'dataset.toml'

# Training settings
epochs = 50
micro_batch_size_per_gpu = 1
pipeline_stages = 1
gradient_accumulation_steps = 4
gradient_clipping = 1.0
warmup_steps = 100

# eval settings
eval_every_n_epochs = 5
eval_before_first_step = true
eval_micro_batch_size_per_gpu = 1
eval_gradient_accumulation_steps = 1

# misc settings
save_every_n_epochs = 5
checkpoint_every_n_minutes = 30
activation_checkpointing = true
partition_method = 'parameters'
save_dtype = 'bfloat16'
caching_batch_size = 1
steps_per_print = 1
video_clip_mode = 'single_middle'

[model]
type = 'hunyuan-video'
transformer_path = '/root/volume/models/hunyuan/hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors'
vae_path = '/root/volume/models/hunyuan/hunyuan_video_vae_bf16.safetensors'
llm_path = '/root/volume/models/llm'
clip_path = '/root/volume/models/clip'
dtype = 'bfloat16'
transformer_dtype = 'float8'
timestep_sample_method = 'logit_normal'

[adapter]
type = 'lora'
rank = 64
dtype = 'bfloat16'

[optimizer]
type = 'adamw_optimi'
lr = 5e-5
betas = [0.9, 0.99]
weight_decay = 0.02
eps = 1e-8
"""
    with open('config.toml', 'w') as f:
        f.write(config_toml_content)

    # Prepare 'dataset.toml'
    dataset_toml_content = """
# Resolution settings.
resolutions = [512]

# Aspect ratio bucketing settings
enable_ar_bucket = true
min_ar = 0.5
max_ar = 2.0
num_ar_buckets = 7

# Frame buckets (1 is for images)
frame_buckets = [1]

[[directory]]
# Path to your dataset
path = '/root/volume/training_data/images'
num_repeats = 5
"""
    with open('dataset.toml', 'w') as f:
        f.write(dataset_toml_content)

    # Run the training command
    os.environ['NCCL_P2P_DISABLE'] = "1"
    os.environ['NCCL_IB_DISABLE'] = "1"
    os.environ['CUDA_VISIBLE_DEVICES'] = "0"

    print("Starting training...")
    subprocess.run(['deepspeed', '--num_gpus=1', 'train.py', '--deepspeed', '--config', 'config.toml'], check=True)
    print("Training completed.")

# --- Entry Points ---
@app.local_entrypoint()
def download_models_entrypoint():
    # First, download all necessary models
    print("Downloading models...")
    download_models.remote()
    print("Models downloaded.")
    
@app.local_entrypoint()
def train_entrypoint():
    # Start training
    train.remote()

...
  Downloading flash_attn-2.7.2.post1.tar.gz (3.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.1/3.1 MB 109.5 MB/s eta 0:00:00
  Preparing metadata (setup.py): started
  Preparing metadata (setup.py): finished with status 'error'
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [20 lines of output]
      fatal: not a git repository (or any of the parent directories): .git
      /tmp/pip-install-kpx_ejik/flash-attn_5dc53f15189c4621b288dc70a19fd132/setup.py:99: UserWarning: flash_attn was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.
        warnings.warn(
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "/tmp/pip-install-kpx_ejik/flash-attn_5dc53f15189c4621b288dc70a19fd132/setup.py", line 177, in <module>
          CUDAExtension(
        File "/usr/local/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 1048, in CUDAExtension
          library_dirs += library_paths(cuda=True)
        File "/usr/local/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 1179, in library_paths
          if (not os.path.exists(_join_cuda_home(lib_dir)) and
        File "/usr/local/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 2223, in _join_cuda_home
          raise EnvironmentError('CUDA_HOME environment variable is not set. '
      OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.


      torch.__version__  = 2.0.1+cu118


      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.
Traceback (most recent call last):
  File "/pkg/modal/_runtime/container_io_manager.py", line 757, in handle_input_exception
    yield
  File "/pkg/modal/_container_entrypoint.py", line 238, in run_input_sync
    res = io_context.call_finalized_function()
  File "/pkg/modal/_runtime/container_io_manager.py", line 181, in call_finalized_function
    res = self.finalized_function.callable(*args, **kwargs)
  File "/root/app.py", line 112, in train
    subprocess.run(['pip', 'install', '-r', 'requirements.txt'], check=True)
  File "/usr/local/lib/python3.10/subprocess.py", line 526, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['pip', 'install', '-r', 'requirements.txt']' returned non-zero exit status 1.

Stopping app - uncaught exception raised locally: CalledProcessError(1, ['pip', 'install', '-r', 'requirements.txt']).
┌─────────────────────────────── Traceback (most recent call last) ────────────────────────────────┐
│ E:\Modal\hunyuan-video-lora-training\app.py:213 in train_entrypoint                              │
│                                                                                                  │
│   212 │   # Start training                                                                       │
│ > 213 │   train.remote()                                                                         │
│   214                                                                                            │
│                                                                                                  │
│ C:\Program Files\Python310\lib\site-packages\synchronicity\synchronizer.py:592 in proxy_method   │
│                                                                                                  │
│ C:\Program Files\Python310\lib\site-packages\synchronicity\combined_types.py:29 in __call__      │
│                                                                                                  │
│               ...Remote call to Modal Function (ta-01JGPSE7JNAC6MVXGSHJ3NDFT0)...                │
│                                                                                                  │
│ /root/app.py:112 in train                                                                        │
│                                                                                                  │
│ > 112 subprocess.run(['pip', 'install', '-r', 'requirements.txt'], check=True)                   │
│                                                                                                  │
│                                                                                                  │
│ /usr/local/lib/python3.10/subprocess.py:526 in run                                               │
│                                                                                                  │
│ > 526 raise CalledProcessError(retcode, process.args,                                            │
│                                                                                                  │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘
CalledProcessError: Command '['pip', 'install', '-r', 'requirements.txt']' returned non-zero exit status 1.