"""
Script to push VINE model to HuggingFace Hub

This script helps you push your trained VINE model to the HuggingFace Hub
for easy sharing and distribution.
"""

import os
import sys
import torch
import argparse
from huggingface_hub import notebook_login
from transformers.pipelines import PIPELINE_REGISTRY

# Add the parent directory to the path to import vine_hf
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

os.environ['OPENAI_API_KEY'] = "dummy-key"
from vine_hf import VineConfig, VineModel, VinePipeline


def push_vine_to_hub(
    model_weights_path: str,
    repo_name: str,
    model_name: str = "openai/clip-vit-base-patch32",
    segmentation_method: str = "grounding_dino_sam2",
    commit_message: str = "Upload VINE model",
    private: bool = False
):
    """
    Push VINE model to HuggingFace Hub.
    
    Args:
        model_weights_path: Path to the trained model weights (.pth file)
        repo_name: Name for the repository (e.g., "username/vine-model")
        model_name: CLIP model backbone name
        segmentation_method: Segmentation method used
        commit_message: Commit message for the push
        private: Whether to create a private repository
    """
    
    print("=== Pushing VINE Model to HuggingFace Hub ===")
    
    # 1. Create configuration
    print(f"Creating configuration with backbone: {model_name}")
    config = VineConfig(
        model_name=model_name,
        segmentation_method=segmentation_method
    )
    
    # 2. Initialize model
    print("Initializing model...")
    model = VineModel(config)
    
    # 3. Load trained weights
    if os.path.exists(model_weights_path):
        print(f"Loading weights from: {model_weights_path}")
        try:
            # Try loading with weights_only=False for compatibility
            weights = torch.load(model_weights_path, map_location='cpu', weights_only=False)
            
            # Handle different weight formats
            if isinstance(weights, dict):
                if 'state_dict' in weights:
                    model.load_state_dict(weights['state_dict'])
                elif 'model' in weights:
                    model.load_state_dict(weights['model'])
                else:
                    model.load_state_dict(weights)
            else:
                # Assume it's the model directly
                model = weights
                
            print("✓ Weights loaded successfully")
        except Exception as e:
            print(f"✗ Error loading weights: {e}")
            print("Please check your weights file format")
            return False
    else:
        print(f"✗ Weights file not found: {model_weights_path}")
        return False
    
    # 4. Register for auto classes
    print("Registering for auto classes...")
    config.register_for_auto_class()
    model.register_for_auto_class("AutoModel")
    
    # 5. Register pipeline
    print("Registering pipeline...")
    PIPELINE_REGISTRY.register_pipeline(
        "vine-video-understanding",
        pipeline_class=VinePipeline,
        pt_model=VineModel,
        type="multimodal",
    )
    
    # 6. Create pipeline instance
    print("Creating pipeline...")
    vine_pipeline = VinePipeline(model=model, tokenizer=None)
    
    try:
        # 7. Push configuration to hub
        print(f"Pushing configuration to {repo_name}...")
        config.push_to_hub(
            repo_name,
            commit_message=f"{commit_message} - config",
            private=private
        )
        print("✓ Configuration pushed successfully")
        
        # 8. Push model to hub
        print(f"Pushing model to {repo_name}...")
        model.push_to_hub(
            repo_name,
            commit_message=f"{commit_message} - model",
            private=private
        )
        print("✓ Model pushed successfully")
        
        # 9. Push pipeline to hub
        print(f"Pushing pipeline to {repo_name}...")
        vine_pipeline.push_to_hub(
            repo_name,
            commit_message=f"{commit_message} - pipeline",
            private=private
        )
        print("✓ Pipeline pushed successfully")
        
        print(f"\n🎉 Successfully pushed VINE model to: https://huggingface.co/{repo_name}")
        print(f"\nTo use your model:")
        print(f"```python")
        print(f"from transformers import pipeline")
        print(f"")
        print(f"vine_pipeline = pipeline(")
        print(f"    'vine-video-understanding',")
        print(f"    model='{repo_name}',")
        print(f"    trust_remote_code=True")
        print(f")")
        print(f"")
        print(f"results = vine_pipeline(")
        print(f"    'path/to/video.mp4',")
        print(f"    categorical_keywords=['human', 'dog', 'frisbee'],")
        print(f"    unary_keywords=['running', 'jumping'],")
        print(f"    binary_keywords=['chasing', 'behind']")
        print(f")")
        print(f"```")
        
        return True
        
    except Exception as e:
        print(f"✗ Error pushing to hub: {e}")
        print("Please check your HuggingFace credentials and repository permissions")
        return False


def main():
    parser = argparse.ArgumentParser(description="Push VINE model to HuggingFace Hub")
    
    parser.add_argument(
        "--weights", 
        type=str, 
        required=True,
        help="Path to the trained model weights (.pth file)"
    )
    
    parser.add_argument(
        "--repo", 
        type=str, 
        required=True,
        help="Repository name (e.g., 'username/vine-model')"
    )
    
    parser.add_argument(
        "--model-name",
        type=str,
        default="openai/clip-vit-base-patch32",
        help="CLIP model backbone name"
    )
    
    parser.add_argument(
        "--segmentation",
        type=str,
        default="grounding_dino_sam2",
        choices=["sam2", "grounding_dino_sam2"],
        help="Segmentation method"
    )
    
    parser.add_argument(
        "--message",
        type=str,
        default="Upload VINE model",
        help="Commit message"
    )
    
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create private repository"
    )
    
    parser.add_argument(
        "--login",
        action="store_true",
        help="Login to HuggingFace Hub first"
    )
    
    args = parser.parse_args()
    
    # Login if requested
    if args.login:
        print("Logging in to HuggingFace Hub...")
        notebook_login()
    
    # Push model
    success = push_vine_to_hub(
        model_weights_path=args.weights,
        repo_name=args.repo,
        model_name=args.model_name,
        segmentation_method=args.segmentation,
        commit_message=args.message,
        private=args.private
    )
    
    if success:
        print("\n✅ Model successfully pushed to HuggingFace Hub!")
    else:
        print("\n❌ Failed to push model to HuggingFace Hub")
        sys.exit(1)


if __name__ == "__main__":
    main()