#!/usr/bin/env python3
"""EMT Compute REST SDK and CLI. Python 3.9+, no third-party dependencies."""
import argparse
import hashlib
import json
import os
import sys
import uuid
import shlex
import re
import subprocess
from pathlib import Path
from urllib.request import Request, urlopen
from urllib.error import HTTPError
from urllib.parse import urlencode, urlparse, quote

class EMTCompute:
    def __init__(self, base_url=None, api_key=None):
        self.base_url=(base_url or os.environ.get('EMT_BASE_URL','https://emtcompute.com')).rstrip('/')
        parsed=urlparse(self.base_url)
        if parsed.username or parsed.password or parsed.query or parsed.fragment or (parsed.scheme!='https' and not(parsed.scheme=='http' and parsed.hostname in ('127.0.0.1','localhost'))):
            raise ValueError('Use HTTPS, except for local development')
        self.api_key=api_key or os.environ.get('EMT_API_KEY')
    def request(self,path,method='GET',payload=None,idempotency_key=None,raw=None,extra_headers=None):
        headers={'Accept':'application/json','User-Agent':'EMT-Compute-Python/0.1'}
        if self.api_key: headers['Authorization']='Bearer '+self.api_key
        if idempotency_key: headers['Idempotency-Key']=idempotency_key
        data=None
        if payload is not None:
            data=json.dumps(payload).encode();headers['Content-Type']='application/json'
        if raw is not None:
            data=raw;headers['Content-Type']='application/octet-stream'
        headers.update(extra_headers or {})
        # Reject cross-origin redirects rather than forwarding a bearer token.
        from urllib.request import HTTPRedirectHandler,build_opener
        class NoRedirect(HTTPRedirectHandler):
            def redirect_request(self,*args,**kwargs): return None
        try:
            with build_opener(NoRedirect).open(Request(self.base_url+'/api/v1/'+path,data=data,headers=headers,method=method),timeout=35) as response:
                return json.load(response)
        except HTTPError as error:
            raw=error.read().decode(errors='replace')
            try: message=json.loads(raw).get('message',str(error))
            except ValueError: message='HTTP %s. Check site access and API credentials.'%error.code
            raise RuntimeError(message) from None
    def catalog(self,**filters):return self.request('catalog?'+urlencode(filters))
    def clusters(self):return self.request('catalog/clusters')
    def cluster_quote(self,cluster,nodes,hours=1):return self.request('cluster-quotes','POST',dict(cluster=cluster,nodes=nodes,hours=hours))
    def rent_cluster(self,quote_id,name,ssh_public_key,idempotency_key):
        if not idempotency_key:raise ValueError('Persist and reuse the idempotency key for this order')
        return self.request('cluster-orders','POST',dict(quote=quote_id,name=name,ssh=ssh_public_key),idempotency_key)
    def cluster_member(self,order_id,pod,action='metrics'):
        if action not in ('logs','metrics','reboot'):raise ValueError('Unsupported member operation')
        return self.request('orders/'+quote(order_id,safe='')+'/members/'+quote(pod,safe='')+'/'+action,'POST' if action=='reboot' else 'GET',{} if action=='reboot' else None)
    def templates(self):return self.request('templates')
    def quote(self,node,count=1,hours=1):return self.request('quotes','POST',dict(node=node,count=count,hours=hours))
    def rent(self,quote_id,name,template,ssh_public_key,idempotency_key,jupyter=False,volume=None,backups=False,ports=10):
        if not idempotency_key:raise ValueError('Persist and reuse an idempotency key for this logical order')
        return self.request('orders','POST',dict(quote=quote_id,name=name,template=template,ssh=ssh_public_key,jupyter=jupyter,backups=backups,ports=ports,**({"volume":volume} if volume else {})),idempotency_key)
    def pods(self):return self.request('orders')
    def refresh(self,pod):return self.request('orders/'+quote(pod,safe='')+'/refresh','POST',{})
    def remove(self,pod):return self.request('orders/'+quote(pod,safe=''),'DELETE')
    def ledger(self):return self.request('ledger')

    def training_jobs(self):return self.request('training/jobs')
    def training_model(self,model):return self.request('training/models/resolve','POST',{'model':model})
    def training_upload(self,path,seed=42,eval_percent=10):
        path=Path(path);size=path.stat().st_size
        if path.suffix.lower()!='.jsonl' or not 0<size<=128*1024*1024:raise ValueError('Use a nonempty JSONL file up to 128 MiB')
        digest=hashlib.sha256()
        with path.open('rb') as f:
            for chunk in iter(lambda:f.read(8*1024*1024),b''):digest.update(chunk)
        item=self.request('training/files','POST',{'name':path.name,'size':size,'sha256':digest.hexdigest()})
        if item['state']=='uploading':
            known={part['part']:part['sha256'] for part in item['parts']}
            with path.open('rb') as f:
                part=0
                for chunk in iter(lambda:f.read(item['chunkBytes']),b''):
                    sha=hashlib.sha256(chunk).hexdigest()
                    if known.get(part)!=sha:self.request('training/files/'+quote(item['id'],safe='')+'/parts/'+str(part),'PUT',raw=chunk,extra_headers={'X-Content-SHA256':sha})
                    part+=1
        report=self.request('training/files/'+quote(item['id'],safe='')+'/complete','POST',{'seed':seed,'evalPercent':eval_percent})
        return {'id':item['id'],'report':report}
    def training_prepare(self,name,dataset,config,idempotency_key):
        if not idempotency_key:raise ValueError('Persist an idempotency key before preparing a job')
        return self.request('training/jobs','POST',{'name':name,'dataset':dataset,'config':config},idempotency_key)
    def training_job(self,job,action=None,payload=None):
        if action not in (None,'events','recipe','stop','resume','start'):raise ValueError('Unsupported training action')
        method='POST' if action in ('stop','resume','start') else 'GET'
        return self.request('training/jobs/'+quote(job,safe='')+('/'+action if action else ''),method,(payload or {}) if method=='POST' else None)

    def monitor(self,order_id,member=None):
        return self.request('orders/'+quote(order_id,safe='')+'/monitor'+(('?'+urlencode({'member':member})) if member else ''))
    def volumes(self):return self.request('volumes')
    def create_volume(self,name,description='',idempotency_key=None):
        if not idempotency_key:raise ValueError('Persist an idempotency key before creating a volume')
        return self.request('volumes','POST',dict(name=name,description=description),idempotency_key)
    def backups(self,order_id,member=None,page=1):return self.request(self.pod_path(order_id,'backups',member)+('&' if member else '?')+'page='+str(page))
    def pod_path(self,order_id,action,member=None):return 'orders/'+quote(order_id,safe='')+'/'+action+(('?'+urlencode({'member':member})) if member else '')
    def backup(self,order_id,name,idempotency_key,member=None):
        if not idempotency_key:raise ValueError('Persist an idempotency key before creating a backup')
        return self.request(self.pod_path(order_id,'backups',member),'POST',dict(name=name,description=''),idempotency_key)
    def configure_backup(self,order_id,path='/root',frequency=6,retention=2,member=None):return self.request(self.pod_path(order_id,'backup-config',member),'POST',dict(path=path,frequency=frequency,retention=retention))
    def safe_stop(self,order_id,path,frequency=6,retention=2,checkpoint_closed=False):
        if not checkpoint_closed:raise ValueError('Flush and close your application checkpoints before safe stop')
        return self.request(self.pod_path(order_id,'safe-stop'),'POST',dict(path=path,frequency=frequency,retention=retention,checkpointClosed=True))
    def connection(self,order_id,member=None):
        info=self.monitor(order_id,member);members=info.get('members',[])
        if len(members)!=1:raise ValueError('Select one cluster member with --member')
        command=members[0].get('ssh_connect_cmd')
        if not command:raise ValueError('SSH connection is not ready')
        return parse_ssh_connection(command)
    def transfer(self,order_id,local,remote,download=False,member=None,identity=None):
        if not re.fullmatch(r'/(?:mnt|root)(?:/[A-Za-z0-9_.-]+)*',remote) or any(p in ('.','..') for p in remote.split('/')):
            raise ValueError('Use an absolute path under /mnt or /root without shell characters')
        target,port=self.connection(order_id,member)
        local_path=Path(local).expanduser().resolve()
        if download and local_path.exists():raise ValueError('Choose a new local destination; existing files are never overwritten')
        if not download and not local_path.exists():raise ValueError('Local source does not exist')
        args=['scp','-r','-P',str(port),'-o','StrictHostKeyChecking=ask']
        if identity:args+=['-i',str(Path(identity).expanduser().resolve())]
        remote_spec=target+':'+remote
        args+=['--']+([remote_spec,str(local_path)] if download else [str(local_path),remote_spec])
        subprocess.run(args,check=True)
        return {'transferred':True,'direction':'download' if download else 'upload','remote':remote,'local':str(local_path)}

def parse_ssh_connection(command):
    """Parse data, never execute a provider-supplied shell command or SSH options."""
    parts=shlex.split(command)
    if not parts or parts.pop(0)!='ssh':raise ValueError('Unsupported SSH connection format')
    target=None;port=22
    while parts:
        value=parts.pop(0)
        if value=='-p' and parts:port=int(parts.pop(0))
        elif re.fullmatch(r'[A-Za-z0-9_-]+@(?:[A-Za-z0-9.-]+|\[[A-Fa-f0-9:]+\])',value) and target is None:target=value
        else:raise ValueError('Unexpected SSH options; inspect the connection in the EMT console')
    if not target or not 1<=port<=65535:raise ValueError('Invalid SSH host or port')
    return target,port

def main():
    parser=argparse.ArgumentParser(description='EMT Compute CLI');sub=parser.add_subparsers(dest='command',required=True)
    ls=sub.add_parser('ls');ls.add_argument('--cuda');ls.add_argument('--nvlink',action='store_true');ls.add_argument('--tier',choices=['secure','spot'])
    for cmd in ['ps','templates','balance','ledger','clusters','volumes','storage-costs','registries','custom-templates']:sub.add_parser(cmd)
    q=sub.add_parser('quote');q.add_argument('node');q.add_argument('--gpus',type=int,required=True);q.add_argument('--hours',type=int,default=1)
    rent=sub.add_parser('rent');rent.add_argument('--quote',required=True);rent.add_argument('--name',required=True);rent.add_argument('--template',required=True);rent.add_argument('--ssh-public-key',required=True);rent.add_argument('--idempotency-key',required=True);rent.add_argument('--jupyter',action='store_true');rent.add_argument('--volume');rent.add_argument('--backups',action='store_true');rent.add_argument('--ports',type=int,default=10)
    for cmd in ['refresh','logs','settle','rm']:
        p=sub.add_parser(cmd);p.add_argument('id')
        if cmd=='rm':p.add_argument('--confirm',action='store_true',required=True)
    cq=sub.add_parser('cluster-quote');cq.add_argument('cluster');cq.add_argument('--nodes',nargs='+',required=True);cq.add_argument('--hours',type=int,default=1)
    cr=sub.add_parser('cluster-rent');cr.add_argument('--quote',required=True);cr.add_argument('--name',required=True);cr.add_argument('--ssh-public-key',required=True);cr.add_argument('--idempotency-key',required=True)
    cm=sub.add_parser('cluster-member');cm.add_argument('id');cm.add_argument('pod');cm.add_argument('action',choices=['logs','metrics','reboot'])
    for cmd in ['monitor','backups']:
        p=sub.add_parser(cmd);p.add_argument('id');p.add_argument('--member')
    v=sub.add_parser('volume-create');v.add_argument('name');v.add_argument('--description',default='');v.add_argument('--idempotency-key',required=True)
    for cmd in ['upload','download']:
        p=sub.add_parser(cmd);p.add_argument('id');p.add_argument('source');p.add_argument('--member');p.add_argument('--identity')
        p.add_argument('--local' if cmd=='download' else '--remote',required=True)
    for cmd in ['backup-config','backup-now','safe-stop']:
        p=sub.add_parser(cmd);p.add_argument('id');p.add_argument('--member');p.add_argument('--path',default='/root/checkpoints');p.add_argument('--frequency',type=int,default=6);p.add_argument('--retention',type=int,default=2)
        if cmd=='backup-now':
            p.add_argument('--name',required=True);p.add_argument('--idempotency-key',required=True)
        if cmd=='safe-stop':p.add_argument('--checkpoint-closed',action='store_true',required=True)
    sub.add_parser('training-jobs')
    tm=sub.add_parser('training-model');tm.add_argument('model')
    tu=sub.add_parser('training-upload');tu.add_argument('file')
    tp=sub.add_parser('training-prepare');tp.add_argument('config');tp.add_argument('--dataset',required=True);tp.add_argument('--name',required=True);tp.add_argument('--idempotency-key',required=True)
    for cmd in ['training-job','training-events','training-recipe','training-stop','training-resume','training-start']:
        p=sub.add_parser(cmd);p.add_argument('id')
        if cmd=='training-start':p.add_argument('--request',required=True,help='JSON with quote, volume, SSH public key and storage charge acknowledgment')
    args=parser.parse_args();client=EMTCompute()
    try:
        if args.command=='training-jobs':result=client.training_jobs()
        elif args.command=='training-model':result=client.training_model(args.model)
        elif args.command=='training-upload':result=client.training_upload(args.file)
        elif args.command=='training-prepare':result=client.training_prepare(args.name,args.dataset,json.loads(Path(args.config).read_text()),args.idempotency_key)
        elif args.command.startswith('training-'):result=client.training_job(args.id,None if args.command=='training-job' else args.command.removeprefix('training-'),json.loads(Path(args.request).read_text()) if args.command=='training-start' else None)
        elif args.command in ('upload','download'):
            result=client.transfer(args.id,args.local if args.command=='download' else args.source,args.source if args.command=='download' else args.remote,args.command=='download',args.member,args.identity)
        elif args.command=='volume-create':result=client.create_volume(args.name,args.description,args.idempotency_key)
        elif args.command=='monitor':result=client.monitor(args.id,args.member)
        elif args.command=='backups':result=client.backups(args.id,args.member)
        elif args.command=='backup-now':result=client.backup(args.id,args.name,args.idempotency_key,args.member)
        elif args.command=='backup-config':result=client.configure_backup(args.id,args.path,args.frequency,args.retention,args.member)
        elif args.command=='safe-stop':result=client.safe_stop(args.id,args.path,args.frequency,args.retention,args.checkpoint_closed)
        elif args.command=='ls':result=client.catalog(**{k:v for k,v in {'cuda':args.cuda,'nvlink':'1' if args.nvlink else None,'tier':args.tier}.items() if v})
        elif args.command=='clusters':result=client.clusters()
        elif args.command=='cluster-quote':result=client.cluster_quote(args.cluster,args.nodes,args.hours)
        elif args.command=='cluster-member':result=client.cluster_member(args.id,args.pod,args.action)
        elif args.command=='cluster-rent':
            with open(args.ssh_public_key,encoding='utf-8') as f:pub=f.read().strip()
            if 'PRIVATE KEY' in pub:raise ValueError('Use a .pub public key file')
            result=client.rent_cluster(args.quote,args.name,pub,args.idempotency_key)
        elif args.command=='quote':result=client.quote(args.node,args.gpus,args.hours)
        elif args.command=='rent':
            with open(args.ssh_public_key,encoding='utf-8') as f:pub=f.read().strip()
            if 'PRIVATE KEY' in pub:raise ValueError('Use a .pub public key file')
            result=client.rent(args.quote,args.name,args.template,pub,args.idempotency_key,args.jupyter,args.volume,args.backups,args.ports)
        elif args.command=='rm':result=client.remove(args.id)
        elif args.command=='refresh':result=client.refresh(args.id)
        elif args.command in ('logs','settle'):result=client.request('orders/'+quote(args.id,safe='')+'/'+args.command,'POST' if args.command=='settle' else 'GET',{} if args.command=='settle' else None)
        else:result=client.request({'ps':'orders','templates':'templates','balance':'me','ledger':'ledger','volumes':'volumes','storage-costs':'storage/costs','registries':'registries','custom-templates':'custom-templates'}[args.command])
        print(json.dumps(result,indent=2,ensure_ascii=False))
    except Exception as e:print('EMT:',str(e),file=sys.stderr);sys.exit(1)
if __name__=='__main__':main()