import os

from dotenv import load_dotenv
load_dotenv()  # Only needed locally


# Prevent libgomp crashes in Spaces
os.environ["OMP_NUM_THREADS"] = "1"

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["HTTP_PROXY"] = ""
os.environ["HTTPS_PROXY"] = ""
os.environ["http_proxy"] = ""
os.environ["https_proxy"] = ""

import threading
import requests
from huggingface_hub import snapshot_download
import gradio as gr
import base64
import io
import json
import uuid

# Import pose transfer related modules
try:
    from diffusers import DDPMScheduler
    from defaults import pose_transfer_C as cfg
    from pose_transfer_train import build_model
    from models import UNet, VariationalAutoencoder
    import torch
    import numpy as np
    import pandas as pd
    from pose_utils import (cords_to_map, draw_pose_from_cords,
                            load_pose_cords_from_strings)
    import random
    from PIL import Image
    from torchvision import transforms
    import copy
except ImportError as e:
    print(f"Import error (models not downloaded yet): {e}")

# ==============================
# CONFIGURATION - MODIFIED: Use temp directory instead of persistent
# ==============================
MODEL_REPO = "recky101/new_l_cfld_model"
# Changed from persistent to temporary directory
MODEL_DIR = "/tmp/models"  # Temporary directory that gets cleared on restart

# ==============================
# GLOBALS
# ==============================
download_thread = None
download_log = []
cancel_download = False
model_ready = False
vae = None
model = None
unet = None
noise_scheduler = None
test_pairs = None
annotation_file = None

# ==============================
# UTILS
# ==============================
def log(msg: str):
    """Append a log message and return the full log as a string."""
    global download_log
    download_log.append(msg)
    print(msg)
    return "\n".join(download_log)

def test_huggingface_connectivity():
    """Test if we can reach Hugging Face servers."""
    try:
        response = requests.get("https://huggingface.co/", timeout=10)
        if response.status_code == 200:
            return True, "✅ Can reach Hugging Face"
        else:
            return False, f"❌ Hugging Face returned status {response.status_code}"
    except Exception as e:
        return False, f"❌ Cannot reach Hugging Face: {str(e)}"

def is_model_ready():
    """Check if required folders already exist in temp directory."""
    # Always return False to force download on each start
    return False

def get_directory_tree(root_dir, indent=""):
    """Recursively build a directory tree string for logs."""
    tree_str = ""
    try:
        items = sorted(os.listdir(root_dir))
    except Exception as e:
        return f"{indent}❌ [Error accessing {root_dir}]: {e}\n"

    for i, item in enumerate(items):
        path = os.path.join(root_dir, item)
        connector = "└── " if i == len(items) - 1 else "├── "
        tree_str += f"{indent}{connector}{item}\n"
        if os.path.isdir(path):
            tree_str += get_directory_tree(path, indent + ("    " if i == len(items) - 1 else "│   "))
    return tree_str

def download_individual_folders_with_retry():
    """Alternative method with robust retry logic."""
    try:
        folders_to_download = ["checkpoints", "pretrained_models", "fashion"]
        max_retries = 3
        
        for folder in folders_to_download:
            if cancel_download:
                log("⛔ Download cancelled during individual folder download.")
                return
                
            folder_path = os.path.join(MODEL_DIR, folder)
            if os.path.exists(folder_path):
                log(f"✅ Folder {folder} already exists, skipping...")
                continue
                
            for attempt in range(max_retries):
                try:
                    log(f"⬇️ Downloading folder: {folder} (Attempt {attempt + 1}/{max_retries})")
                    
                    snapshot_download(
                        repo_id=MODEL_REPO,
                        local_dir=folder_path,
                        resume_download=True,
                        local_dir_use_symlinks=False,
                        allow_patterns=f"{folder}/*"
                    )
                    log(f"✅ Successfully downloaded {folder}")
                    break
                    
                except Exception as e:
                    if attempt == max_retries - 1:
                        log(f"❌ Failed to download {folder} after {max_retries} attempts: {str(e)}")
                    else:
                        log(f"⚠️ Attempt {attempt + 1} failed for {folder}: {str(e)}")
                        import time
                        time.sleep(5)
            
        log("✅ All folders processed.")
        
    except Exception as e:
        log(f"❌ Individual folder download failed: {str(e)}")

def download_models():
    """Download models with logging - will download on every start."""
    global cancel_download, model_ready, vae, model, unet, noise_scheduler, test_pairs, annotation_file
    
    try:
        # Clear previous downloads if they exist
        import shutil
        if os.path.exists(MODEL_DIR):
            shutil.rmtree(MODEL_DIR)
            log(f"🧹 Cleared previous model directory: {MODEL_DIR}")
        
        # Test connectivity first
        success, message = test_huggingface_connectivity()
        log(message)
        if not success:
            log("🌐 Please check your internet connection and try again")
            return

        os.makedirs(MODEL_DIR, exist_ok=True)

        log("⬇️ Starting model download from Hugging Face Hub...")
        log("ℹ️ Models will be downloaded to temporary storage and will be cleared on restart.")
        
        # Add retry logic
        max_retries = 3
        for attempt in range(max_retries):
            try:
                if cancel_download:
                    log("⛔ Download cancelled during process.")
                    return
                    
                log(f"🔄 Attempt {attempt + 1}/{max_retries}")
                
                snapshot_download(
                    repo_id=MODEL_REPO,
                    local_dir=MODEL_DIR,
                    resume_download=True,
                    local_dir_use_symlinks=False,
                )
                break
                
            except Exception as e:
                if attempt == max_retries - 1:
                    raise e
                log(f"⚠️ Attempt {attempt + 1} failed: {str(e)}")
                import time
                time.sleep(10)
        
        if cancel_download:
            log("⛔ Download cancelled during process.")
            return
            
        log("✅ Download completed successfully.")
        log("📂 Listing downloaded directory structure...")
        tree = get_directory_tree(MODEL_DIR)
        log(f"\n{tree}")
        model_ready = True
        initialize_models()

    except Exception as e:
        error_msg = str(e)
        log(f"❌ Download failed after {max_retries} attempts: {error_msg}")
        log("💡 Trying alternative download method...")
        download_individual_folders_with_retry()

def initialize_models():
    """Initialize the pose transfer models after download."""
    global vae, model, unet, noise_scheduler, test_pairs, annotation_file
    
    try:
        log("🔄 Initializing pose transfer models...")
        
        # DEBUG: Print entire directory tree with sizes
        log("📂 Verifying model files in temporary directory...")
        for root, dirs, files in os.walk(MODEL_DIR):
            level = root.replace(MODEL_DIR, "").count(os.sep)
            indent = " " * 4 * (level)
            log(f"{indent}{os.path.basename(root)}/")
            subindent = " " * 4 * (level + 1)
            for f in files:
                size = os.path.getsize(os.path.join(root, f)) / (1024*1024)
                log(f"{subindent}{f} ({size:.2f} MB)")
        
        log("🔄 Initializing pose transfer models...")
        
        # Initialize models
        noise_scheduler = DDPMScheduler.from_pretrained(os.path.join(MODEL_DIR, "pretrained_models/scheduler/scheduler_config.json"))
        log("✅ Noise scheduler initialized")
        
        vae = VariationalAutoencoder(pretrained_path=os.path.join(MODEL_DIR, "pretrained_models/vae")).eval().requires_grad_(False).cuda()
        log("✅ VAE initialized")
        
        model = build_model(cfg).eval().requires_grad_(False).cuda()
        log("✅ Main model initialized")
        
        unet = UNet(cfg).eval().requires_grad_(False).cuda()
        log("✅ UNet initialized")
        
        # Load model weights
        model.load_state_dict(torch.load(
            os.path.join(MODEL_DIR, "checkpoints", "pytorch_model.bin"), map_location="cpu"
        ), strict=False)
        log("✅ Model weights loaded")
        
        unet.load_state_dict(torch.load(
            os.path.join(MODEL_DIR, "checkpoints", "pytorch_model_1-001.bin"), map_location="cpu"
        ), strict=False)
        log("✅ UNet weights loaded")
        
        # Load test data
        test_pairs_path = os.path.join(MODEL_DIR, "fashion", "fasion-resize-pairs-test.csv")
        test_pairs = pd.read_csv(test_pairs_path)
        log("✅ Test pairs loaded")
        
        annotation_file_path = os.path.join(MODEL_DIR, "fashion", "fasion-resize-annotation-test.csv")
        annotation_file = pd.read_csv(annotation_file_path, sep=':')
        annotation_file = annotation_file.set_index('name')
        log("✅ Annotation file loaded")
        
        log("✅ Models initialized successfully")
        
    except Exception as e:
        log(f"❌ Error initializing models: {str(e)}")

def build_pose_img(annotation_file, img_path):
    """Build pose image from annotation file."""
    log(f"📄 img_path: {img_path}")
    log(f"📄 basename(img_path): {os.path.basename(img_path)}")
    log(f"📄 Index Sample: {annotation_file.index[:5]}")
    log(f"📄 Does key exist?: {os.path.basename(img_path) in annotation_file.index}")

    string = annotation_file.loc[os.path.basename(img_path)]
    array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
    pose_map = torch.tensor(cords_to_map(array, (256, 256), (256, 176)).transpose(2, 0, 1), dtype=torch.float32)
    pose_img = torch.tensor(draw_pose_from_cords(array, (256, 256), (256, 176)).transpose(2, 0, 1) / 255., dtype=torch.float32)
    pose_img = torch.cat([pose_img, pose_map], dim=0)
    return pose_img

def pose_transfer(source_image, test_pair_index):
    test_pair_index = int(test_pair_index)
    """Perform pose transfer from source image to target pose."""
    global vae, model, unet, noise_scheduler, test_pairs, annotation_file
    
    if not model_ready:
        raise ValueError("Models not ready. Please download models first.")
    
    if test_pair_index < 0 or test_pair_index >= len(test_pairs):
        raise ValueError(f"Test pair index must be between 0 and {len(test_pairs)-1}")
    
    # Get target image path
    img_to_path = test_pairs.iloc[test_pair_index]["to"]
    log(f"🔄 img_to_path: {img_to_path}")
    
    # Build pose image
    pose_img_tensor = build_pose_img(annotation_file, img_to_path).unsqueeze(0)
    
    # Transform source image
    trans = transforms.Compose([
        transforms.Resize([256, 256], interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    img_from_tensor = trans(source_image).unsqueeze(0)
    
    # Perform pose transfer
    with torch.no_grad():
        c_new, down_block_additional_residuals, up_block_additional_residuals = model({
            "img_cond": img_from_tensor.cuda(), "pose_img": pose_img_tensor.cuda()})
        noisy_latents = torch.randn((1, 4, 64, 64)).cuda()
        weight_dtype = torch.float32
        bsz = 1

        c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[bsz:]])
        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample]).to(dtype=weight_dtype)
                                           for sample in down_block_additional_residuals]
        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v]).to(dtype=weight_dtype)
                                         for k, v in up_block_additional_residuals.items()}

        noise_scheduler.set_timesteps(cfg.TEST.NUM_INFERENCE_STEPS)
        for t in noise_scheduler.timesteps:
            inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents], dim=0)
            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
            noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
                              down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
                              up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))

            noise_pred_uc, noise_pred_down, noise_pred_full = noise_pred.chunk(3)
            noise_pred = noise_pred_uc + \
                         cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
                         cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down)
            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]

        sampling_imgs = vae.decode(noisy_latents) * 0.5 + 0.5  # denormalize
        sampling_imgs = sampling_imgs.clamp(0, 1)
    
    # Convert to PIL image
    output_img = Image.fromarray(
        (sampling_imgs[0] * 255.)
        .permute((1, 2, 0))
        .long()
        .cpu()
        .numpy()
        .astype(np.uint8)
    )
    
    log("✅ Pose transfer completed successfully")
    log(f"🔄 output_img size: {output_img.size}")

    return output_img

def start_download():
    """Start model download in a separate thread - will always download fresh."""
    global download_thread, download_log, cancel_download
    if download_thread and download_thread.is_alive():
        return "⚠️ Download already running..."

    download_log = []
    cancel_download = False
    download_thread = threading.Thread(target=download_models)
    download_thread.start()
    return "📥 Download started (fresh download each time)..."

def get_download_status():
    """Get the latest log status."""
    global download_thread
    status = "\n".join(download_log) if download_log else "Preparing download..."
    if download_thread and download_thread.is_alive():
        return status + "\n\n⏳ Download in progress..."
    elif model_ready:
        return status + "\n\n✅ Models are ready and loaded in memory (temporary storage)."
    return status

def cancel_download_fn():
    """Cancel request handler."""
    global cancel_download
    cancel_download = True
    log("⛔ Download cancelled by user.")
    return "Download cancelled."

# ==============================
# AUTO-START DOWNLOAD ON STARTUP
# ==============================

# Start download automatically when the app starts
print("🚀 Starting model download on startup...")
download_thread = threading.Thread(target=download_models)
download_thread.start()

# ==============================
# GRADIO UI ONLY (NO API ENDPOINTS)
# ==============================
def gradio_pose_transfer(source_image, test_pair_index):
    test_pair_index = int(test_pair_index)
    """Gradio interface for pose transfer."""
    try:
        if not model_ready:
            return None, "Models not ready. Please wait for download to complete."
        
        if test_pair_index < 0 or test_pair_index >= len(test_pairs):
            return None, f"Test pair index must be between 0 and {len(test_pairs)-1}"
        
        # Perform pose transfer
        output_image = pose_transfer(source_image, test_pair_index)
        
        return output_image, "Pose transfer successful"
        
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        log(f"❌ Error during pose transfer:\n{tb}")
        log(f"❌ the value received is as follow: \n{test_pair_index}")
        return None, f"Error during pose transfer: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("## 🧩 Model Downloader & Pose Transfer")
    gr.Markdown(f"**Model Source:** [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})")
    gr.Markdown("**Storage:** Models are downloaded to temporary storage and will be cleared on restart.")
    gr.Markdown("**Status:** Download starts automatically on app launch.")

    with gr.Tab("Model Download"):
        with gr.Row():
            start_btn = gr.Button("🔄 Re-download Models")
            cancel_btn = gr.Button("❌ Cancel Download")

        status_box = gr.Textbox(
            label="Download Logs",
            lines=25,
            interactive=False,
            placeholder="Models downloading on startup..."
        )

        # Button bindings
        start_btn.click(fn=start_download, inputs=None, outputs=status_box)
        cancel_btn.click(fn=cancel_download_fn, inputs=None, outputs=status_box)

        # Periodic refresh of logs
        demo.load(fn=get_download_status, inputs=None, outputs=status_box, every=2)

    with gr.Tab("Pose Transfer"):
        gr.Markdown("## Pose Transfer")
        
        with gr.Row():
            with gr.Column():
                source_image = gr.Image(label="Source Image", type="pil")
                test_pair_index = gr.Number(
                    label="Test Pair Index", 
                    value=0, 
                    minimum=0, 
                    maximum=4039
                )
                generate_btn = gr.Button("🚀 Generate Pose Transfer")
            
            with gr.Column():
                output_image = gr.Image(label="Output Image", type="pil")
                status_message = gr.Textbox(label="Status", interactive=False)
        
        generate_btn.click(
            fn=gradio_pose_transfer,
            inputs=[source_image, test_pair_index],
            outputs=[output_image, status_message]
        )

from fastapi import FastAPI, Form, UploadFile
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import base64
from io import BytesIO
import uvicorn
import cloudinary
import cloudinary.uploader

# Configure Cloudinary from environment variables
cloudinary.config(
    cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"),
    api_key=os.environ.get("CLOUDINARY_API_KEY"),
    api_secret=os.environ.get("CLOUDINARY_API_SECRET")
)

app = FastAPI()

# Allow CORS (needed for Postman / frontends)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/pose-transfer")
async def pose_transfer_api(
    test_pair_index: int = Form(...),
    source_image: UploadFile = None,
    image_base64: str = Form(None)
):
    try:
        # Get input image
        if source_image:
            image_bytes = await source_image.read()
            image = Image.open(BytesIO(image_bytes)).convert("RGB")
        elif image_base64:
            if "," in image_base64:
                image_base64 = image_base64.split(",", 1)[1]
            image_bytes = base64.b64decode(image_base64)
            image = Image.open(BytesIO(image_bytes)).convert("RGB")
        else:
            return JSONResponse(
                content={"status": "error", "message": "No image provided"},
                status_code=400
            )

        # Check if models are ready
        if not model_ready:
            return JSONResponse(
                content={"status": "error", "message": "Models are still downloading. Please wait."},
                status_code=503
            )

        # Check if Cloudinary is configured
        if not all([cloudinary.config().cloud_name, cloudinary.config().api_key, cloudinary.config().api_secret]):
            return JSONResponse(
                content={"status": "error", "message": "Cloudinary not properly configured"},
                status_code=500
            )

        # 🔹 Call your pose transfer model
        output_image = pose_transfer(image, test_pair_index)

        # Save image to buffer
        buffer = BytesIO()
        output_image.save(buffer, format="PNG")
        buffer.seek(0)
        
        # Upload to Cloudinary
        upload_result = cloudinary.uploader.upload(
            buffer,
            folder="pose_transfer",  # Optional folder in Cloudinary
            public_id=f"pose_transfer_{uuid.uuid4().hex[:8]}",  # Unique ID
            overwrite=True,
            resource_type="image"
        )
        
        # Get the URL from Cloudinary response
        image_url = upload_result.get('secure_url', upload_result.get('url'))
        
        return JSONResponse(
            content={
                "status": "success",
                "message": "Pose transfer completed and image uploaded to Cloudinary",
                "image_url": image_url
            }
        )
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        return JSONResponse(
            content={"status": "error", "message": str(e), "traceback": tb},
            status_code=500
        )

if __name__ == "__main__":
    # Note: Models will be downloaded automatically on startup via the thread started above
    # No need to check if models exist since we always download fresh
    
    # Launch both Gradio and FastAPI
    # You might want to run them on different ports or use a different approach
    # For simplicity, we'll just run the FastAPI app which includes the Gradio interface
    uvicorn.run(app, host="0.0.0.0", port=7860)