#!/usr/bin/env python3
"""
HyperSynapse-SAS Chat Template Applier
Usage:
    python apply_template.py --template universal --local-dir ./path/to/model
    python apply_template.py --template deepseek --model-id Solstice-AI/my-model --push
"""

import argparse
import json
import os
import sys
from huggingface_hub import HfApi, hf_hub_download

TEMPLATES = {
    "universal": "templates/universal.jinja",
    "deepseek": "templates/deepseek.jinja",
    "qwen": "templates/qwen.jinja",
    "glm": "templates/glm.jinja",
    "llama3": "templates/llama3.jinja",
    "llama": "templates/llama3.jinja",
    "mistral": "templates/mistral.jinja",
    "gemma": "templates/gemma.jinja",
}

def load_template(name):
    base_dir = os.path.dirname(os.path.abspath(__file__))
    rel_path = TEMPLATES.get(name.lower())
    if not rel_path:
        raise ValueError(f"Unknown template: {name}. Available: {list(TEMPLATES.keys())}")
    full_path = os.path.join(base_dir, rel_path)
    with open(full_path, "r", encoding="utf-8") as f:
        return f.read()

def apply_local(local_dir, template_name):
    tc_path = os.path.join(local_dir, "tokenizer_config.json")
    if not os.path.exists(tc_path):
        print(f"Error: {tc_path} not found.")
        sys.exit(1)
    
    template_code = load_template(template_name)
    with open(tc_path, "r", encoding="utf-8") as f:
        tc = json.load(f)
    
    tc["chat_template"] = template_code
    with open(tc_path, "w", encoding="utf-8") as f:
        json.dump(tc, f, indent=2, ensure_ascii=False)
        f.write("\n")
    
    # Also save chat_template.jinja standalone
    jinja_path = os.path.join(local_dir, "chat_template.jinja")
    with open(jinja_path, "w", encoding="utf-8") as f:
        f.write(template_code)
    
    print(f" Successfully applied {template_name} template to {local_dir}")

def apply_hub(model_id, template_name, push=False, token=None):
    api = HfApi(token=token)
    print(f"Downloading tokenizer_config.json from {model_id}...")
    tc_path = hf_hub_download(repo_id=model_id, filename="tokenizer_config.json", token=token)
    with open(tc_path, "r", encoding="utf-8") as f:
        tc = json.load(f)
    
    template_code = load_template(template_name)
    tc["chat_template"] = template_code
    
    if push:
        print(f"Pushing updated tokenizer_config.json and chat_template.jinja to {model_id}...")
        api.upload_file(
            path_or_fileobj=json.dumps(tc, indent=2, ensure_ascii=False).encode("utf-8"),
            path_in_repo="tokenizer_config.json",
            repo_id=model_id,
            commit_message=f"feat: apply HyperSynapse-SAS {template_name} chat template"
        )
        api.upload_file(
            path_or_fileobj=template_code.encode("utf-8"),
            path_in_repo="chat_template.jinja",
            repo_id=model_id,
            commit_message=f"feat: add HyperSynapse-SAS {template_name} chat_template.jinja"
        )
        print(f" Successfully pushed to {model_id}!")
    else:
        print("Dry run completed. Pass --push to commit changes to Hugging Face Hub.")

def main():
    parser = argparse.ArgumentParser(description="Apply HyperSynapse-SAS Chat Templates")
    parser.add_argument("--template", "-t", required=True, choices=list(TEMPLATES.keys()), help="Template family to apply")
    parser.add_argument("--local-dir", "-l", help="Path to local model directory")
    parser.add_argument("--model-id", "-m", help="Hugging Face Model ID (e.g. org/model)")
    parser.add_argument("--push", action="store_true", help="Push to Hugging Face Hub")
    parser.add_argument("--token", help="Hugging Face User Access Token")
    
    args = parser.parse_args()
    if args.local_dir:
        apply_local(args.local_dir, args.template)
    elif args.model_id:
        apply_hub(args.model_id, args.template, push=args.push, token=args.token)
    else:
        print("Please provide either --local-dir or --model-id.")
        sys.exit(1)

if __name__ == "__main__":
    main()