#!/usr/bin/env python3
from collections import defaultdict
from pathlib import Path

from bs4 import BeautifulSoup
from local_file_picker import local_file_picker

from nicegui import ui, app, html, run

from datetime import datetime

import zmq
import zmq.asyncio

from loguru import logger
import sys

from anonymiser import Anonymizer
import pydicom

import requests
import time

from multiprocessing import Manager, Queue

import watchfiles


from typing import Optional

from fastapi import Request
from fastapi.responses import RedirectResponse
from starlette.middleware.base import BaseHTTPMiddleware
#import typer

#typer_app = typer.Typer()

#logger.add(sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO")

#async def pick_file() -> None:
#    result = await local_file_picker('~', multiple=True)
#    ui.notify(f'You chose {result}')
#    ui.html("<h2>Uploader tool</h2>")
#
#    ui.chat_message(f'{result}!',
#                name='Robot',
#                stamp='now',
#                avatar='https://robohash.org/ui')
#
#
#@ui.page('/')
#def index():
#    ui.html("<h2>Uploader tool</h2>")
#    ui.button('Choose file', on_click=pick_file, icon='folder')

EXPORT_PATH = Path("./test/export")
PROCESS_PATH = Path("./test/to_process")
ANON_PATH = Path("./test/anon")

EXPORT_PATH.mkdir(exist_ok=True, parents=True)
PROCESS_PATH.mkdir(exist_ok=True, parents=True)
EXPORT_PATH.mkdir(exist_ok=True, parents=True)

BASE_SITE_URL = "http://localhost:8000"

LOGIN_URL = f"{BASE_SITE_URL}/accounts/login/"

LOGIN_SUCCESS = False

rqst = None


use_zmq = False
loaded_files: dict = {}
uploaded_files: dict = {}
duplicate_series = set()

if use_zmq:
    context = zmq.asyncio.Context()
    socket = context.socket(zmq.PULL)
    #socket.bind('tcp://localhost:5575')
    try:
        socket.bind('tcp://localhost:9975')
    except zmq.error.ZMQError as e:
        logger.debug("send reload")
        socket = context.socket(zmq.PUSH)
        socket.connect('tcp://localhost:5575')
        socket.send_string("loaded")
        sys.exit(0)

    poller = zmq.asyncio.Poller()
    poller.register(socket, zmq.POLLIN)


anonymizer = Anonymizer(ANON_PATH)

def load_files(q: Queue):
    loaded_files: dict = {}
    loaded_series = defaultdict(list)
    loaded_series_data = {}
    # Move files
    src_path = EXPORT_PATH

    to_process = []
    for file in src_path.iterdir():
        logger.debug(f"Moving {file} to {PROCESS_PATH}")

        file_to_process = PROCESS_PATH / file.name
        file.rename(file_to_process)

        to_process.append(file_to_process)

    to_process_len = len(to_process)
    for n, file in enumerate(to_process):
        dataset, output_file = anonymizer.anonymize_file(file, remove_original=True)
        #with processed_files:
        #    ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {file.name} -> {output_file.name}")

        loaded_files[output_file] = (dataset.StudyDescription, dataset.SeriesDescription, str(dataset.SeriesInstanceUID))

        loaded_series[dataset.SeriesInstanceUID].append(output_file)
        loaded_series_data[dataset.SeriesInstanceUID] = (dataset.StudyDescription, dataset.SeriesDescription)

        q.put_nowait((n + 1) / to_process_len)

        #time.sleep(0.01)

    return loaded_files, loaded_series, loaded_series_data 
    #load_series_view()

def load_series_view(loaded_series, loaded_series_data):
    logger.debug("load series view")
    logger.debug(loaded_series)
    for key in loaded_series:
        logger.debug(f"load series: {key}")
        with series:
            with ui.card():
                ui.label(loaded_series_data[key][0])
                ui.label(loaded_series_data[key][1])
                ui.label(f"Images: {len(loaded_series[key])}")

    series_title.visible = True


def reload_anonymized(q: Queue):
    loaded_files = {}
    loaded_series = defaultdict(list)
    loaded_series_data = {}

    #processed_files.clear()


    to_process = list(ANON_PATH.glob("*.dcm"))
    to_process_len = len(to_process)
    for n, file in enumerate(to_process):
        with pydicom.dcmread(file) as dataset:
            loaded_files[file] = (dataset.StudyDescription, dataset.SeriesDescription, str(dataset.SeriesInstanceUID))

            loaded_series[dataset.SeriesInstanceUID].append(file)
            loaded_series_data[dataset.SeriesInstanceUID] = (dataset.StudyDescription, dataset.SeriesDescription)

        #with processed_files:
        #    ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {file.name} -> Reloaded")


        q.put_nowait(f"{(n + 1) / to_process_len * 100:.0f}%")
        logger.debug(file)

        #time.sleep(0.01)
    
    logger.debug("end reload")
    logger.debug(loaded_files)
    logger.debug(loaded_series)
    logger.debug("-----")

    return loaded_files, loaded_series, loaded_series_data 

    #load_series_view()

def upload_files(q, rqst):
    #global rqst#, uploaded_files

    uploaded_files = {}

    files_to_upload = []
    for file in ANON_PATH.iterdir():
        files_to_upload.append(("files", open(str(file), "rb")))

    if not files_to_upload:
        ui.notify('No files to upload', color='negative')
        return

    # chunck files
    n = 10
    chunked_files = [
        files_to_upload[i : i + n] for i in range(0, len(files_to_upload), n)
    ]


    upload_file_list = []
    duplicate_file_list = []
    failed = []
    duplicate_series = set()

    logger.debug("START UPLOAD")
    logger.debug(files_to_upload)

    to_process_len = len(chunked_files)

    for n, files in enumerate(chunked_files):

        def upload_files_(files):
            rqst.headers["X-CSRFToken"] = rqst.cookies["csrftoken"]
            # print(self.rqst.headers)
            # print(self.rqst.cookies)
            resp = rqst.post(
                f"{BASE_SITE_URL}/api/atlas/upload_dicom",
                # data=data,
                files=files,
            )

            logger.debug(f"Resp: {resp}")
            logger.debug(f"{resp.content}")
            return resp

        # progress_dialog.Update(n, f"Uploading batch {n}/{len(chunked_files)}")

        logger.debug(f"n: {n}")
        # try to upload the files

        for i in range(3):
            logger.error(f"i: {i}")
            resp = upload_files_(files)
            if resp.status_code == 200:
                upload_file_list.extend(resp.json()["uploaded"])
                duplicate_file_list.extend(resp.json()["duplicates"])
                failed.extend(resp.json()["failed"])
                duplicate_series.update(resp.json()["duplicate_series"])

                break

            logger.debug(f"n: {n} fail (attempt {i})")

        q.put_nowait((n + 1) / to_process_len)
        # progress_dialog.Destroy()
    print(upload_file_list)
    print("dup", duplicate_file_list)
    print("failed", failed)

    for f, hash in upload_file_list:
        uploaded_files[f] = "success"

    for f, hash in duplicate_file_list:
        uploaded_files[f] = "duplicate"

    for f in failed:
        uploaded_files[f] = "failed"


    return uploaded_files, upload_file_list, duplicate_file_list, failed, duplicate_series


dark = ui.dark_mode()
dark.enable()

ui.page_title('Uploader tool')
ui.label('Uploader tool').classes("text-h2")
user_label = ui.label("User: None")

login_button = ui.button("LOGIN", on_click=lambda: ui.navigate.to("/login"))
login_button.bind_visibility_from(globals(), "LOGIN_SUCCESS", backward=lambda v: not v)

updates = ui.list()

series_title = ui.label('Series to upload').classes("text-h3")
series_title.visible = False

series = ui.row()


async def upload_files_start():
    global uploaded_files, duplicate_series
    upload_progress.visible = True
    uploaded_files, upload_file_list, duplicate_file_list, failed, duplicate_series = await run.cpu_bound(upload_files, upload_queue, rqst)
    upload_progress.visible = False

    ui.notify(f"Uploaded {len(upload_file_list)} files", color='positive')
    ui.notify(f"Duplicate {len(duplicate_file_list)} files", color='warning')
    ui.notify(f"Failed {len(failed)} files", color='negative')

    uploaded_files_ui.refresh()

upload_button = ui.button("Upload", on_click=upload_files_start)
upload_button.bind_visibility_from(globals(), "LOGIN_SUCCESS")

if use_zmq:
    async def read_loop() -> None:
        while not app.is_stopped:
            events = await poller.poll()
            if socket in dict(events):
                data = await socket.recv()
                msg = data.decode()

                logger.debug(f'Received msg {msg}')

                match msg:
                    case "loaded":
                        with updates:
                            ui.item(f"{datetime.now().strftime('%H:%M:%S')} - {msg}")
                        logger.debug("loaded")
                        load_files_start()

                        
                    case _:
                        logger.error("invalid message")
    app.on_startup(read_loop)

# Create a queue to communicate with the heavy computation process
queue = Manager().Queue()
# Update the progress bar on the main process
ui.timer(0.1, callback=lambda: progressbar.set_value(queue.get() if not queue.empty() else progressbar.value))

# Create a queue to communicate with the heavy computation process
upload_queue = Manager().Queue()
# Update the progress bar on the main process
ui.timer(0.1, callback=lambda: upload_progressbar.set_value(upload_queue.get() if not upload_queue.empty() else upload_progressbar.value))
upload_progress = ui.row().classes("w-full place-content-center bg-blue-900")

with upload_progress:
    ui.label("Uploading files")
    upload_progressbar = ui.linear_progress(value=0).props('instant-feedback')
upload_progress.visible = False

anon_progress = ui.row().classes("w-full place-content-center bg-blue-900")

with anon_progress:
    ui.label("Anonymizing files")
    progressbar = ui.linear_progress(value=0).props('instant-feedback')
anon_progress.visible = False


async def load_files_start():
    global loaded_files
    anon_progress.visible = True
    loaded_files, loaded_series, loaded_series_data = await run.cpu_bound(load_files, queue)
    anon_progress.visible = False
    load_series_view(loaded_series, loaded_series_data)
    loaded_files_ui.refresh()

async def reload_anonymized_start():
    global loaded_files
    anon_progress.visible = True
    loaded_files, loaded_series, loaded_series_data = await run.cpu_bound(reload_anonymized, queue)
    anon_progress.visible = False
    load_series_view(loaded_series, loaded_series_data)
    loaded_files_ui.refresh()

with ui.expansion('Extra', icon='build').classes('w-full'):
    ui.button("Load files", on_click=load_files_start)
    ui.button("Reload anon", on_click=reload_anonymized_start)


@ui.page('/login')
def login() -> Optional[RedirectResponse]:
    def try_login() -> None:  # local function to avoid passing username and password as arguments
        logger.debug(username.value)
        #if passwords.get(username.value) == password.value:
        #    app.storage.user.update({'username': username.value, 'authenticated': True})
        #    ui.navigate.to(app.storage.user.get('referrer_path', '/'))  # go back to where the user wanted to go
        #else:
        #    ui.notify('Wrong username or password', color='negative')
        user = username.value
        pw = password.value

        global rqst, LOGIN_SUCCESS
        rqst = requests.session()
        try: 
            rsp = rqst.get(LOGIN_URL)
        except requests.exceptions.ConnectionError:
            ui.notify('Connection error!', color='negative')
            return

        token = (
            BeautifulSoup(rsp.content, "html.parser").find("input").attrs["value"]
        )  # , attr={"name": "csrfmiddlewaretoken"}).attrs("value")

        # token = rsp.cookies["csrftoken"]
        # header = {"X-CSRFToken": token}
        # cookies = {"csrftoken": token}

        data = {
            "username": user,
            "password": pw,
            "csrfmiddlewaretoken": token,
            "next": "/",
        }

        rqst.headers["Referer"] = LOGIN_URL

        rsp = rqst.post(
            LOGIN_URL,
            data=data,
            # headers=header,
            # cookies=cookies
        )

        soup = BeautifulSoup(rsp.content, "html.parser")

        if soup.find("button") and soup.find("button").find(string="Login"):
            print("login fail")
            pass
        else:
            print("login success")
            LOGIN_SUCCESS = True
            user_label.clear()
            user_label.text = f"User: {user}"
        # print(rqst.headers)

        rqst.headers["X-CSRFToken"] = token

        print(rqst.headers)
        print(rqst.cookies)
        new = rqst.get(f"{BASE_SITE_URL}/accounts/api/atlas/get_cases_user")


        if LOGIN_SUCCESS:
            ui.notify('Logged in!', color='positive')
            ui.navigate.to("/")
            #login_button.visible = False
            pass
            #self.parent.upload_button.Enable()
            #self.parent.login_button.Hide()
            #self.parent.panel.Layout()
            #self.Destroy()

            # self.parent.login_button.SetBackgroundColour(wx.BLUE)
            # self.parent.login_button.SetForegroundColour(wx.WHITE)
        else:
            ui.notify('Wrong username or password', color='negative')
            pass
            #self.login_fail()

    #if app.storage.user.get('authenticated', False):
    #    return RedirectResponse('/')
    with ui.card().classes('absolute-center'):
        username = ui.input('Username').on('keydown.enter', try_login)
        password = ui.input('Password', password=True, password_toggle_button=True).on('keydown.enter', try_login)
        ui.button('Log in', on_click=try_login)
    return None

@ui.refreshable
def loaded_files_ui() -> None:
    with ui.expansion('Loaded files', icon='drive_folder_upload').classes('w-full'):
        ui.label(f"Items: {len(loaded_files)}")
        with ui.list():
            for file, data in loaded_files.items():
                ui.item(f"{file} - {data[0]} - {data[1]}")

@ui.refreshable
def uploaded_files_ui() -> None:
    with ui.expansion('Uploaded files', icon='file_upload').classes('w-full'):
        ui.label(f"Items: {len(uploaded_files)}")
        with ui.list():
            for series in duplicate_series:
                ui.item(f"Duplicate series: {series}", on_click=lambda: ui.navigate.to(f'{BASE_SITE_URL}{series}', new_tab=True)).classes("text-red-500")
        with ui.list():
            for file, data in uploaded_files.items():
                ui.item(f"{file} - {data}")

with ui.expansion('File status', icon='view_list').classes('w-full'):
    loaded_files_ui()
    uploaded_files_ui()

    #with ui.expansion('Processed files!', icon='work').classes('w-full'):
    #    processed_files = ui.list()

with ui.expansion("Settings", icon="settings").classes("w-full"):
    ui.label("Settings").classes("text-h4")
    ui.label(f"Export path: {EXPORT_PATH}")
    ui.label(f"Process path: {PROCESS_PATH}")
    ui.label(f"Anon path: {ANON_PATH}")

ui.run(storage_secret="123456")



#if __name__ in {"__main__", "__mp_main__"}:
#    typer_app()