from .Table import Table
from .Base import Base
from ..common import exceptions
from ..common.api_request import make_request
from tqdm.auto import tqdm
import time
import logging


import warnings
import pathlib
import os
from urllib.parse import quote as quote_uri

from ..common.retryable_upload import perform_resumable_upload, perform_standard_upload
from ..common.util import convert_data_to_parquet, get_warning


class Notebook(Base):
    def __init__(self, name, *, workflow=None, properties=None):
        from .Workflow import Workflow  # avoid circular import

        if not workflow:
            if len(name.split(".")) == 3:
                workflow = Workflow(".".join(name.split(".")[0:2]))
                name = name.split(".")[-1]
            elif os.getenv("REDIVIS_DEFAULT_WORKFLOW"):
                workflow = Workflow(os.getenv("REDIVIS_DEFAULT_WORKFLOW"))
            else:
                raise exceptions.ValueError(
                    "Invalid notebook specifier, must be the fully qualified reference if no workflow is specified"
                )

        if isinstance(workflow, str):
            workflow = Workflow(workflow)

        self.workflow = workflow
        self.name = name

        self.qualified_reference = (properties or {}).get(
            "qualifiedReference", f"{self.workflow.qualified_reference}.{self.name}"
        )
        self.scoped_reference = (properties or {}).get("scopedReference", self.name)
        self.uri = (properties or {}).get(
            "uri", f"/notebooks/{quote_uri(self.qualified_reference, '')}"
        )
        self.properties = properties

    def __repr__(self):
        return f"<Notebook {self.qualified_reference}>"

    def get(self):
        self.properties = make_request(
            method="GET",
            path=self.uri,
        )
        self.uri = self.properties["uri"]

        return self

    def exists(self):
        try:
            make_request(method="HEAD", path=self.uri)
            return True
        except exceptions.NotFoundError:
            return False

    def run(self, *, wait_for_finish=True):
        self.properties = make_request(
            method="POST",
            path=f"{self.uri}/run",
        )
        self.uri = self.properties["uri"]

        if wait_for_finish:
            while True:
                time.sleep(2)
                self.get()
                current_job = self.properties.get("currentJob") or self.properties.get(
                    "lastRunJob"
                )
                if current_job and current_job["status"] in [
                    "completed",
                    "failed",
                ]:
                    if current_job["status"] == "failed":
                        raise exceptions.JobError(
                            message=current_job.get("errorMessage"),
                            status=current_job.get("status"),
                            kind=current_job.get("kind"),
                        )
                    break
                else:
                    logging.debug("Notebook is still in progress...")

        return self

    def stop(self):
        self.properties = make_request(
            method="POST",
            path=f"{self.uri}/stop",
        )
        self.uri = self.properties["uri"]
        return self

    def referenced_tables(self):
        self.get()
        return [
            Table(source_table["qualifiedReference"], properties=source_table)
            for source_table in self.properties["referencedTables"]
        ]

    def source_tables(self):
        warnings.warn(
            get_warning("source_tables_deprecation"), FutureWarning, stacklevel=2
        )
        return self.referenced_tables()

    def source_table(self):
        self.get()
        return Table(
            self.properties["sourceTable"]["qualifiedReference"],
            properties=self.properties["sourceTable"],
        )

    def update(self, *, name=None, source_table=None):
        payload = {}

        if name is not None:
            payload["name"] = name
        if source_table is not None:
            if isinstance(source_table, Table):
                payload["sourceTable"] = source_table.qualified_reference
            else:
                payload["sourceTable"] = source_table

        response = make_request(
            method="PATCH",
            path=f"{self.uri}",
            payload=payload,
        )
        self.properties = response
        self.uri = self.properties["uri"]
        return self

    def output_table(self):
        self.get()
        return Table(
            self.properties["outputTable"]["qualifiedReference"],
            properties=self.properties["outputTable"],
        )

    def create_output_table(
        self, data=None, *, name=None, append=False, geography_variables=None
    ):
        if not (self.properties or {}).get("currentJob"):
            self.get()

        current_notebook_job_id = self.properties.get("currentJob")["id"]

        should_remove_tempfile = True
        temp_file_path = None
        try:
            if isinstance(data, str) or isinstance(data, pathlib.PurePath):
                data = str(data)
                if data.endswith(".parquet"):
                    should_remove_tempfile = False
                    temp_file_path = data
                elif data.endswith(".csv"):
                    import pandas

                    data = pandas.read_csv(data)
                else:
                    raise exceptions.ValueError(
                        "Only paths to parquet files (ending in .parquet) or CSV files (ending in .csv) are supported when a string argument is provided"
                    )

            if not isinstance(data, str):
                temp_file_path = convert_data_to_parquet(data)

            size = os.stat(temp_file_path).st_size
            pbar_bytes = tqdm(
                total=size, unit="B", leave=False, unit_scale=True, mininterval=0.1
            )

            res = make_request(
                method="POST",
                path=f"/notebookJobs/{current_notebook_job_id}/tempUploads",
                payload={"tempUploads": [{"size": size, "resumable": size > 5e7}]},
            )
            temp_upload = res["results"][0]

            with open(temp_file_path, "rb") as f:
                if temp_upload["resumable"]:
                    perform_resumable_upload(
                        data=f,
                        progressbar=pbar_bytes,
                        temp_upload_url=temp_upload["url"],
                    )
                else:
                    perform_standard_upload(
                        data=f,
                        temp_upload_url=temp_upload["url"],
                        progressbar=pbar_bytes,
                    )

            pbar_bytes.close()
            if should_remove_tempfile:
                os.remove(temp_file_path)

            res = make_request(
                method="PUT",
                path=f"/notebookJobs/{current_notebook_job_id}/outputTable",
                payload={
                    "name": name,
                    "append": append,
                    "geographyVariables": geography_variables,
                    "tempUploadId": temp_upload["id"],
                },
            )

            return Table(name=res["name"], properties=res)
        except Exception as e:
            if (
                temp_file_path
                and os.path.exists(temp_file_path)
                and should_remove_tempfile
            ):
                os.remove(temp_file_path)
            raise e
