from io import BufferedReader, BytesIO
from typing import Any, Union
from lxml import html
import requests
import logging
import backoff
import websocket
import json
import zipfile

logger = logging.getLogger(__name__)


class EscriptoriumConnector:
    def __init__(
        self,
        base_url: str,
        api_url: str,
        username: str,
        password: str,
        project: str = None,
    ):
        # Make sure the urls terminates with a front slash
        self.api_url = api_url if api_url[-1] == "/" else api_url + "/"
        self.base_url = base_url if base_url[-1] == "/" else base_url + "/"

        self.headers = {"Accept": "application/json"}
        self.project = project

        self.session_requests = requests.session()
        login_url = f"""{self.base_url}login/"""
        result = self.session_requests.get(login_url)
        tree = html.fromstring(result.text)
        authenticity_token = list(
            set(tree.xpath("//input[@name='csrfmiddlewaretoken']/@value"))
        )[0]
        payload = {
            "username": username,
            "password": password,
            "csrfmiddlewaretoken": authenticity_token,
        }
        result = self.session_requests.post(
            login_url, data=payload, headers={**self.headers, "referer": login_url}
        )
        self.cookie = "; ".join(
            [
                f"""{k}={v}"""
                for k, v in self.session_requests.cookies.get_dict().items()
            ]
        )
        result = self.session_requests.get(
            self.base_url + "profile/apikey/", headers=self.headers
        )
        tree = html.fromstring(result.text)
        api_key = list(set(tree.xpath("//button[@id='api-key-clipboard']/@data-key")))[
            0
        ]
        self.headers["Authorization"] = f"""Token {api_key}"""

    def __on_message(self, ws, message):
        logging.debug(message)

    def __on_error(self, ws, error):
        logging.debug(error)

    def __on_close(self, ws, close_status_code, close_msg):
        logging.debug("### websocket closed ###")
        logging.debug(close_status_code)
        logging.debug(close_msg)

    def __on_open(self, ws):
        logging.debug("### websocket opened ###")

        # def run(*args):
        #     for i in range(3):
        #         time.sleep(1)
        #         ws.send("Hello %d" % i)
        #     time.sleep(1)
        #     ws.close()
        #     print("thread terminating...")

        # thread.start_new_thread(run, ())

    @backoff.on_exception(backoff.expo, requests.exceptions.HTTPError, max_time=60)
    def __get_url(self, url: str) -> requests.Response:
        r = requests.get(url, headers=self.headers)
        r.raise_for_status()
        return r

    @backoff.on_exception(backoff.expo, requests.exceptions.HTTPError, max_time=60)
    def __post_url(
        self, url: str, payload: object, files: object = None
    ) -> requests.Response:
        r = (
            requests.post(url, data=payload, files=files, headers=self.headers)
            if files is not None
            else requests.post(url, data=payload, headers=self.headers)
        )
        r.raise_for_status()
        return r

    @backoff.on_exception(backoff.expo, requests.exceptions.HTTPError, max_time=60)
    def __put_url(
        self, url: str, payload: object, files: object = None
    ) -> requests.Response:
        r = (
            requests.put(url, data=payload, files=files, headers=self.headers)
            if files is not None
            else requests.put(url, data=payload, headers=self.headers)
        )
        r.raise_for_status()
        return r

    @backoff.on_exception(backoff.expo, requests.exceptions.HTTPError, max_time=60)
    def __delete_url(self, url: str) -> requests.Response:
        r = requests.delete(url, headers=self.headers)
        r.raise_for_status()
        return r

    def get_documents(self):
        r = self.__get_url(f"{self.api_url}documents/")
        info = r.json()
        documents = info["results"]
        while info["next"] is not None:
            r = self.__get_url(info["next"])
            info = r.json()
            documents = documents + info["results"]

        return documents

    def get_document(self, pk: int):
        r = self.__get_url(f"{self.api_url}documents/{pk}/")
        return r.json()

    def get_document_parts(self, doc_pk: int):
        return self.get_document_images(doc_pk)

    def get_document_part(self, doc_pk: int, part_pk: int):
        r = self.__get_url(f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/")
        return r.json()

    def get_document_part_line(self, doc_pk: int, part_pk: int, line_pk: int):
        r = self.__get_url(
            f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/lines/{line_pk}/"
        )
        return r.json()

    def get_document_part_region(self, doc_pk: int, part_pk: int, region_pk: int):
        regions = self.get_document_part_regions(doc_pk, part_pk)
        region = [x for x in regions if x["pk"] == region_pk]
        return region[0] if region else None

    def get_document_part_line_transcription(
        self, doc_pk: int, part_pk: int, line_pk: int, line_transcription_pk: int
    ):
        transcriptions = self.get_document_part_line_transcriptions(
            doc_pk, part_pk, line_pk
        )
        transcription = [x for x in transcriptions if x["pk"] == line_transcription_pk]
        return transcription[0] if transcription else None

    def get_document_part_line_transcription_by_transcription(
        self, doc_pk: int, part_pk: int, line_pk: int, transcription_pk: int
    ):
        transcriptions = self.get_document_part_line_transcriptions(
            doc_pk, part_pk, line_pk
        )
        transcription = [
            x for x in transcriptions if x["transcription"] == transcription_pk
        ]
        return transcription[0] if transcription else None

    def get_document_part_line_transcriptions(
        self, doc_pk: int, part_pk: int, line_pk: int
    ):
        line = self.get_document_part_line(doc_pk, part_pk, line_pk)
        return line["transcriptions"]

    def get_document_part_regions(self, doc_pk: int, part_pk: int):
        r = self.__get_url(f"{self.api_url}documents/{doc_pk}/parts/{part_pk}")
        part = r.json()
        return part["regions"]

    def get_document_transcription(self, doc_pk: int, transcription_pk: int):
        r = self.__get_url(
            f"{self.api_url}documents/{doc_pk}/transcriptions/{transcription_pk}"
        )
        return r.json()

    def create_document_transcription(self, doc_pk: int, transcription_name: str):
        r = self.__post_url(
            f"{self.api_url}documents/{doc_pk}/transcriptions/",
            {"name": transcription_name},
        )
        return r.json()

    def get_document_transcriptions(self, doc_pk: int):
        r = self.__get_url(f"{self.api_url}documents/{doc_pk}/transcriptions/")
        return r.json()

    def create_document_line_transcription(
        self,
        doc_pk: int,
        parts_pk: int,
        line_pk: int,
        transcription_pk: int,
        transcription_content: str,
        graphs: Union[Any, None],
    ):
        # Do I need a "transcription" field too? I don't know what it means.
        payload = {
            "pk": transcription_pk,
            "line": line_pk,
            "content": transcription_content,
        }
        if graphs is not None:
            payload["graphs"] = graphs
        r = self.__post_url(
            f"{self.api_url}documents/{doc_pk}/parts/{parts_pk}/transcriptions/",
            payload,
        )
        return r.json()

    def download_part_alto_transcription(
        self,
        document_pk: int,
        part_pk: Union[list[int], int],
        transcription_pk: int,
    ) -> Union[list[bytes], None]:
        return self.__download_part_output_transcription(
            document_pk, part_pk, transcription_pk, "alto"
        )

    def download_part_pagexml_transcription(
        self,
        document_pk: int,
        part_pk: Union[list[int], int],
        transcription_pk: int,
    ) -> Union[list[bytes], None]:
        return self.__download_part_output_transcription(
            document_pk, part_pk, transcription_pk, "pagexml"
        )

    def download_part_text_transcription(
        self,
        document_pk: int,
        part_pk: Union[list[int], int],
        transcription_pk: int,
    ) -> Union[list[bytes], None]:
        return self.__download_part_output_transcription(
            document_pk, part_pk, transcription_pk, "text"
        )

    def __download_part_output_transcription(
        self,
        document_pk: int,
        part_pk: Union[list[int], int],
        transcription_pk: int,
        output_type: str,
    ) -> Union[list[bytes], None]:
        if self.cookie is None:
            raise Exception("Must use websockets to download ALTO exports")

        download_link = None
        ws = websocket.WebSocket()
        ws.connect(
            f"{self.base_url.replace('http', 'ws')}ws/notif/", cookie=self.cookie
        )
        self.__post_url(
            f"{self.api_url}documents/{document_pk}/export/",
            {
                "task": "export",
                "document": document_pk,
                "parts": part_pk,
                "transcription": transcription_pk,
                "file_format": output_type,
            },
        )

        message = ws.recv()
        ws.close()
        logging.debug(message)
        msg = json.loads(message)
        if "export" in msg["text"].lower():
            for entry in msg["links"]:
                if entry["text"].lower() == "download":
                    download_link = entry["src"]

        if download_link is None:
            logging.warning(
                f"Did not receive a link to download ALTO export for {document_pk}, {part_pk}, {transcription_pk}"
            )
            return None
        alto_request = self.__get_url(f"{self.base_url}{download_link}")

        if alto_request.status_code != 200:
            return None

        zip = zipfile.ZipFile(BytesIO(alto_request.content))
        return [zip.read(x) for x in zip.infolist()]

    def upload_part_transcription(
        self,
        document_pk: int,
        part_pk: int,
        transcription_name: str,
        filename: str,
        file_data: BufferedReader,
    ):
        """Upload a txt, PageXML, or ALTO file.

        Args:
            document_pk (int): Document PK
            part_pk (int): Part PK
            transcription_name (str): Transcription name
            filename (str): Filename
            file_data (BufferedReader): File data as a BufferedReader

        Returns:
            null: Nothing
        """

        return self.__post_url(
            f"{self.api_url}documents/{document_pk}/imports/",
            {
                "task": "import-xml",
                "name": transcription_name,
                "document": document_pk,
                "parts": part_pk,
            },
            {"upload_file": (filename, file_data)},
        )

    def get_document_part_lines(self, doc_pk: int, part_pk: int):
        r = self.__get_url(f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/lines/")
        line_info = r.json()
        lines = line_info["results"]
        while line_info["next"] is not None:
            r = self.__get_url(line_info["next"])
            line_info = r.json()
            lines = lines + line_info["results"]

        return lines

    def create_document_part_line(self, doc_pk: int, part_pk: int, new_line: object):
        r = self.__post_url(
            f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/lines/", new_line
        )
        return r

    def delete_document_part_line(self, doc_pk: int, part_pk: int, line_pk: int):
        r = self.__delete_url(
            f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/lines/{line_pk}"
        )
        return r

    def get_document_part_transcriptions(self, doc_pk: int, part_pk: int):
        get_url = f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/transcriptions/"
        r = self.__get_url(get_url)
        transcriptions_info = r.json()
        transcriptions = transcriptions_info["results"]
        while transcriptions_info["next"] is not None:
            r = self.__get_url(transcriptions_info["next"])
            transcriptions_info = r.json()
            transcriptions = transcriptions + transcriptions_info["results"]
        return transcriptions

    def create_document_part_transcription(
        self, doc_pk: int, part_pk: int, transcription: object
    ):
        r = self.__post_url(
            f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/transcriptions/",
            transcription,
        )
        return r

    def create_document_part_region(self, doc_pk: int, part_pk: int, region: object):
        if not isinstance(region["box"], str):
            region["box"] = json.dumps(region["box"])

        r = self.__post_url(
            f"{self.api_url}documents/{doc_pk}/parts/{part_pk}/blocks/",
            region,
        )
        return r

    def get_line_types(self):
        r = self.__get_url(f"{self.api_url}types/line/")
        line_type_info = r.json()
        line_types = line_type_info["results"]
        while line_type_info["next"] is not None:
            r = self.__get_url(line_type_info["next"])
            line_type_info = r.json()
            line_types = line_types + line_type_info["results"]

        return line_types

    def create_line_type(self, line_type: object):
        r = self.__post_url(f"{self.api_url}types/line/", line_type)
        return r

    def get_region_types(self):
        r = self.__get_url(f"{self.api_url}types/block/")
        block_type_info = r.json()
        block_types = block_type_info["results"]
        while block_type_info["next"] is not None:
            r = self.__get_url(block_type_info["next"])
            block_type_info = r.json()
            block_types = block_types + block_type_info["results"]

        return block_types

    def create_region_type(self, region_type: object):
        r = self.__post_url(f"{self.api_url}types/block/", region_type)
        return r

    def get_document_images(self, document_pk: int):
        r = self.__get_url(f"{self.api_url}documents/{document_pk}/parts/")
        image_info = r.json()
        image_names = image_info["results"]
        while image_info["next"] is not None:
            r = self.__get_url(image_info["next"])
            image_info = r.json()
            image_names = image_names + image_info["results"]

        return image_names

    def delete_document_parts(self, document_pk: int, start: int, end: int):
        parts = self.get_document_images(document_pk)
        for part in parts[start:end]:
            r = self.__delete_url(
                f'{self.api_url}documents/{document_pk}/parts/{part["pk"]}/'
            )

    def get_image(self, img_url: str):
        r = self.__get_url(f"{self.base_url}{img_url}")
        return r.content

    def create_document(self, doc_data: object):
        if self.project:
            doc_data["project"] = self.project
        return self.__post_url(f"{self.api_url}documents/", doc_data)

    def create_image(
        self, document_pk: int, image_data_info: object, image_data: bytes
    ):
        return self.__post_url(
            f"{self.api_url}documents/{document_pk}/parts/",
            image_data_info,
            {"image": (image_data_info["filename"], image_data)},
        )


if __name__ == "__main__":
    source_url = "https://www.escriptorium.fr/"
    source_api = f"{source_url}api/"
    source_token = ""
    source = EscriptoriumConnector(source_url, source_api, source_token)
    print(source.get_documents())
