import os
import sys
import logging
from huggingface_hub import HfApi, create_repo, upload_folder
from getpass import getpass

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

# --- Configuration ---
# The model weights and adapters to upload.
# BASE_MODEL_DIR should contain the merged INT8 checkpoint.
# ADAPTERS_DIR should contain the persona adapter subdirectories.
BASE_MODEL_DIR = "artifacts/deploy/ankahi-gemma4-e4b-int8"  # This is the INT8 quantized model
ADAPTERS_DIR = "artifacts/stage2"
# Note: The original BF16 model in punjabi_gemma/ankahi is ~16GB, the INT8 is ~11GB.
# We upload the INT8 version.

# Hugging Face repository details
# Choose a meaningful repository name, e.g., your username/org + model name
REPO_ID = "bhriguverma/ankahi-gemma-aac-int8"  # Example: use your HF username or organization here

# --- Script Logic ---

def upload_model_to_hf(token: str):
    """Uploads the model and adapters to Hugging Face."""
    api = HfApi()

    try:
        log.info(f"Attempting to create or connect to repository: {REPO_ID}")
        # Create repository if it doesn't exist
        create_repo(REPO_ID, private=True, token=token, exist_ok=True)
        log.info(f"Repository '{REPO_ID}' ready.")
    except Exception as e:
        log.error(f"Failed to connect to Hugging Face Hub or create repository. Error: {e}")
        sys.exit(1)

    repo_url = f"https://huggingface.co/{REPO_ID}"
    log.info(f"Uploading model components to: {repo_url}")

    # Upload the INT8 base model
    try:
        log.info(f"Uploading base model from: {BASE_MODEL_DIR}")
        upload_folder(
            repo_id=REPO_ID,
            folder_path=BASE_MODEL_DIR,
            token=token,
            commit_message="Upload INT8 quantized base Gemma 4 model for Ankahi",
            repo_type="model",
            commit_description="Upload INT8 quantized base Gemma 4 model for Ankahi",
        )
        log.info(f"Successfully uploaded base model from {BASE_MODEL_DIR}")
    except Exception as e:
        log.error(f"Failed to upload base model from {BASE_MODEL_DIR}: {e}")
        sys.exit(1)

    # Upload the persona adapters
    try:
        log.info(f"Uploading persona adapters from: {ADAPTERS_DIR}")
        upload_folder(
            repo_id=REPO_ID,
            folder_path=ADAPTERS_DIR,
            token=token,
            commit_message="Upload persona adapters for Ankahi",
            repo_type="model",
            commit_description="Upload persona adapters for Ankahi",
        )
        log.info(f"Successfully uploaded persona adapters from {ADAPTERS_DIR}")
    except Exception as e:
        log.error(f"Failed to upload persona adapters from {ADAPTERS_DIR}: {e}")
        sys.exit(1)
        
    # Upload other relevant files like config, requirements, etc.
    try:
        log.info("Uploading supplementary files (config, requirements, etc.)...")
        supplementary_files = [
            "README.md",
            "ankahi_backend/requirements.txt",
            "ankahi_backend/config.py",
            "RESULTS.md",
            "ANKAHI_DEVELOPER_HANDOVER.md"
        ]
        
        for file_path in supplementary_files:
            if os.path.exists(file_path):
                log.info(f"Uploading {file_path}...")
                if os.path.dirname(file_path):
                    folder_path = os.path.dirname(file_path)
                    path_in_repo = os.path.basename(file_path)
                else:
                    folder_path = '.'
                    path_in_repo = os.path.basename(file_path)

                upload_folder(
                    repo_id=REPO_ID,
                    folder_path=folder_path,
                    path_in_repo=path_in_repo,
                    token=token,
                    commit_message=f"Upload {os.path.basename(file_path)}",
                    repo_type="model",
                    commit_description=f"Upload {os.path.basename(file_path)}",
                    discard_local_dir=True
                )
                log.info(f"Uploaded {file_path}")
            else:
                log.warning(f"Supplementary file not found, skipping: {file_path}")
        log.info("Supplementary files upload complete.")
    except Exception as e:
        log.error(f"Failed to upload supplementary files: {e}")

    log.info("\n--- Upload complete ---")
    log.info(f"Model and adapters uploaded to: {repo_url}")
    log.info("Please check your Hugging Face repository for the uploaded files.")
    log.info("Remember to set your repository to public if intended.")

if __name__ == "__main__":
    # WARNING: Hardcoded token as requested. Do not share this file publicly!
    hf_token = "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"
    
    upload_model_to_hf(hf_token)