import os
import re
import warnings
from enum import Enum
from typing import Union, List

import polars as pl
import holoviews as hv
from bokeh.models.formatters import BasicTickFormatter

import datashader as ds
import datashader.transfer_functions as tf

from genomeshader.genomeshader import *

warnings.simplefilter(action='ignore', category=FutureWarning)

hv.extension("bokeh")
hv.output(backend="bokeh")

warnings.resetwarnings()

base_colors = {
    'A': '#00F100',
    'C': '#341BFF',
    'G': '#D37D2B',
    'T': '#FF0030',
    'N': '#CCCCCC'
}


class GenomeBuild(Enum):
    GRCh38 = 'GRCh38'
    chm13v2_0 = 'chm13v2.0'


class GenomeShader:
    def __init__(self,
                 session_name: str = None,
                 gcs_session_dir: str = None,
                 genome_build: GenomeBuild = GenomeBuild.GRCh38):
        self._validate_session_name(session_name)
        self.session_name = session_name

        if gcs_session_dir is None:
            if 'GOOGLE_BUCKET' in os.environ:
                bucket = os.environ['GOOGLE_BUCKET']
                gcs_session_dir = f"{bucket}/GenomeShader/{session_name}"
            else:
                raise ValueError(
                    "gcs_session_dir is None and "
                    "GOOGLE_BUCKET is not set in environment variables"
                )

        self._validate_gcs_session_dir(gcs_session_dir)
        self.gcs_session_dir = gcs_session_dir

        self.genome_build: GenomeBuild = genome_build

        self._session = _init()

    def _validate_gcs_session_dir(self, gcs_session_dir: str):
        gcs_pattern = re.compile(
            r'^gs://[a-z0-9][a-z0-9._-]{1,61}[a-z0-9]/'  # bucket
            r'([^/]+/)*'  # folders (optional)
            r'[^/]*$'  # file (optional)
        )

        if not gcs_pattern.match(gcs_session_dir):
            raise ValueError("Invalid GCS path")

    def _validate_session_name(self, session_name: str):
        session_pattern = re.compile("^[a-zA-Z0-9_]+$")

        if not session_pattern.match(session_name):
            raise ValueError(
                "session_name contains special characters or whitespace"
            )

    def __str__(self):
        return (
            f'GenomeShader:\n'
            f' - session_name: {self.session_name}\n'
            f' - gcs_session_dir: {self.gcs_session_dir}\n'
            f' - genome_build: {self.genome_build}\n'
        )

    def get_version(self):
        """
        This function returns the version of the library.

        Returns:
            str: The version of the library.
        """
        return self._session.version()

    def get_session_name(self):
        """
        This function returns the name of the current session.

        Returns:
            str: The name of the current session.
        """
        return self.session_name

    def attach_reads(self, gcs_paths: Union[str, List[str]]):
        """
        This function attaches reads from the provided GCS paths to the
        current session. The GCS paths can be a single string or a list.
        Each GCS path can be a direct path to a .bam or .cram file, or a
        directory containing .bam and/or .cram files.

        Args:
            gcs_paths (Union[str, List[str]]): The GCS paths to attach reads.
        """
        if isinstance(gcs_paths, str):
            gcs_paths = [gcs_paths]  # Convert single string to list

        for gcs_path in gcs_paths:
            if gcs_path.endswith(".bam") or gcs_path.endswith(".cram"):
                self._session.attach_reads([gcs_path])
            else:
                bams = _gcs_list_files_of_type(gcs_path, ".bam")
                crams = _gcs_list_files_of_type(gcs_path, ".cram")

                self._session.attach_reads(bams)
                self._session.attach_reads(crams)

    def attach_loci(self, loci: Union[str, List[str]]):
        """
        Attaches loci to the current session from the provided list.
        The loci can be a single string or a list of strings.

        Args:
            loci (Union[str, List[str]]): Loci to be attached.
        """
        if isinstance(loci, str):
            self._session.attach_loci([loci])
        else:
            self._session.attach_loci(loci)

    def stage(self, use_cache: bool = True):
        """
        This function stages the current session. It prepares the session
        for subsequent operations like attaching reads or loci.
        """
        self._session.stage(use_cache)

    def get_locus(self, locus: str) -> pl.DataFrame:
        """
        This function retrieves the data for a staged locus from the
        current session.

        Args:
            locus (str): The locus to retrieve data for.

        Returns:
            pl.DataFrame: The data for the specified locus.
        """
        return self._session.get_locus(locus)

    def show(self,
             locus_or_dataframe: Union[str, pl.DataFrame],
             width: int = 980,
             height: int = 400,
             vertical: bool = False,
             expand: bool = False) -> hv.Rectangles:
        """
        Visualizes genomic data in a rectangular format.

        Args:
            locus_or_dataframe (str or DataFrame): Genomic locus to visualize.
            Can either be specified as a locus string (e.g. 'chr:start-stop')
            or a Polars DataFrame (usually from the get_locus() method,
            optionally modified by the user).
            width (int, optional): Visualization width. Defaults to 980.
            height (int, optional): Visualization height. Defaults to 400.
            expand (bool, optional): If True, expands each sample to show all
            reads. Defaults to False.

        Returns:
            hv.Rectangles: HoloViews Rectangles object for visualization.
        """

        if isinstance(locus_or_dataframe, str):
            df = self.get_locus(locus_or_dataframe)
        elif isinstance(locus_or_dataframe, pl.DataFrame):
            df = locus_or_dataframe
        else:
            raise ValueError("locus_or_dataframe must be a locus string or a"
                             "Polars DataFrame.")

        y0s = []
        y0 = 0
        if expand:
            query_name = None
            for row in df.iter_rows(named=True):
                if query_name is None:
                    query_name = row['query_name']
                    y0 = 0

                if query_name != row['query_name']:
                    query_name = row['query_name']
                    y0 += 1

                y0s.append(y0)
        else:
            sample_name = None
            for row in df.iter_rows(named=True):
                if sample_name is None:
                    sample_name = row['sample_name']
                    y0 = 0

                if sample_name != row['sample_name']:
                    sample_name = row['sample_name']
                    y0 += 1

                y0s.append(y0)

        df = df.with_columns(pl.Series(name="read_num", values=y0s))
        df = df.with_columns(pl.Series(name="height", values=[1.0]*len(y0s)))

        df = df.with_columns(
            pl.col("read_num").alias("y0") * -1 - pl.col("height") / 2
        )
        df = df.with_columns(
            pl.col("read_num").alias("y1") * -1 + pl.col("height") / 2
        )

        boxes = hv.Rectangles((
            list(df["reference_start"]),
            list(df["y0"]),
            list(df["reference_end"]),
            list(df["y1"]),
            list(df["element_type"])
        ), vdims="element_type")

        return boxes.opts(
            width=width,
            height=height,

            xlabel="",
            xformatter=BasicTickFormatter(use_scientific=False),
            yaxis=None,
            ylim=(-18, 1),

            title="Read visualization",
            fontscale=1.3,
            color="element_type",

            line_width=0,

            tools=['xwheel_zoom', 'ywheel_zoom', 'pan'],
            active_tools=['xwheel_zoom', 'pan'],
            default_tools=['reset', 'save']
        )

    def reset(self):
        self._session.reset()

    def print(self):
        self._session.print()


def init(session_name: str,
         gcs_session_dir: str = None,
         genome_build: GenomeBuild = GenomeBuild.GRCh38) -> GenomeShader:
    session = GenomeShader(session_name=session_name,
                           gcs_session_dir=gcs_session_dir,
                           genome_build=genome_build)

    return session
