# setup_models.py
import os
import shutil
from huggingface_hub import hf_hub_download, snapshot_download

# Configuration
REPO_ID = "ShadowGard3n/AgroVision-Models"
DEST_MODELS_DIR = "models"

# Ensure the destination models directory exists
os.makedirs(DEST_MODELS_DIR, exist_ok=True)

print(f"--- Starting Download from {REPO_ID} ---")

# 1. Download final_output.csv to the root directory
try:
    print("Downloading final_output.csv...")
    hf_hub_download(
        repo_id=REPO_ID,
        filename="final_output.csv",
        local_dir=".",
        local_dir_use_symlinks=False
    )
    print("✅ final_output.csv downloaded.")
except Exception as e:
    print(f"❌ Error downloading csv: {e}")

# 2. Download 'market_models' folder and move files to 'models/'
try:
    print("Downloading market_models...")
    market_path = snapshot_download(
        repo_id=REPO_ID,
        allow_patterns="market_models/*",
        local_dir="temp_download",
        local_dir_use_symlinks=False
    )
    
    # Move files from temp_download/market_models to models/
    source_dir = os.path.join(market_path, "market_models")
    if os.path.exists(source_dir):
        for file_name in os.listdir(source_dir):
            full_file_name = os.path.join(source_dir, file_name)
            if os.path.isfile(full_file_name):
                shutil.move(full_file_name, DEST_MODELS_DIR)
    print("✅ Market models moved to /models.")
except Exception as e:
    print(f"❌ Error downloading market models: {e}")

# 3. Download 'other_models' folder and move files to 'models/'
try:
    print("Downloading other_models...")
    other_path = snapshot_download(
        repo_id=REPO_ID,
        allow_patterns="other_models/*",
        local_dir="temp_download_other",
        local_dir_use_symlinks=False
    )
    
    # Move files from temp_download_other/other_models to models/
    source_dir = os.path.join(other_path, "other_models")
    if os.path.exists(source_dir):
        for file_name in os.listdir(source_dir):
            full_file_name = os.path.join(source_dir, file_name)
            if os.path.isfile(full_file_name):
                shutil.move(full_file_name, DEST_MODELS_DIR)
    print("✅ Other models moved to /models.")
except Exception as e:
    print(f"❌ Error downloading other models: {e}")

# Cleanup temp folders
shutil.rmtree("temp_download", ignore_errors=True)
shutil.rmtree("temp_download_other", ignore_errors=True)

print("--- Setup Complete ---")