import cv2
import uiautomator2 as u2
import numpy as np
import sys, os, time

delim = "; "
filter_scroll_steps = 2

def clickBrand(d, brand_name):
    print("Try finding brand %s......" % brand_name)
    if gotoProductFilter(d, 10) is None:
        d(scrollable=True).scroll.vert.toBeginning()
    else:
        d.press("back")

    while True:
        names = d(className="android.support.v7.widget.RecyclerView").child(resourceId="org.c2h4.afei.beauty:id/tv_name")
        for name in names:
            if name.get_text() == brand_name:
                name.click()
                print("Brand %s is clicked!" % brand_name)
                return
        names = d(className="android.support.v7.widget.RecyclerView").child(resourceId="org.c2h4.afei.beauty:id/tv_name_en")
        for name in names:
            if name.get_text() == brand_name:
                name.click()
                print("Brand %s is clicked!" % brand_name)
                return
        d.swipe_ext("up", scale=0.5)
        time.sleep(1)

def getProductFilterUiObj(d):
    return d(resourceId="org.c2h4.afei.beauty:id/tv_filter")

def adjustProductFilterPosition(d):
    top = d.window_size()[0]
    filter = getProductFilterUiObj(d)
    if filter.count == 0:
        print("No product filter is found!!!!!!")
        return
    while True:
        if filter.info['bounds']['top'] > top / 2:
            d.swipe_ext("up", scale=0.2)
        else:
            break
    return filter

def gotoProductFilter(d, time_cnt=None):
    cnt = 0
    while d(resourceId="org.c2h4.afei.beauty:id/tv_filter").count == 0:
        d.swipe_ext("up", scale=0.5)
        time.sleep(1)
        cnt += 1
        if time_cnt is not None and cnt > time_cnt:
            return
    return adjustProductFilterPosition(d)

def getAllFilterTypes(d):
    filter = adjustProductFilterPosition(d)
    filter.click()
    ftype_names = []
    filter_pane = d(scrollable=True)
    if filter_pane.count not in [0, 1]:
        print("filter pane is invalid(with count %d)!" % filter_pane.count)
        return
    while True:
        ftypes = d(resourceId="org.c2h4.afei.beauty:id/tv_type")
        new_flag = False
        for ftype in ftypes:
            if ftype.get_text() == "全部品类":
                continue
            elif ftype.get_text() in ftype_names:
                continue
            else:
                new_flag = True
                ftype_names.append(ftype.get_text())
        if new_flag and filter_pane.count:
            filter_pane.scroll.vert(steps=filter_scroll_steps)
        else:
            break
    if filter_pane.count:
        filter_pane.scroll.vert.backward()

    return ftype_names

def filterProduct(d, ftype_name):
    print("Select the filter type %s" % ftype_name)
    ftype = d(resourceId="org.c2h4.afei.beauty:id/tv_type", text=ftype_name)
    cnt = 0
    while ftype.count == 0:
        if cnt < 10:
            d(scrollable=True).scroll.vert(steps=filter_scroll_steps)
        else:
            d(scrollable=True).scroll.vert.backward(steps=filter_scroll_steps)
        ftype = d(resourceId="org.c2h4.afei.beauty:id/tv_type", text=ftype_name)
        cnt += 1
    ftype.click()


def cropImg(img, out, bounds):
    delta = 20 + 28
    cropped = img[
              (int(bounds['top']) + delta) : (int(bounds['bottom']) - delta),
              (int(bounds['left']) + delta) : (int(bounds['right']) - delta)
              ]  # 裁剪坐标为[y0:y1, x0:x1]
    cv2.imwrite(out, cropped)

def getProductImg(d, path, pos):
    image = d.screenshot(format='opencv')
    cropImg(image, path, pos)



def getProductKey(product, names = None):
    keyobj = None
    name = product.child(resourceId="org.c2h4.afei.beauty:id/tv_name")
    if name.count == 0:
        name = ""
    else:
        keyobj = name
        name = name.get_text()

    name_en = product.child(resourceId="org.c2h4.afei.beauty:id/tv_name_en")
    if name_en.count == 0:
        name_en = ""
    else:
        if keyobj is None:
            keyobj = name_en
        name_en = name_en.get_text()

    key = ""
    if name != "":
        key = name
    elif name_en != "":
        key = name_en

    if names is not None:
        names.append(name)
        names.append(name_en)

    return key, keyobj

def getProductDetailBasic(product, done):
    names = []
    key, keyobj = getProductKey(product, names)
    if key in done or key == "":
        return key, keyobj, None

    rate_score = keyobj.sibling(resourceId="org.c2h4.afei.beauty:id/rate_score")
    if rate_score.count == 0:
        rate_score = ""
    else:
        rate_score = rate_score.get_text()

    asess_num = keyobj.sibling(resourceId="org.c2h4.afei.beauty:id/tv_asess_num")
    if asess_num.count == 0:
        asess_num = ""
    else:
        asess_num = asess_num.get_text()

    address = keyobj.sibling(resourceId="org.c2h4.afei.beauty:id/tv_address")
    if address.count == 0:
        address = ""
    else:
        address = address.get_text()
    return key, keyobj, names + [rate_score, asess_num, address]

def getProductDetail(d, brand_name, product, done,
                     debug_trigger=False  # used for debug
                     ):
    if debug_trigger:
        d.press("back")
        raise RuntimeError("No response from APP")


    try:
        key, keyobj, basicinfo = getProductDetailBasic(product, done)
        if basicinfo is None:
            return None
    except Exception as e:
        print("ERROR: %s", str(e))
        raise RuntimeError("No response from APP")

    keyobj.click()
    img = d(resourceId="org.c2h4.afei.beauty:id/iv_image")
    wait_cnt = 0
    while img.count == 0:
        time.sleep(0.2)
        wait_cnt += 1
        if (wait_cnt == 300):
            raise RuntimeError("No response from APP")
        img = d(resourceId="org.c2h4.afei.beauty:id/iv_image")

    # wait half second for the image to be stable to take screenshot,
    # otherwise possibly the image will be different somehow
    time.sleep(0.5)

    path = brand_name + '/' + key + ".jpg"
    getProductImg(d, path, img.info['bounds'])
    effects = d(resourceId="org.c2h4.afei.beauty:id/rl_effect").child(className="android.widget.TextView")
    temp = []
    for effect in effects:
        if effect.get_text().startswith("功效"):
            continue
        else:
            temp.append(effect.get_text())
    effects = " ".join(temp)
    d.press("back")
    time.sleep(1)

    done.append(key)
    basicinfo.append(effects)
    return basicinfo

def saveProduct(f, res, product_type):
    res.append(product_type)
    print(delim.join(res))
    f.write(delim.join(res) + "\n")

def resetRuntime(d, brand_name, ftype_name, done):
    print("!!!!!!!! start reset runtime environment !!!!!!!!!!!!!")
    # re-select the target app
    d.press("home")
    time.sleep(2)
    d.press("recent")
    time.sleep(2)
    center = [ele / 2 for ele in d.window_size()]
    d.click(center[0], center[1])
    time.sleep(2)


    # go to the page where crashed
    clickBrand(d, brand_name)
    time.sleep(2)
    gotoProductFilter(d).click()
    time.sleep(2)
    filterProduct(d, ftype_name)
    time.sleep(2)

    # goto the last product
    if len(done) > 0:
        ready = False
        while True:
            products = getAllProductUiObj(d)
            for product in products:
                key, keyobj = getProductKey(product)
                if key == done[-1]:
                    ready = True
                    break
            if ready:
                break
            else:
                d.swipe_ext("up", scale=0.5)
                time.sleep(0.5)
    print("!!!!!!!! reset runtime environment successfully !!!!!!!!")

def getAllProductUiObj(d):
    return d(resourceId="org.c2h4.afei.beauty:id/rv_container").child(className="android.widget.RelativeLayout")

def getAllProducts(d, f, brand_name, product_type, starting_product_name = None):
    done = []
    nocnt = 0
    scraped_cnt = 0
    scroll_scale = 0.1

    if starting_product_name is not None:
        gotoSpecificProduct(d, starting_product_name)
    while True:
        products = getAllProductUiObj(d)
        product = None
        for temp in products:
            if temp.info['bounds']['top'] == temp.info['visibleBounds']['top']:
                product = temp
                break
        try:
            res = getProductDetail(d, brand_name, product, done)
        except RuntimeError:
            resetRuntime(d, brand_name, product_type, done)
            nocnt = 0
            continue
        if res is not None:
            saveProduct(f, res, product_type)
            scraped_cnt += 1
            if scraped_cnt % 10 == 0:
                f.flush()
            nocnt = 0
        else:
            nocnt += 1
        if nocnt == int(1 / scroll_scale) + 1:
            break
        d.swipe_ext("up", scale=scroll_scale)

    # handle the last few products in the list
    products = getAllProductUiObj(d)
    pidx = 1
    while pidx < products.count:
        try:
            res = getProductDetail(d, brand_name, products[pidx], done)
            pidx += 1
        except RuntimeError:
            resetRuntime(d, brand_name, product_type, done)
        if res is not None:
            saveProduct(f, res, product_type)
            scraped_cnt += 1

    return scraped_cnt

def gotoSpecificProduct(d, productName):
    print("Goto product of %s......" % productName)
    while True:
        products = getAllProductUiObj(d)
        for product in products:
            key, keyObj = getProductKey(product)
            # print(key)
            if key == productName:
                return
        d.swipe_ext("up", scale=0.4)

def getProductByType(argv):
    brand_name = argv[1]
    scraped_brands = []
    with open("done.txt", "r") as f:
        for line in f.readlines():
            scraped_brands.append(line.strip())
    if brand_name in scraped_brands:
        print("Brand %s is already scraped!" % brand_name)
        return

    if not os.path.exists(brand_name):
        os.makedirs(brand_name)

    starting_category_name = None
    starting_product_name = None
    if len(argv) == 3:
        starting_category_name = argv[2]
        with open(brand_name + "/products.csv", 'r') as f:
            temp = None
            for line in f.readlines():
                temp = line
            last = temp.split("; ")
            if last[-1].strip() == starting_category_name:
                if last[0] != "":
                    starting_product_name = last[0]
                else:
                    starting_product_name = last[1]

    d = u2.connect_usb('d52196830204')
    print(d.app_current())
    d.freeze_rotation()

    if getProductFilterUiObj(d).count == 0:
        clickBrand(d, brand_name)
        time.sleep(2)
        gotoProductFilter(d)
    else:
        adjustProductFilterPosition(d)

    scraped_cnt = 0
    with open(brand_name+"/products.csv", 'a+') as f:
        temp_types = getAllFilterTypes(d)

        ftype_names = []
        if starting_category_name is not None:
            starting_flag = False
            for temp_type in temp_types:
                if starting_flag:
                    ftype_names.append(temp_type)
                elif temp_type == starting_category_name:
                    ftype_names.append(temp_type)
                    starting_flag = True
        else:
            ftype_names = temp_types
        print("All the brand sub-types:")
        for ftype_name in ftype_names:
            print(ftype_name)

        for ftype_name in ftype_names:
            # if ftype_name != "粉底·BB霜":
            #     continue
            filterProduct(d, ftype_name)
            time.sleep(1)
            if ftype_names.index(ftype_name) == 0 and starting_product_name is not None:
                temp_cnt = getAllProducts(d, f, brand_name, ftype_name, starting_product_name)
            else:
                temp_cnt = getAllProducts(d, f, brand_name, ftype_name)
            f.flush()
            print("--------------------- Brand Sub-summury ---------------------")
            print("%s %s products are scraped from %s" % (temp_cnt, ftype_name, brand_name))
            print("-------------------------------------------------------------")
            scraped_cnt += temp_cnt
            if ftype_names.index(ftype_name) != len(ftype_names) - 1:
                # start select another type
                filter = getProductFilterUiObj(d)
                filter.click()
    print("--------------------- Brand Summury ---------------------")
    print("%s products are scraped from %s" % (scraped_cnt, brand_name))
    print("---------------------------------------------------------")

    # return to the brand list page in case
    # you want to continue scraping anthoer brand
    d.press("back")
    d.freeze_rotation(False)

if __name__=='__main__':
    getProductByType(sys.argv)