#!/usr/bin/env python3
"""MathHub CLI. Standard library only; one account session, no API tokens."""
import argparse
import getpass
import hashlib
import http.cookiejar
import io
import json
import os
import stat
import sys
import urllib.error
import urllib.parse
import urllib.request
import uuid
import zipfile
from pathlib import Path, PurePosixPath


def emit(obj):
    print(json.dumps(obj, ensure_ascii=False, indent=2))


class Client:
    def __init__(self, base):
        self.base = base.rstrip('/')
        parsed = urllib.parse.urlsplit(self.base)
        if parsed.scheme not in {'http', 'https'} or not parsed.hostname or parsed.username or parsed.password:
            raise ValueError('Use an http(s) MathHub origin without credentials.')
        root = Path(os.environ.get('MATHHUB_CONFIG', str(Path.home() / '.config/mathhub')))
        root.mkdir(parents=True, exist_ok=True, mode=0o700)
        root.chmod(0o700)
        self.path = root / (hashlib.sha256(self.base.encode()).hexdigest()[:20] + '.cookies')
        self.jar = http.cookiejar.LWPCookieJar(str(self.path))
        if self.path.exists():
            self.jar.load(ignore_discard=True, ignore_expires=True)
        self.opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.jar))
        self.csrf = None

    def request(self, path, method='GET', data=None, content_type=None, raw=False):
        headers = {'User-Agent': 'MathHub-CLI/0.1', 'Accept': 'application/json', 'Origin': self.base}
        if method != 'GET':
            if not self.csrf:
                self.csrf = self.request('/api/session')['csrf']
            headers['X-CSRF-Token'] = self.csrf
            if content_type:
                headers['Content-Type'] = content_type
            else:
                headers['Content-Type'] = 'application/json'
                data = json.dumps(data or {}).encode()
        req = urllib.request.Request(self.base + path, method=method, data=data, headers=headers)
        try:
            with self.opener.open(req, timeout=90) as res:
                body = res.read(36 * 1024 * 1024 + 1)
                if len(body) > 36 * 1024 * 1024:
                    raise ValueError('Response exceeds CLI size limit.')
        except urllib.error.HTTPError as e:
            try:
                message = json.loads(e.read(10000))['error']
            except Exception:
                message = 'HTTP ' + str(e.code)
            raise ValueError(message) from e
        self.jar.save(ignore_discard=True, ignore_expires=True); self.path.chmod(0o600)
        if raw:
            return body
        out = json.loads(body)
        if isinstance(out, dict) and out.get('csrf'):
            self.csrf = out['csrf']
        return out


def safe_path(n):
    p = PurePosixPath(n)
    if not n or p.is_absolute() or '\\' in n or ':' in n or any(x in ('', '.', '..') for x in n.split('/')):
        raise ValueError('Unsafe package path: ' + n)
    return n


def package(source):
    p = Path(source)
    if p.is_symlink():
        raise ValueError('Symlink packages are not supported.')
    if p.is_file():
        if p.stat().st_size > 32 * 1024 * 1024:
            raise ValueError('ZIP exceeds 32 MiB.')
        return p.read_bytes()
    if not p.is_dir():
        raise ValueError('Expected a ZIP file or prepared work directory.')
    files = []
    for x in p.rglob('*'):
        if x.is_symlink():
            raise ValueError('Package contains a symlink: ' + str(x))
        if x.is_file():
            if any(part.startswith('.') for part in x.relative_to(p).parts):
                continue
            files.append(x)
    if len(files) > 200 or sum(x.stat().st_size for x in files) > 32 * 1024 * 1024:
        raise ValueError('Package exceeds limits.')
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as z:
        for x in files:
            z.writestr(safe_path(x.relative_to(p).as_posix()), x.read_bytes())
    return buf.getvalue()


def multipart(source, fields):
    boundary = 'MathHub-' + uuid.uuid4().hex
    parts = []
    for k, v in fields.items():
        parts += [f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode()]
    parts += [f'--{boundary}\r\nContent-Disposition: form-data; name="bundle"; filename="work.zip"\r\nContent-Type: application/zip\r\n\r\n'.encode(), package(source), f'\r\n--{boundary}--\r\n'.encode()]
    return b''.join(parts), 'multipart/form-data; boundary=' + boundary


def pull(client, wid, version, output):
    m = client.request('/api/works/' + wid + (f'/versions/{version}' if version else ''))
    raw = client.request(f'/api/works/{wid}/versions/{m["version"]}/bundle', raw=True)
    out = Path(output)
    if out.exists():
        raise ValueError('Output already exists; refusing to overwrite: ' + str(out))
    files = {}
    with zipfile.ZipFile(io.BytesIO(raw)) as z:
        if len(z.infolist()) > 201:
            raise ValueError('Too many files in response.')
        total = 0
        for item in z.infolist():
            n = safe_path(item.filename)
            if stat.S_ISLNK(item.external_attr >> 16) or item.is_dir() or n in files:
                raise ValueError('Unsafe or duplicate ZIP entry.')
            total += item.file_size
            if total > 33 * 1024 * 1024 or item.file_size > 12 * 1024 * 1024:
                raise ValueError('ZIP response exceeds limits.')
            files[n] = z.read(item)
    expected = {f['path']: f for f in m['files']}
    if set(expected) != set(files):
        raise ValueError('Package inventory mismatch.')
    h = hashlib.sha256()
    for n, b in sorted(files.items()):
        if len(b) != expected[n]['bytes'] or hashlib.sha256(b).hexdigest() != expected[n]['sha256']:
            raise ValueError('File hash mismatch: ' + n)
        h.update(n.encode()); h.update(b'\0'); h.update(hashlib.sha256(b).digest())
    if h.hexdigest() != m['sha256']:
        raise ValueError('Snapshot hash mismatch.')
    out.mkdir(parents=True, mode=0o700)
    for n, b in files.items():
        f = out / n; f.parent.mkdir(parents=True, exist_ok=True); f.write_bytes(b)
    return {'work_id': wid, 'version': m['version'], 'output': str(out.resolve()), 'sha256': m['sha256'], 'verified_hashes': True}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--url', default=os.environ.get('MATHHUB_URL', 'https://mathhub.eggtea.org'))
    sub = p.add_subparsers(dest='command', required=True)
    for cmd in ['login', 'register']:
        s = sub.add_parser(cmd); s.add_argument('--username', required=True); s.add_argument('--display-name')
    for cmd in ['logout', 'whoami']:
        sub.add_parser(cmd)
    s = sub.add_parser('list'); s.add_argument('--type'); s.add_argument('--field'); s.add_argument('--hide-demos', action='store_true'); s.add_argument('--offset',type=int,default=0)
    for cmd in ['read', 'pull', 'cite']:
        s = sub.add_parser(cmd); s.add_argument('work'); s.add_argument('--version',type=int)
        if cmd == 'pull':
            s.add_argument('--output',required=True)
    for cmd in ['validate', 'upload']:
        s = sub.add_parser(cmd); s.add_argument('source'); s.add_argument('--note',default='Initial publication'); s.add_argument('--demo',action='store_true')
    for cmd in ['push','propose']:
        s = sub.add_parser(cmd); s.add_argument('work'); s.add_argument('source'); s.add_argument('--base-version',type=int,required=True); s.add_argument('--note',required=True)
        if cmd=='push':
            s.add_argument('--resolves',default='')
        else:
            s.add_argument('--title',required=True)
    s=sub.add_parser('comment');s.add_argument('work');s.add_argument('--version',type=int,required=True);s.add_argument('--file',required=True)
    for cmd in ['comments','history','revisions']:
        s=sub.add_parser(cmd);s.add_argument('work')
    for cmd in ['merge','close']:
        s=sub.add_parser(cmd);s.add_argument('work');s.add_argument('revision')
    s=sub.add_parser('compare');s.add_argument('work');s.add_argument('--from-version',type=int,required=True);s.add_argument('--to-version',type=int,required=True)
    s=sub.add_parser('collaborator');s.add_argument('work');s.add_argument('username');s.add_argument('--remove',action='store_true')
    s=sub.add_parser('archive');s.add_argument('work');s.add_argument('--undo',action='store_true')
    a=p.parse_args();c=Client(a.url);cmd=a.command
    if cmd in {'login','register'}:
        password=getpass.getpass('Unique password: ')
        b={'username':a.username,'password':password}
        if cmd=='register':
            if getpass.getpass('Confirm password: ')!=password:
                raise ValueError('Passwords do not match.')
            if input('Accept the site privacy policy and publishing terms? [yes/no] ').strip().lower()!='yes':
                raise ValueError('Registration cancelled.')
            b.update(display_name=a.display_name or a.username,accept_terms=True)
        result=c.request('/api/auth/'+cmd,'POST',b);result.pop('csrf',None);emit(result)
    elif cmd=='logout':
        emit(c.request('/api/auth/logout','POST'))
    elif cmd=='whoami':
        r=c.request('/api/session');r.pop('csrf',None);emit(r)
    elif cmd=='list':
        qs=urllib.parse.urlencode({k:v for k,v in {'type':a.type,'field':a.field,'demo':'hide' if a.hide_demos else None,'offset':a.offset}.items() if v is not None});emit(c.request('/api/works?'+qs))
    elif cmd in {'read','cite'}:
        path='/api/works/'+a.work
        if cmd=='cite' and not a.version:
            a.version=c.request(path)['version']
        if a.version:
            path+=f'/versions/{a.version}'
        if cmd=='cite':
            path+='/citation'
        emit(c.request(path))
    elif cmd=='pull':
        emit(pull(c,a.work,a.version,a.output))
    elif cmd in {'validate','upload','push','propose'}:
        fields={'change_note':a.note,'is_demo':str(getattr(a,'demo',False)).lower()}
        if cmd in {'push','propose'}:
            fields['base_version']=a.base_version
        if cmd=='propose':
            fields['title']=a.title
        if cmd=='push':
            fields['resolves']=a.resolves
        data,ct=multipart(a.source,fields)
        path='/api/bundles/validate' if cmd=='validate' else '/api/works' if cmd=='upload' else f'/api/works/{a.work}/'+('versions' if cmd=='push' else 'revisions')
        emit(c.request(path,'POST',data,ct))
    elif cmd=='comment':
        b=json.loads(Path(a.file).read_text());b['version']=a.version;emit(c.request(f'/api/works/{a.work}/comments','POST',b))
    elif cmd in {'comments','history','revisions'}:
        emit(c.request(f'/api/works/{a.work}/'+('versions' if cmd=='history' else cmd)))
    elif cmd in {'merge','close'}:
        emit(c.request(f'/api/works/{a.work}/revisions/{a.revision}/{cmd}','POST'))
    elif cmd=='compare':
        emit(c.request(f'/api/works/{a.work}/compare?from={a.from_version}&to={a.to_version}'))
    elif cmd=='collaborator':
        emit(c.request(f'/api/works/{a.work}/collaborators','POST',{'username':a.username,'remove':a.remove}))
    elif cmd=='archive':
        emit(c.request(f'/api/works/{a.work}','PATCH',{'archived':not a.undo}))


if __name__=='__main__':
    try:
        main()
    except (ValueError,OSError,urllib.error.URLError,zipfile.BadZipFile) as e:
        print(json.dumps({'error':str(e)},ensure_ascii=False),file=sys.stderr);sys.exit(1)