import pandas as pd
import requests
import re
import numpy as np
import typing
import os

from io import BytesIO

from ovo.app.components.molstar_custom_component.color_picker import ColorPicker, mix_colors
from ovo.app.components.molstar_custom_component.dataclasses import (
    ContigSegment,
)


class ContigsParser:
    """Class for parsing the contigs definition."""

    def get_chain_and_indices_fixed_region(self, region: str):
        """Get the chain and region residue indices.
        e.g.: From A123-131 return (A, 123, 131)."""
        chain_inp = "".join(re.findall(r"[a-zA-Z]+", region[0]))
        first_idx_inp = "".join(re.findall(r"[0-9]+", region[0]))
        if len(region) > 1:
            end_idx_inp = "".join(re.findall(r"[0-9]+", region[1]))
        else:
            end_idx_inp = first_idx_inp

        assert len(chain_inp) <= 2, f"Chains must have a maximum of 2 characters. {region}"
        assert first_idx_inp != "", f"Expected to find a number value for the contig {region}."
        assert int(first_idx_inp) <= int(end_idx_inp), f"Expected {first_idx_inp} <= {end_idx_inp} for {region}."
        assert len(region) < 3, f"Expected just 1-2 numbers in {region}"

        return chain_inp, first_idx_inp, end_idx_inp

    def get_output_chain(self, has_generated_before_chain_break: bool):
        """Get the output chain. Currently, we assume that anything with a "generated region" until there is a chain break
        results in chain A, otherwise returns chain B."""
        return "A" if has_generated_before_chain_break else "B"

    def parse_region_generated(self, region: str, last_parsed_res: int, has_generated_before_chain_break: bool):
        """Parse regions that are defined to be generated by RFdiffusion."""
        split_region = region.split("-")
        out_res_start = last_parsed_res + 1

        if len(split_region) == 1:
            split_region.append(split_region[0])

        assert len(split_region) <= 2 and len(split_region) > 0, f"Expected one or two coordinates, got: {region}"
        # we do not want the case of 10-20
        assert split_region[0] == split_region[1], f"Expected fixed region length, got: {region}"

        # without postprocessing it is not possible to guess the final boundaries
        out_res_end = out_res_start + int(split_region[0]) - 1

        return ContigSegment(
            value=region,
            length=int(split_region[0]),
            type="generated",
            input_res_start=-1,  # we do not care about input here as it does not exist
            input_res_end=-1,
            input_res_chain="?",
            out_res_start=last_parsed_res + 1,
            out_res_end=out_res_end,
            out_res_chain=self.get_output_chain(has_generated_before_chain_break),
        )

    def parse_region_fixed_no_mapping(self, region: str):
        """Parse fixed regions with an identity mapping."""
        split_region = region.split("-")
        first_chain_inp, first_idx_inp, end_idx_inp = self.get_chain_and_indices_fixed_region(split_region)

        return ContigSegment(
            value=region,
            type="fixed",
            length=int(end_idx_inp) - int(first_idx_inp) + 1,
            input_res_start=int(first_idx_inp),
            input_res_end=int(end_idx_inp),
            input_res_chain=first_chain_inp,  # should be OK as long as we assume there are no things such as A100-B150
            out_res_start=int(first_idx_inp),
            out_res_end=int(end_idx_inp),
            out_res_chain=first_chain_inp,  # should be OK as long as chain breaks are used properly
        )

    def parse_region_fixed(
        self,
        region: str,
        ref_indices: list[list[typing.Union[str, np.int64, bool]]],
        new_indices: list[list[typing.Union[str, np.int64, bool]]],
        has_generated_before_chain_break: bool,
    ):
        """Parse fixed regions with corresponding mapping defined in the .trb file."""
        split_region = region.split("-")
        first_chain_inp, first_idx_inp, end_idx_inp = self.get_chain_and_indices_fixed_region(split_region)

        # now let's find the corresponding values in the reference/new indices list (otherwise return -1)
        ref_indices_first_index = next(
            (
                i
                for i, item in enumerate(ref_indices)
                if item[0] == first_chain_inp and item[1] == np.int64(first_idx_inp) and item[2] == False
            ),
            -1,
        )

        ref_indices_last_index = next(
            (
                i
                for i, item in enumerate(ref_indices)
                if item[0] == first_chain_inp and item[1] == np.int64(end_idx_inp) and item[2] == False
            ),
            -1,
        )

        assert ref_indices_first_index > -1, f"Could not find mapping for residue {first_chain_inp}{first_idx_inp}"
        assert ref_indices_last_index > -1, f"Could not find mapping for residue {first_chain_inp}{end_idx_inp}"

        # Set the flag to already used.
        new_indices[ref_indices_first_index][2] = True
        new_indices[ref_indices_last_index][2] = True

        ref_indices[ref_indices_first_index][2] = True
        ref_indices[ref_indices_last_index][2] = True

        # now we have the indices so we might go on and use them to find the mapping
        # the mapping will use the same indices
        first_res_mapping = new_indices[ref_indices_first_index]
        last_res_mapping = new_indices[ref_indices_last_index]

        return ContigSegment(
            value=region,
            type="fixed",
            length=int(end_idx_inp) - int(first_idx_inp) + 1,
            input_res_start=int(first_idx_inp),
            input_res_end=int(end_idx_inp),
            input_res_chain=first_chain_inp,  # should be OK as long as we assume there are no things such as A100-B150
            out_res_start=int(first_res_mapping[1]),
            out_res_end=int(last_res_mapping[1]),
            out_res_chain=self.get_output_chain(
                has_generated_before_chain_break
            ),  # should be OK as long as chain breaks are used properly
        )

    def parse_contigs_str(self, contigs_string: str = "", include_generated=False) -> list[ContigSegment]:
        return self._parse_contigs_with_tuples(contigs=contigs_string, include_generated=include_generated)

    def parse_contigs_trb(self, contigs_trb: str | dict = "") -> list[ContigSegment]:
        """Parse the contigs from a given .trb file or .trb parsed dictionary.
        Can be either an http(s) URL, a system path to the .trb file or parsed .trb file content directly (checked in this order).
        """
        # notice that in this case we are using the contigs from the .trb file, not the original ones that were defined by the user
        # first, we check if "trb" is an URL, then we check for a file and lastly we treat "trb" as the unpacked pickle
        if "http://" in contigs_trb or "https://" in contigs_trb:
            r = requests.get(contigs_trb)
            trb_file = pd.read_pickle(BytesIO(r.content))
        elif (isinstance(contigs_trb, str) or isinstance(contigs_trb, os.PathLike)) and os.path.exists(contigs_trb):
            trb_file = pd.read_pickle(contigs_trb)
        elif isinstance(contigs_trb, dict):
            trb_file = contigs_trb
        else:
            raise RuntimeError(
                f"Provided an invalid .trb definition. Cannot parse: {contigs_trb}, {type(contigs_trb)}."
            )

        # contigs: str = trb_file["config"]["contigmap"]["contigs"][0].replace(" ", "/0 ")
        contigsList: list[str] = trb_file["sampled_mask"]
        newContigs = []

        for idx, contig in enumerate(contigsList):
            contig = contig.removesuffix("/0")
            newContigs.append(contig)

        contigs = "/0 ".join(newContigs)

        if "complex_con_ref_pdb_idx" in trb_file:
            ref_idx = trb_file["complex_con_ref_pdb_idx"]
        else:
            assert "con_ref_pdb_idx" in trb_file, "Could not find con_ref_pdb_idx in the .trb file"
            ref_idx = trb_file["con_ref_pdb_idx"]

        if "complex_con_hal_pdb_idx" in trb_file:
            hal_idx = trb_file["complex_con_hal_pdb_idx"]
        else:
            assert "con_hal_pdb_idx" in trb_file, "Could not find con_hal_pdb_idx in the .trb file"
            hal_idx = trb_file["con_hal_pdb_idx"]

        return self.parse_contigs_ref(contigs, ref_idx=ref_idx, hal_idx=hal_idx)

    def parse_contigs_ref(self, contigs: str, ref_idx: list[tuple], hal_idx: list[tuple]):
        """Parse the contigs and return segments with output numbering."""
        reference_tuples = []
        for tuple in ref_idx:
            reference_tuples.append([tuple[0], tuple[1], False])

        mapped_tuples = []
        for tuple in hal_idx:
            mapped_tuples.append([tuple[0], tuple[1], False])

        return self._parse_contigs_with_tuples(
            contigs=contigs, mapped_tuples=mapped_tuples, reference_tuples=reference_tuples, include_generated=True
        )

    def _parse_contigs_with_tuples(
        self, contigs: str, mapped_tuples: list = [], reference_tuples: list = [], include_generated=True
    ) -> list[ContigSegment]:
        """
        Parse the contigs from the .trb file. If no .trb file path is defined, then the parser treats fixed regions with no mapping.

        Parameters
        ----------
            contigs: str
                Contigs string definition (e.g.: A30-54/10/A70-75).
            mapped_tuples: list
                Mapped residues from the .trb contigs definition.
            reference_tuples: list
                Mapped reference residues from the .trb contigs definition.

        Returns
        -------
            Parsed contig regions
        """
        if not contigs:
            return []

        parsed_segments: list[ContigSegment] = []
        color_picker = ColorPicker()

        last_parsed_residue = 0
        has_generated_before_chain_break = False

        # Here, we check if there is a generated segment in each "bigger region" separated by a chain break.
        # For example, for these regions    ['A57-67/10-10/A34-42/10-10/A90-102/10-10/A126-137', 'H1-50', 'L1-114', 'A50-67/20-20/A70-77', 'H60-80']
        # we expect these results:          [True, False, False, True, False]
        split_regions_by_chainbreak = contigs.split("/0 ")
        regions_contain_generated = []
        for split_reg in split_regions_by_chainbreak:
            found = False

            for seg in split_reg.split("/"):
                seg = seg.replace("0 ", "").replace(" ", "")

                if not any(c.isalpha() for c in seg):  # this means there is no letter -> a generated specification
                    found = True
                    break

            regions_contain_generated.append(found)

        chainbreak_idx = 0
        has_generated_before_chain_break = regions_contain_generated[chainbreak_idx]

        split_segments = contigs.split("/")
        for segment in split_segments:
            if segment.startswith("0 "):
                # For each of the bigger regions, we have to check if there was a generated segment inside to
                chainbreak_idx += 1
                has_generated_before_chain_break = regions_contain_generated[chainbreak_idx]

                last_parsed_residue = 0

                new_segment_split = segment.split(" ")
                assert len(new_segment_split) == 2, f"Expected a format of '/0 <chain>..-..', got {segment}"

                segment = new_segment_split[1]

            if not any(c.isalpha() for c in segment):  # this means there is no letter -> a generated specification
                if not include_generated:
                    # for contigs defined in a string, we do not parse the generated regions
                    continue
                parsed_segment = self.parse_region_generated(
                    segment, last_parsed_residue, has_generated_before_chain_break
                )
            else:
                if len(reference_tuples) == 0:
                    # there is no mapping in this case
                    parsed_segment = self.parse_region_fixed_no_mapping(segment)
                else:
                    parsed_segment = self.parse_region_fixed(
                        segment, reference_tuples, mapped_tuples, has_generated_before_chain_break
                    )

            last_parsed_residue = parsed_segment.out_res_end
            parsed_segments.append(parsed_segment)

        # assign colors on segments sorted in original order, to get a consistent coloring independent of reordering
        sorted_segments = sorted(parsed_segments, key=lambda x: (x.input_res_chain, x.input_res_start, x.out_res_start))
        for seg in sorted_segments:
            if seg.type == "generated":
                seg.color = "#44aa44"  # fallback color - will be changed below
            else:
                seg.color = color_picker((seg.input_res_start, seg.input_res_end, seg.input_res_chain))

        # assign color to generated segments based on neighboring fixed segments
        for prev_seg, seg, next_seg in zip(
            [None] + parsed_segments[:-1], parsed_segments, parsed_segments[1:] + [None]
        ):
            if seg.type == "generated":
                if prev_seg:
                    # mid generated segment
                    seg.color = mix_colors(mix_colors(prev_seg.color, "#ffffff"), "#ffffff")
                elif next_seg:
                    # start segment (N term)
                    seg.color = "#eeeeee"

        return parsed_segments
