# Use variables from training cell instead of hardcoded values
num_classes = model_info.get('num_classes', len(model_info.get('classes', [])))
classes_list = model_info.get('classes', [])
current_date = datetime.now().strftime("%Y-%m-%d")

model_card = f"""---

print("📤 Uploading enhanced model card to Hugging Face...")

api.upload_file(
    path_or_fileobj=enhanced_model_card.encode('utf-8'),
    path_in_repo="README.md",  # Exact case-sensitive filename
    repo_id=target_repo_id,
    commit_message="Overwrite README.md with comprehensive model card and training statistics",
    token=hf_token,
    create_pr=False  # Ensure direct commit, not PR
)

print("✅ Enhanced model card uploaded successfully!")
print(f"🌐 View your model at: https://huggingface.co/{target_repo_id}")

try:
    files = api.list_repo_files(repo_id=target_repo_id, token=hf_token)
    if "README.md" in files:
        print("✅ README.md confirmed on repository")
    else:
        print("⚠️ README.md not found in repository files")
        print(f"📋 Available files: {files}")
except Exception as verify_error:
    print(f"⚠️ Could not verify upload: {verify_error}")

# Get HF token and API (should be available from previous cells)
hf_token = os.getenv('HF_TOKEN')
if not hf_token:
    print("❌ HF_TOKEN not found. Please run the setup cells first.")
    raise ValueError("HF_TOKEN is required")

api = HfApi()
me = api.whoami(token=hf_token)

# Use repo_id from training cell if available, otherwise create it
if 'repo_id' in globals():
    target_repo_id = repo_id
    print(f"✅ Using repo_id from training cell: {target_repo_id}")
else:
    target_repo_id = f"{me['name']}/hair-classifier"
    print(f"⚠️ No repo_id found, using default: {target_repo_id}")

# Use model_info from training cell if available
if 'model_info' in globals():
    model_info_for_card = model_info
    print(f"✅ Using model_info from training cell: {len(model_info_for_card.get('classes', []))} classes")
elif 'data_module' in globals():
    model_info_for_card = {
        'repo_id': target_repo_id,
        'classes': data_module.classes if hasattr(data_module, 'classes') else [],
        'num_classes': len(data_module.classes) if hasattr(data_module, 'classes') else 0
    }
    print(f"✅ Created model_info from data_module: {len(model_info_for_card.get('classes', []))} classes")
else:
    print("⚠️ No model info found from training cell. Using defaults.")
    model_info_for_card = {
        'repo_id': target_repo_id,
        'classes': ['1A', '1B', '1C'],  # Fallback based on your training output
        'num_classes': 3
    }

# Extract training statistics from trainer if available
training_stats_for_card = {}
if 'trainer' in globals() and hasattr(trainer, 'logged_metrics'):
    metrics = trainer.logged_metrics
    training_stats_for_card = {
        'val_accuracy': f"{metrics.get('val_accuracy', 0):.1%}",
        'val_f1': f"{metrics.get('val_f1', 0):.1%}",
        'epochs_trained': trainer.current_epoch + 1 if hasattr(trainer, 'current_epoch') else 'N/A',
        'best_epoch': trainer.current_epoch + 1 if hasattr(trainer, 'current_epoch') else 'N/A'
    }
    print(f"✅ Using training stats from trainer: {training_stats_for_card}")
else:
    # Fallback to your actual training results
    training_stats_for_card = {
        'val_accuracy': '76.7%',
        'val_f1': '31.1%',
        'epochs_trained': 25,
        'best_epoch': 4
    }
    print(f"⚠️ Using fallback training stats: {training_stats_for_card}")

upload_model_card(target_repo_id, model_info_for_card, training_stats_for_card, hf_token)