#!/usr/bin/env python3

import argparse
import os
import requests
import json
import subprocess

class RoiResearch(object):
    def __init__(self, token=None):
        self.token = token

    def login(self, username, password):
        url = "https://themblem.com/api/v1/login/"
        data = {
            "username": username,
            "password": password,
        }
        response = requests.post(url, json=data)
        token = response.json()['token']
        self.token = token
        return token

    def fetch_scan_data(self, last_id=None):
        x = "/api/v1/scan-data/?limit=1000"
        while True:
            url = "https://themblem.com" + x
            print("fetching", url)
            response = requests.get(url, headers={"Authorization": f"Token {self.token}"})
            r = response.json()
            meta = r['meta']
            if meta['next'] is None:
                break
            x = meta['next']
            for sd in r['objects']:
                if last_id and sd['id'] <= last_id:
                    return
                yield sd

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--token", '-t', default='3ebd8c33-f46e-4b06-bda8-4c0f5f5eb530', type=str)
    parser.add_argument("--username", "-u", type=str)
    parser.add_argument("--password", "-p", type=str)
    return parser.parse_args()

def get_last_id(list_file):
    ret = None
    if not os.path.exists(list_file):
        return ret
    with open(list_file, "r") as f:
        for line in f:
            if not line.strip():
                continue
            id = json.loads(line)['id']
            if not ret or id > ret:
                ret = id
    return ret

def get_all_ids(list_file):
    ret = []
    with open(list_file, "r") as f:
        for line in f:
            ret.append(json.loads(line)['id'])
    return ret
    
def main():
    args = parse_args()
    rr = RoiResearch(args.token)
    if args.username and args.password:
        print(rr.login(args.username, args.password))
    list_file = "data/roi/list.txt"
    last_id = get_last_id(list_file)
    print("last_id", last_id)
    with open(list_file, "a") as f:
        for sd in rr.fetch_scan_data(last_id):
            if not sd['image'] or not sd['labels']:
                print("skipping", sd['id'])
                continue
            print("new id", sd['id'])
            line = json.dumps({
                'id': sd['id'],
                'image': sd['image'],
                'labels': sd['labels'],
            })
            f.write(line + "\n")
    all_ids = get_all_ids(list_file)
    for id in all_ids:
        outd = f"data/roi/samples/{id}"
        if os.path.exists(f"{outd}/{id}.json"):
            print(f"skipping {id}, json already exists")
            continue
        cmd = f'../scripts/emcli get-scan-data {id} -o {outd}'
        print(cmd)
        subprocess.call(cmd, shell=True)

if __name__ == "__main__":
    main()