import hashlib
import numpy as np
import os
import re
import requests
import tarfile
import time

from absl import logging
from typing import List, NoReturn, Tuple


class MMSeqs2Runner:

    r"""Runner object

    Fetches sequence alignment and templates from MMSeqs2 server
    Based on the function run_mmseqs2 from ColabFold (sokrypton/ColabFold)
    Version 62d7558c91a9809712b022faf9d91d8b183c328c

    Relevant publications
    ----------
    * "Clustering huge protein sequence sets in linear time"
      https://doi.org/10.1038/s41467-018-04964-5
    * "MMseqs2 enables sensitive protein sequence searching for the analysis
      of massive data sets"
      https://doi.org/10.1038/nbt.3988

    Private variables
    ----------
    self.job: Job ID (five-char string)
    self.seq: Sequence to search
    self.host_url: URL address to ping for data
    self.t_url: URL address to ping for templates from PDB
    self.n_templates = Number of templates to fetch (default=20)
    self.path: Path to use
    self.tarfile: Compressed file archive to download
    """

    def __init__(
        self,
        job: str,
        seq: str,
        host_url: str = "https://a3m.mmseqs.com",
        t_url: str = "https://a3m-templates.mmseqs.com/template",
        path_suffix: str = "env",
        n_templates: int = 20,
    ):

        r"""Initialize runner object

        Parameters
        ----------
        job : Job name
        seq : Amino acid sequence
        host_url : Website to ping for sequence data
        t_url : Website to ping for template info
        path_suffix : Suffix for path info

        """

        # Clean up sequence
        self.seq = self._cleanseq(seq.upper())

        # Come up with unique job ID for MMSeqs
        self.job = self._define_jobname(job)

        # Save everything else
        self.host_url = host_url
        self.t_url = t_url
        self.n_templates = n_templates

        self.path = "_".join((self.job, path_suffix))

        if not os.path.isdir(self.path):
            os.system(f"mkdir { self.path }")

        self.tarfile = f"{ self.path }/out.tar.gz"

    def _cleanseq(self, seq) -> str:

        r"""Cleans the sequence to remove whitespace and noncanonical letters

        Parameters
        ----------
        seq : Amino acid sequence (only all 20 here)

        Returns
        ----------
        Cleaned up amin acid sequence

        """

        if any([aa in seq for aa in "BJOUXZ"]):
            logging.warning("Sequence contains non-canonical amino acids!")
            logging.warning("Removing B, J, O, U, X, and Z from sequence")
            seq = re.sub(r"[BJOUXZ]", "", seq)

        return re.sub(r"[^A-Z]", "", "".join(seq.split()))

    def _define_jobname(self, job: str) -> str:

        r"""Provides a unique five-digit identifier for the job name

        Parameters
        ----------
        job : Job name

        Returns
        ----------
        Defined job name

        """

        return "_".join(
            (
                re.sub(r"\W+", "", "".join(job.split())),
                hashlib.sha1(self.seq.encode()).hexdigest()[:5],
            )
        )

    def _submit(self) -> dict:

        r"""Submit job to MMSeqs2 server

        Parameters
        ----------
        None

        Returns
        ----------
        None

        """

        data = {"q": f">101\n{ self.seq }", "mode": "env"}

        res = requests.post(f"{ self.host_url }/ticket/msa", data=data)

        try:
            out = res.json()

        except ValueError:
            out = {"status": "UNKNOWN"}

        return out

    def _status(self, idx: str) -> dict:

        r"""Check status of job

        Parameters
        ----------
        idx : Index assigned by MMSeqs2 server

        Returns
        ----------
        None

        """

        res = requests.get(f"{ self.host_url }/ticket/{ idx }")

        try:
            out = res.json()

        except ValueError:
            out = {"status": "UNKNOWN"}

        return out

    def _download(self, idx: str, path: str) -> NoReturn:

        r"""Download job outputs

        Parameters
        ----------
        idx : Index assigned by MMSeqs2 server
        path : Path to download data

        Returns
        ----------
        None

        """

        res = requests.get(f"{ self.host_url }/result/download/{ idx }")

        with open(path, "wb") as out:
            out.write(res.content)

    def _search_mmseqs2(self) -> NoReturn:

        r"""Run the search and download results
        Heavily modified from ColabFold

        Parameters
        ----------
        None

        Returns
        ----------
        None

        """

        if os.path.isfile(self.tarfile):
            return

        out = self._submit()

        time.sleep(5 + np.random.randint(0, 5))
        while out["status"] in ["UNKNOWN", "RATELIMIT"]:
            # resubmit
            time.sleep(5 + np.random.randint(0, 5))
            out = self._submit()

        logging.debug(f"ID: { out[ 'id' ] }")

        while out["status"] in ["UNKNOWN", "RUNNING", "PENDING"]:
            time.sleep(5 + np.random.randint(0, 5))
            out = self._status(out["id"])

        if out["status"] == "COMPLETE":
            self._download(out["id"], self.tarfile)

        elif out["status"] == "ERROR":
            raise RuntimeError(
                " ".join(
                    (
                        "MMseqs2 API is giving errors.",
                        "Please confirm your input is a valid protein sequence.",
                        "If error persists, please try again in an hour.",
                    )
                )
            )

    def process_templates(self, templates: List[str] = []) -> str:

        r"""Process templates and fetch from MMSeqs2 server

        Parameters
        ----------
        use_templates : True/False whether to use templates
        max_templates : Maximum number of templates to use

        Returns
        ----------
        Directory containing templates (empty if not using templates)

        """

        path = f"{ self.job }_env/templates_101"
        if os.path.isdir(path):
            os.system(f"rm { path }")

        # templates = {}
        logging.info("\t".join(("seq", "pdb", "cid", "evalue")))

        pdbs = []
        with open(f"{ self.path }/pdb70.m8", "r") as infile:

            for line in infile:

                sl = line.rstrip().split()
                pdb = sl[1]
                if pdb in templates:
                    pdbs.append(sl[1])
                    logging.info(f"{ sl[0] }\t{ sl[1] }\t{ sl[2] }\t{ sl[10] }")

        if len(pdbs) == 0:
            logging.warning("No templates found.")
            return ""

        else:

            if not os.path.isdir(path):
                os.mkdir(path)

            pdbs = [t for t in pdbs if t in templates]

            if len(templates) == 0 or len(pdbs) == 0:
                pdbs = ",".join(templates[: self.n_templates])
            else:
                pdbs = ",".join(pdbs[: self.n_templates])

            os.system(f"curl -v { self.t_url }/{ pdbs } |tar xzf - -C { path }/")

            os.system(f"cp { path }/pdb70_a3m.ffindex { path }/pdb70_cs219.ffindex")

            os.system(f"touch { path }/pdb70_cs219.ffdata")

            return path

    def _process_alignment(
        self, a3m_files: list, templates: List[str] = []
    ) -> Tuple[str, str]:

        r"""Process sequence alignment
        (modified from ColabFold)

        Parameters
        ----------
        a3m_files : List of files to parse
        token : Token to look for when parsing

        Returns
        ----------
        Tuple with [0] string with alignment, and [1] path to template

        """

        a3m_lines = ""

        for a3m_file in a3m_files:
            for line in open(os.path.join(self.path, a3m_file), "r"):
                if len(line) > 0:
                    a3m_lines += line.replace("\x00", "")

        return a3m_lines, self.process_templates(templates)

    def run_job(self, templates: List[str] = []) -> Tuple[str, str]:

        r"""
        Run sequence alignments using MMseqs2

        Parameters
        ----------
        use_templates: Whether to use templates

        Returns
        ----------
        Tuple with [0] string with alignment, and [1] path to template

        """

        self._search_mmseqs2()

        a3m_files = ["uniref.a3m", "bfd.mgnify30.metaeuk30.smag30.a3m"]

        # extract a3m files
        if not os.path.isfile(os.path.join(self.path, a3m_files[0])):
            with tarfile.open(self.tarfile) as tar_gz:
                tar_gz.extractall(self.path)

        return self._process_alignment(a3m_files, templates)