#!/usr/bin/env python3
from __future__ import annotations
import argparse, hashlib, html.parser, json, os, shutil, stat, sys, tempfile, time
import urllib.error, urllib.parse, urllib.request
from pathlib import Path

SELF_URL='https://static.afi-d.ru/afi.py'
BASE_URL='https://static.afi-d.ru/'
TTL=86400
UA='afi-static-client/1.0'

def sp(): return Path(__file__).resolve()
def cp(): return sp().parent/'cache.json'
def stp(): return sp().parent/'.afi-state.json'
def load(p,d):
    try: return json.loads(p.read_text(encoding='utf-8'))
    except Exception: return d

def atomic(p,b,mode=None):
    p.parent.mkdir(parents=True,exist_ok=True)
    fd,n=tempfile.mkstemp(prefix='.'+p.name+'.',dir=p.parent)
    try:
        with os.fdopen(fd,'wb') as f: f.write(b); f.flush(); os.fsync(f.fileno())
        if mode is not None: os.chmod(n,mode)
        os.replace(n,p)
    finally:
        try: os.unlink(n)
        except FileNotFoundError: pass

def req(url):
    return urllib.request.urlopen(urllib.request.Request(url,headers={'User-Agent':UA}),timeout=30)

def update(force=False):
    s=load(stp(),{}); now=int(time.time())
    if not force and now-int(s.get('last_update_check',0) or 0)<TTL: return False
    try:
        with req(SELF_URL) as r: remote=r.read()
    except Exception as e:
        print(f'afi: update check failed: {e}',file=sys.stderr); return False
    if not remote.startswith(b'#!/usr/bin/env python3'):
        print('afi: invalid update payload',file=sys.stderr); return False
    try: compile(remote.decode(),'afi.py','exec')
    except Exception as e:
        print(f'afi: invalid update syntax: {e}',file=sys.stderr); return False
    s={'last_update_check':now,'remote_sha256':hashlib.sha256(remote).hexdigest()}
    atomic(stp(),(json.dumps(s,indent=2)+'\n').encode())
    cur=sp().read_bytes()
    if hashlib.sha256(cur).digest()==hashlib.sha256(remote).digest(): return False
    mode=stat.S_IMODE(sp().stat().st_mode)|stat.S_IXUSR
    atomic(sp(),remote,mode)
    return True

class P(html.parser.HTMLParser):
    def __init__(self): super().__init__(); self.hrefs=[]
    def handle_starttag(self,t,a):
        if t.lower()=='a':
            h=dict(a).get('href')
            if h: self.hrefs.append(h)

def norm(x,d=False):
    x=urllib.parse.unquote(x).replace('\\','/').lstrip('/')
    parts=[p for p in x.split('/') if p not in ('','.')]
    if '..' in parts: raise ValueError('path traversal')
    y='/'.join(parts)
    return y+'/' if d and y else y

def crawl():
    base=urllib.parse.urlsplit(BASE_URL); q=[BASE_URL]; seen=set(); dirs=set(); files=set()
    while q:
        u=q.pop(0)
        if u in seen: continue
        seen.add(u)
        try:
            with req(u) as r: body=r.read().decode('utf-8','replace')
        except Exception as e:
            print(f'afi: skip {u}: {e}',file=sys.stderr); continue
        p=P(); p.feed(body)
        for h in p.hrefs:
            if h.startswith(('#','?','mailto:','javascript:')): continue
            a=urllib.parse.urlsplit(urllib.parse.urljoin(u,h))
            if a.netloc!=base.netloc or a.query or a.fragment: continue
            if not a.path.startswith(base.path): continue
            isd=a.path.endswith('/')
            try: rel=norm(a.path[len(base.path):],isd)
            except ValueError: continue
            if not rel: continue
            if isd:
                if rel not in dirs: dirs.add(rel); q.append(urllib.parse.urljoin(BASE_URL,rel))
            else: files.add(rel)
    now=int(time.time()); data={'schema':1,'base_url':BASE_URL,'generated_at':now,'directories':sorted(dirs),'files':sorted(files)}
    atomic(cp(),(json.dumps(data,ensure_ascii=False,indent=2)+'\n').encode())
    return data

def cache(force=False):
    c=load(cp(),{})
    if force or not c or int(time.time())-int(c.get('generated_at',0) or 0)>=TTL:
        print('afi: refreshing index...',file=sys.stderr); c=crawl()
    return c

def url_for(p):
    p=norm(p,p.endswith('/'))
    q='/'.join(urllib.parse.quote(x) for x in p.split('/') if x)
    if p.endswith('/'): q+='/'
    return urllib.parse.urljoin(BASE_URL,q)

def one(remote,dst):
    dst.parent.mkdir(parents=True,exist_ok=True)
    fd,n=tempfile.mkstemp(prefix='.'+dst.name+'.',dir=dst.parent)
    try:
        with os.fdopen(fd,'wb') as f, req(url_for(remote)) as r: shutil.copyfileobj(r,f,1024*1024); f.flush(); os.fsync(f.fileno())
        os.replace(n,dst); print(dst)
    finally:
        try: os.unlink(n)
        except FileNotFoundError: pass

def download(sel,c):
    sel=norm(sel,sel.endswith('/')); files=set(c['files']); dirs=set(c['directories'])
    if sel in files: one(sel,Path.cwd()/Path(sel).name); return
    d=sel.rstrip('/')+'/' if sel else ''
    if d and d not in dirs: raise SystemExit(f'afi: not found: {sel}')
    m=sorted(x for x in files if x.startswith(d))
    if not m: raise SystemExit('afi: empty or missing directory')
    root=Path.cwd()/(Path(d.rstrip('/')).name if d else '')
    for x in m: one(x,root.joinpath(*x[len(d):].split('/')))

COMP=r'''_afi_complete(){
 local cur="${COMP_WORDS[COMP_CWORD]}"
 if [[ $COMP_CWORD -eq 1 ]]; then mapfile -t COMPREPLY < <(compgen -W "download update refresh install help" -- "$cur"); return; fi
 if [[ "${COMP_WORDS[1]}" == download ]]; then mapfile -t COMPREPLY < <(afi _complete "$cur" 2>/dev/null); compopt -o nospace; fi
}
complete -F _afi_complete afi
'''

def install():
    b=Path.home()/'.local/bin'; c=Path.home()/'.local/share/bash-completion/completions'; b.mkdir(parents=True,exist_ok=True); c.mkdir(parents=True,exist_ok=True)
    t=b/'afi'
    try:
        if t.exists() or t.is_symlink(): t.unlink()
        t.symlink_to(sp())
    except OSError: shutil.copy2(sp(),t); t.chmod(t.stat().st_mode|stat.S_IXUSR)
    atomic(c/'afi',COMP.encode())
    rc=Path.home()/'.bashrc'; line='export PATH="$HOME/.local/bin:$PATH"'; txt=rc.read_text(encoding='utf-8') if rc.exists() else ''
    if line not in txt:
        with rc.open('a',encoding='utf-8') as f: f.write('\n# afi\n'+line+'\n')
    crawl(); print('Installed. Run: source ~/.bashrc && source ~/.local/share/bash-completion/completions/afi')

def main():
    a=argparse.ArgumentParser(prog='afi'); a.add_argument('--update',action='store_true'); s=a.add_subparsers(dest='cmd')
    d=s.add_parser('download'); d.add_argument('path'); s.add_parser('update'); s.add_parser('refresh'); s.add_parser('install'); x=s.add_parser('_complete'); x.add_argument('current',nargs='?',default='')
    z=a.parse_args(); force=z.update or z.cmd=='update'
    if update(force): os.execv(str(sp()),[str(sp()),*sys.argv[1:]])
    if z.cmd=='_complete':
        c=load(cp(),{}); cur=z.current.lstrip('/'); [print(i) for i in list(c.get('directories',[]))+list(c.get('files',[])) if i.startswith(cur)]; return
    if z.cmd=='install': install(); return
    if z.cmd=='refresh': crawl(); return
    c=cache(z.cmd=='update')
    if z.cmd=='download': download(z.path,c)
    elif z.cmd=='update': print('afi: updated')
    else: a.print_help()
if __name__=='__main__': main()