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}",
        local_dir_use_symlinks=False,
        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_id, model_type, ignore_patterns in models_to_download_from_hf_repos:
        hf_download_repo.call(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.call(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():
    # First, download all necessary models
    print("Downloading models...")
    download_models.call()
    print("Models downloaded.")
    
@app.local_entrypoint()
def train():
    # Start training
    train.call()

✓ Created objects.
├── 🔨 Created mount E:\Modal\hunyuan-video-lora-training\app.py
├── 🔨 Created function hf_download_repo.
├── 🔨 Created function hf_download_file.
├── 🔨 Created function download_models.
└── 🔨 Created function train.
Downloading models...
Stopping app - uncaught exception raised locally: AttributeError("'LocalEntrypoint' object has no attribute 'call'").
┌─────────────────────────────── Traceback (most recent call last) ────────────────────────────────┐
│ E:\Modal\hunyuan-video-lora-training\app.py:203 in download_models                               │
│                                                                                                  │
│   202 │   print("Downloading models...")                                                         │
│ > 203 │   download_models.call()                                                                 │
│   204 │   print("Models downloaded.")                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘
AttributeError: 'LocalEntrypoint' object has no attribute 'call'

E:\Modal\hunyuan-video-lora-training>modal run app.py::download_models
✓ Initialized. View run at https://modal.com/apps/pussy/main/ap-Lk0rgQkQo0zLfKAT3YoINb
✓ Created objects.
├── 🔨 Created mount E:\Modal\hunyuan-video-lora-training\app.py
├── 🔨 Created function hf_download_repo.
├── 🔨 Created function hf_download_file.
├── 🔨 Created function download_models.
└── 🔨 Created function train.
Downloading models...
Stopping app - uncaught exception raised locally: AttributeError("'LocalEntrypoint' object has no attribute 'call'").
┌─────────────────────────────── Traceback (most recent call last) ────────────────────────────────┐
│ E:\Modal\hunyuan-video-lora-training\app.py:200 in download_models                               │
│                                                                                                  │
│   199 │   print("Downloading models...")                                                         │
│ > 200 │   download_models.call()                                                                 │
│   201 │   print("Models downloaded.")                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘
AttributeError: 'LocalEntrypoint' object has no attribute 'call'