"""Optional finite Kaggle continuation using a notebook-scoped Kaggle Secret.

Never embeds, writes or logs the credential. Enabled only by the launcher after
the user authorizes the Secret and its binding to this notebook.
"""
import json
import os
from pathlib import Path
import shutil
import time
from common import dump, sha, now

BACKUP_SLUG='faithdemo/think72-training-checkpoints'

def api_client():
    if not os.environ.get('KAGGLE_API_TOKEN'):
        from kaggle_secrets import UserSecretsClient
        os.environ['KAGGLE_API_TOKEN']=UserSecretsClient().get_secret('THINK72_KAGGLE_TOKEN')
    from kaggle.api.kaggle_api_extended import KaggleApi
    api=KaggleApi();api.authenticate()
    assert api.config_values.get('username')=='faithdemo'
    return api

def backup_checkpoint(out, checkpoint, force=False):
    if os.environ.get('THINK72_CLOUD_RESUME')!='1':return False
    out=Path(out); ledger=out/'cloud_backup_status.json'
    previous=json.loads(ledger.read_text()) if ledger.exists() else {}
    if not force and time.time()-previous.get('last_success_epoch',0)<7200:return True
    api=api_client()
    stage=out.parent/'cloud_checkpoint';stage.mkdir(exist_ok=True)
    target=stage/'resume.pt'
    if target.exists():
        assert target.resolve().parent==stage.resolve()
        target.unlink()
    os.link(checkpoint,target)
    step=json.loads((out/'checkpoints/latest.json').read_text())['step']
    state={'step':step,'checkpoint_sha256':sha(target),'bytes':target.stat().st_size,
           'session_index':int(os.environ.get('THINK72_SESSION_INDEX','1')),'created_utc':now()}
    dump(stage/'cloud_state.json',state)
    dump(stage/'dataset-metadata.json',{
        'id':BACKUP_SLUG,'title':'Think72 Training Checkpoints','licenses':[{'name':'other'}],
        'description':'Private resumable 72M Qwen3-style model training state: model weights, optimizer, scaler, RNG and data cursor. Related token snapshot: https://www.kaggle.com/datasets/faithdemo/think72-english-tokens-v1 . Original project contributions Copyright (c) 2026 sraivante, Apache-2.0; training data retains third-party terms including CC-BY-NC-4.0. No blanket commercial license is asserted.'})
    exists=True
    try:api.dataset_status(BACKUP_SLUG)
    except Exception as exc:
        response=getattr(exc,'response',None)
        if response is not None and response.status_code in (403,404):
            # Kaggle returns 403 for absent private dataset slugs. Confirm the
            # authenticated owner's matching datasets before attempting create.
            owned=api.dataset_list(mine=True,search='Think72 Training Checkpoints') or []
            if any(getattr(item,'ref',None)==BACKUP_SLUG for item in owned):raise
            exists=False
        else:raise
    if exists:
        result=api.dataset_create_version(str(stage),f'Checkpoint step {step}',quiet=True,
                                         convert_to_csv=False,delete_old_versions=False)
    else:result=api.dataset_create_new(str(stage),public=False,quiet=True,convert_to_csv=False)
    error=getattr(result,'error',None)
    assert not error,'Checkpoint dataset upload did not succeed'
    for _ in range(120):
        status=api.dataset_status(BACKUP_SLUG)
        if status=='ready':break
        if status in ('error','failed'):raise RuntimeError('Checkpoint dataset processing failed')
        time.sleep(5)
    else:raise TimeoutError('Checkpoint dataset not ready')
    dump(ledger,{**state,'dataset':BACKUP_SLUG,'last_success_epoch':time.time(),'status':'READY'})
    print('CLOUD_CHECKPOINT_READY',step,flush=True)
    return True

def continue_if_needed(root, source_script):
    root=Path(root);out=root/'run'
    status=json.loads((out/'run_status.json').read_text())
    latest=json.loads((out/'checkpoints/latest.json').read_text())
    backup_checkpoint(out,out/'checkpoints'/latest['path'],force=True)
    if status['status']=='COMPLETE':return
    assert status['status']=='CHECKPOINTED_NEEDS_RESUME'
    index=int(os.environ.get('THINK72_SESSION_INDEX','1'))
    assert index<4,'Automatic continuation limit reached; inspect progress before further sessions'
    next_dir=root/'next_session';next_dir.mkdir(exist_ok=True)
    shutil.copyfile(source_script,next_dir/'run.py')
    dump(next_dir/'kernel-metadata.json',{
        'id':'faithdemo/think72-qwen3-t4x2','title':'Think72 Qwen3 T4x2',
        'code_file':'run.py','language':'python','kernel_type':'script','is_private':True,
        'enable_gpu':True,'enable_internet':True,'machine_shape':'NvidiaTeslaT4',
        'dataset_sources':['faithdemo/think72-english-tokens-v1',BACKUP_SLUG],
        'kernel_sources':[],'competition_sources':[],'model_sources':[],
        'docker_image_pinning_type':'latest'})
    api=api_client()
    result=api.kernels_push(str(next_dir),timeout='43200',acc='NvidiaTeslaT4')
    assert not getattr(result,'error',None),'Continuation request failed'
    report={'from_step':status['step'],'next_session_index':index+1,
            'requested_accelerator':'NvidiaTeslaT4','url':getattr(result,'url',None),
            'submitted_utc':now(),'target_tokens':1_500_000_000}
    dump(root/'continuation_submitted.json',report)
    print('NEXT_T4X2_SESSION_SUBMITTED',json.dumps(report),flush=True)