import sys, datetime, random, re, cv2
from os.path import dirname, realpath
sys.path.append(dirname(realpath(__file__)) + "/../")
from util.db_conn import Postgres
from util.minio_conn import HuMinio
from util import findMaxDt
import base64
from io import BytesIO
import pandas as pd
from PIL import Image
import pdfplumber


PG = Postgres("infiniflow", "docgpt")
MINIO = HuMinio("infiniflow")
def set_thumbnail(did, base64):
    sql = f"""
    update doc_info set thumbnail_base64='{base64}' 
    where
    did={did}
    """
    PG.update(sql)


def collect(comm, mod, tm):
    sql = f"""
    select 
    did, uid, doc_name, location, updated_at
    from doc_info
    where
    updated_at >= '{tm}'
    and MOD(did, {comm}) = {mod}
    and is_deleted=false
    and type <> 'folder'
    and thumbnail_base64=''
    order by updated_at asc
    limit 10
    """
    docs = PG.select(sql)
    if len(docs) == 0:return pd.DataFrame()

    mtm = str(docs["updated_at"].max())[:19]
    print("TOTAL:", len(docs), "To: ", mtm)
    return docs


def build(row):
    if not re.search(r"\.(pdf|jpg|jpeg|png|gif|svg|apng|icon|ico|webp|mpg|mpeg|avi|rm|rmvb|mov|wmv|mp4)$",
              row["doc_name"].lower().strip()):
        set_thumbnail(row["did"], "_")
        return

    def thumbnail(img, SIZE=128):
        w,h = img.size
        p = SIZE/max(w, h)
        w, h = int(w*p), int(h*p)
        img.thumbnail((w, h))
        buffered = BytesIO()
        try:
            img.save(buffered, format="JPEG")
        except Exception as e:
            try:
                img.save(buffered, format="PNG")
            except Exception as ee:
                pass
        return base64.b64encode(buffered.getvalue()).decode("utf-8")


    iobytes = BytesIO(MINIO.get("%s-upload"%str(row["uid"]), row["location"]))
    if re.search(r"\.pdf$", row["doc_name"].lower().strip()):
        pdf = pdfplumber.open(iobytes)
        img = pdf.pages[0].to_image().annotated
        set_thumbnail(row["did"], thumbnail(img))

    if re.search(r"\.(jpg|jpeg|png|gif|svg|apng|webp|icon|ico)$", row["doc_name"].lower().strip()):
        img = Image.open(iobytes)
        set_thumbnail(row["did"], thumbnail(img))

    if re.search(r"\.(mpg|mpeg|avi|rm|rmvb|mov|wmv|mp4)$", row["doc_name"].lower().strip()):
        url  = MINIO.get_presigned_url("%s-upload"%str(row["uid"]),
                                       row["location"],
                                       expires=datetime.timedelta(seconds=60)
                                      )
        cap = cv2.VideoCapture(url)
        succ = cap.isOpened()
        i = random.randint(1, 11)
        while succ:
            ret, frame = cap.read()
            if not ret: break
            if i > 0:
                i -= 1
                continue
            img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            print(img.size)
            set_thumbnail(row["did"], thumbnail(img))
        cap.release()
        cv2.destroyAllWindows()


def main(comm, mod):
    global model
    tm_fnm = f"res/thumbnail-{comm}-{mod}.tm"
    tm = findMaxDt(tm_fnm)
    rows = collect(comm, mod, tm)
    if len(rows) == 0:return

    tmf = open(tm_fnm, "a+")
    for _, r in rows.iterrows():
        build(r)
        tmf.write(str(r["updated_at"]) + "\n")
    tmf.close()


if __name__ == "__main__":
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    main(comm.Get_size(), comm.Get_rank())