"""
Script to push VINE model to video-fm organization on HuggingFace Hub

This script pushes the VINE architecture (config, model, pipeline) and model weights
to the video-fm organization for easy sharing and distribution.
"""

import os
import sys
import torch
import argparse
from pathlib import Path
from huggingface_hub import HfApi, login
from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModel
from safetensors.torch import save_file

# Add src/ to sys.path so LASER, video-sam2, GroundingDINO are importable
current_dir = Path(__file__).resolve().parent
src_dir = current_dir.parent / "src"
if src_dir.is_dir() and str(src_dir) not in sys.path:
    sys.path.insert(0, str(src_dir))

os.environ['OPENAI_API_KEY'] = "dummy-key"

# Import from vine_hf package
from vine_hf import VineConfig, VineModel, VinePipeline


def push_vine_to_video_fm(
    source_repo_or_path: str = "KevinX-Penn28/testing",
    target_repo: str = "video-fm/vine",
    model_name: str = "openai/clip-vit-base-patch32",
    commit_message: str = "Upload VINE model architecture and weights",
    private: bool = False,
    use_local_weights: bool = False,
):
    """
    Push VINE model to video-fm organization on HuggingFace Hub.

    Args:
        source_repo_or_path: Source HF repo or local path with model weights
        target_repo: Target repository (e.g., "video-fm/vine")
        model_name: CLIP model backbone name
        commit_message: Commit message for the push
        private: Whether to create a private repository
        use_local_weights: If True, source_repo_or_path is a local file path
    """

    print("=" * 70)
    print("🚀 Pushing VINE Model to HuggingFace Hub - video-fm Organization")
    print("=" * 70)

    # 1. Create configuration
    print(f"\n📝 Creating configuration with backbone: {model_name}")
    config = VineConfig(
        model_name=model_name,
        segmentation_method="grounding_dino_sam2",
        use_hf_repo=not use_local_weights,
        model_repo=source_repo_or_path if not use_local_weights else None,
        local_dir=str(Path(source_repo_or_path).parent) if use_local_weights else None,
        local_filename=Path(source_repo_or_path).name if use_local_weights else None,
    )

    # 2. Initialize model (will automatically load weights from source)
    print(f"\n🔧 Initializing model and loading weights from: {source_repo_or_path}")
    model = VineModel(config)
    print("✓ Model initialized with weights loaded")

    # 3. Register for auto classes
    print("\n📋 Registering for auto classes...")
    config.register_for_auto_class()
    model.register_for_auto_class("AutoModel")
    print("✓ Registered for AutoModel and AutoConfig")

    # 4. Register pipeline
    print("\n🔌 Registering custom pipeline...")
    try:
        PIPELINE_REGISTRY.register_pipeline(
            "vine-video-understanding",
            pipeline_class=VinePipeline,
            pt_model=VineModel,
            type="multimodal",
        )
        print("✓ Pipeline registered")
    except Exception as e:
        print(f"⚠ Pipeline registration: {e} (may already be registered)")

    try:
        # 5. Push configuration to hub
        print(f"\n⬆️  Pushing configuration to {target_repo}...")
        config.push_to_hub(
            target_repo,
            commit_message=f"{commit_message} - config",
            private=private
        )
        print("✓ Configuration pushed successfully")

        # 6. Push model to hub
        print(f"\n⬆️  Pushing model to {target_repo}...")
        model.push_to_hub(
            target_repo,
            commit_message=f"{commit_message} - model and weights",
            private=private
        )
        print("✓ Model and weights pushed successfully")

        # 7. Copy additional necessary files to the repo
        print(f"\n📦 Uploading additional architecture files...")
        api = HfApi()

        # Upload flattening.py and vis_utils.py as they're imported by the model
        current_dir = Path(__file__).parent
        additional_files = [
            "flattening.py",
            "vis_utils.py",
        ]

        for filename in additional_files:
            file_path = current_dir / filename
            if file_path.exists():
                api.upload_file(
                    path_or_fileobj=str(file_path),
                    path_in_repo=filename,
                    repo_id=target_repo,
                    commit_message=f"Add {filename}",
                )
                print(f"✓ Uploaded {filename}")
            else:
                print(f"⚠ Warning: {filename} not found at {file_path}")

        # 8. Upload README if it exists
        readme_path = current_dir / "README.md"
        if readme_path.exists():
            api.upload_file(
                path_or_fileobj=str(readme_path),
                path_in_repo="README.md",
                repo_id=target_repo,
                commit_message="Add README documentation",
            )
            print("✓ Uploaded README.md")

        print("\n" + "=" * 70)
        print("🎉 Successfully pushed VINE model to HuggingFace Hub!")
        print("=" * 70)
        print(f"\n📍 Model URL: https://huggingface.co/{target_repo}")
        print(f"\n📚 To use your model:")
        print(f"""
```python
from transformers import AutoModel, AutoConfig
from vine_hf import VineConfig, VineModel, VinePipeline

# Option 1: Load with AutoModel
model = AutoModel.from_pretrained('{target_repo}', trust_remote_code=True)

# Option 2: Load with VineModel directly
config = VineConfig.from_pretrained('{target_repo}')
model = VineModel.from_pretrained('{target_repo}')

# Option 3: Use with pipeline
from transformers import pipeline

vine_pipeline = pipeline(
    'vine-video-understanding',
    model='{target_repo}',
    trust_remote_code=True
)

results = vine_pipeline(
    'path/to/video.mp4',
    categorical_keywords=['human', 'dog', 'frisbee'],
    unary_keywords=['running', 'jumping'],
    binary_keywords=['chasing', 'behind']
)
```
""")

        return True

    except Exception as e:
        print(f"\n❌ Error pushing to hub: {e}")
        import traceback
        traceback.print_exc()
        print("\nPlease check:")
        print("  - HuggingFace credentials (run: huggingface-cli login)")
        print("  - Repository permissions for video-fm organization")
        print("  - Network connectivity")
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Push VINE model to video-fm organization on HuggingFace Hub"
    )

    parser.add_argument(
        "--source",
        type=str,
        default="KevinX-Penn28/testing",
        help="Source HF repo or local path with model weights (default: KevinX-Penn28/testing)"
    )

    parser.add_argument(
        "--target",
        type=str,
        default="video-fm/vine",
        help="Target repository in video-fm org (default: video-fm/vine)"
    )

    parser.add_argument(
        "--model-name",
        type=str,
        default="openai/clip-vit-base-patch32",
        help="CLIP model backbone name"
    )

    parser.add_argument(
        "--message",
        type=str,
        default="Upload VINE model architecture and weights",
        help="Commit message"
    )

    parser.add_argument(
        "--private",
        action="store_true",
        help="Create private repository"
    )

    parser.add_argument(
        "--local-weights",
        action="store_true",
        help="Use local weights file instead of HF repo"
    )

    args = parser.parse_args()

    # Check login status
    try:
        api = HfApi()
        user_info = api.whoami()
        print(f"✓ Logged in as: {user_info['name']}")

        # Check if user has access to video-fm org
        orgs = [org['name'] for org in user_info.get('orgs', [])]
        if 'video-fm' in orgs:
            print(f"✓ Confirmed access to video-fm organization")
        else:
            print(f"⚠ Warning: You may not have access to video-fm organization")
            print(f"  Your organizations: {orgs}")
    except Exception as e:
        print(f"❌ Not logged in to HuggingFace. Please run: huggingface-cli login")
        print(f"   Or use: python -c 'from huggingface_hub import login; login()'")
        sys.exit(1)

    # Push model
    success = push_vine_to_video_fm(
        source_repo_or_path=args.source,
        target_repo=args.target,
        model_name=args.model_name,
        commit_message=args.message,
        private=args.private,
        use_local_weights=args.local_weights,
    )

    if success:
        print("\n✅ Successfully completed!")
        sys.exit(0)
    else:
        print("\n❌ Push failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()