import http.server
import logging
import textwrap
from itertools import groupby
from operator import itemgetter
from pathlib import Path
from typing import Dict, Union

import numpy as np
from Bio.SeqUtils import seq1
from jinja2 import Environment, FileSystemLoader

from abcfold.output.alphafold3 import AlphafoldOutput
from abcfold.output.boltz import BoltzOutput
from abcfold.output.chai import ChaiOutput
from abcfold.output.file_handlers import ConfidenceJsonFile, NpzFile
from abcfold.output.protenix import ProtenixOutput
from abcfold.plots.pae_plot import create_pae_plots
from abcfold.plots.plddt_plot import plot_plddt
from abcfold.scripts.ipsae import Ipsae

PORT = 8000
logger = logging.getLogger("logger")


def get_plddt_regions(plddts: Union[np.ndarray, list]) -> dict:
    """
    Get the pLDDT regions for the model
    """
    if not isinstance(plddts, np.ndarray):
        plddts = np.array(plddts)

    regions = {}
    # replace none values with -1
    plddts = np.where(plddts == None, -1, plddts) # noqa F401

    v_low = np.where((0 <= plddts) & (plddts <= 50))[0]
    regions["v_low"] = get_regions_helper(v_low)
    low = np.where((plddts > 50) & (plddts < 70))[0]
    regions["low"] = get_regions_helper(low)
    confident = np.where((plddts >= 70) & (plddts < 90))[0]
    regions["confident"] = get_regions_helper(confident)
    v_confident = np.where(plddts >= 90)[0]
    regions["v_high"] = get_regions_helper(v_confident)

    return regions


def get_regions_helper(indices):
    """
    Get the regions from the indices
    """
    regions = []
    for _, g in groupby(enumerate(indices), lambda x: x[0] - x[1]):
        group = map(itemgetter(1), g)
        group = list(map(int, group))
        regions.append((group[0], group[-1]))
    return regions


def get_model_sequence_data(cif_objs) -> dict:
    """
    Get the sequence for each chain and ligand in the model, used internally
    for plotting

    Args:
        cif_objs : A list of CifFile objs

    Returns:
        dict : Chain ID and sequence data
    """
    sequence_data: dict = {}
    for cif_obj in cif_objs:
        sequence_data_ = {}
        for chain in cif_obj.get_chains():
            if cif_obj.check_ligand(chain):
                if chain.id not in sequence_data_:
                    sequence_data_[chain.id] = ""
                sequence_data_[chain.id] += "".join(
                    [atom.id[0] for residue in chain for atom in residue]
                )
            elif cif_obj.check_other(chain, ["dna"]):
                sequence_data_[chain.id] = "".join(
                    [residue.get_resname()[-1] for residue in chain]
                )
            elif cif_obj.check_other(chain, ["rna"]):
                sequence_data_[chain.id] = "".join(
                    [residue.get_resname() for residue in chain]
                )
            else:

                sequence_data_[chain.id] = "".join(
                    [seq1(residue.get_resname()) for residue in chain]
                )
        sequence_data = {
            chain_id: sorted(
                [sequence_data_[chain_id], sequence_data.get(chain_id, "")],
                reverse=True,
            )[0]
            for chain_id in sequence_data_
        }

    return sequence_data


def get_model_data(model,
                   plot_dict,
                   method,
                   plddt_scores,
                   pae_obj,
                   score_file,
                   output_dir):
    """
    Get the model data for the output page

    Args:
        model (CifFile): Model object
        plot_dict (dict): Dictionary of plots
        method (str): Method used to generate the model
        score_file (str): Path to the file containing model scores
        pae_obj (ConfidenceJsonFile): Obj containing PAE data
        output_dir (Path): Path to the output directory
    """
    regions = get_plddt_regions(plddt_scores)
    ptm_score, iptm_score = parse_scores(score_file)
    full_model_path = Path(model.pathway)
    model_path = full_model_path.relative_to(output_dir).as_posix()

    ipsae_score = []
    try:
        ipsae = Ipsae(model, pae_obj)
        distances = ipsae.get_cb_distance()
        ipsae_scores = ipsae.compute_iptm_ipsae(distances)

        # Currently displaying all chain variations, might alter to max in future
        for k, v in ipsae_scores["ipsae_d0res_asym"].items():
            for k2, v2 in v.items():
                ipsae_score.append(f"{k}{k2}:{np.round(v2, 4)}")

        # Output full ipsae scores for each model as we are only outputting d0res_asym
        ipsae_out = full_model_path.parent / f"{Path(model_path).stem}_ipsae.csv"
        ipsae.output_results(pdockq_scores=None,
                             pdockq2_scores=None,
                             lis_scores=None,
                             ipsae_scores=ipsae_scores,
                             verbose=False,
                             output_csv=ipsae_out)
    except ValueError:
        logger.error("ValueError when calculating ipSAE score, bypassing ipSAE")
        pass

    model_data = {
        "model_id": model.name,
        "model_source": method,
        "model_path": model_path,
        "plddt_regions": regions,
        "avg_plddt": model.average_plddt,
        "h_score": model.h_score,
        "ptm_score": ptm_score,
        "iptm_score": iptm_score,
        "ipsae_score": ",".join(ipsae_score),
        "residue_clashes": model.clashes_residues,
        "atom_clashes": model.clashes,
        "pae_path": Path(plot_dict[model.pathway.as_posix()])
        .relative_to(output_dir)
        .as_posix(),
    }
    return model_data


class NoCacheHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
    def end_headers(self):
        self.send_header(
            "Cache-Control", "no-store, no-cache, must-revalidate, max-age=0"
        )
        self.send_header("Pragma", "no-cache")
        self.send_header("Expires", "0")
        super().end_headers()


def plots(outputs: list, output_dir: Path):
    """
    Generate plots for the output of the different programs

    Args:
        outputs (list): List of output objects

    """
    pathway_plots = create_pae_plots(outputs, output_dir=output_dir)
    plddt_plot_input: Dict[str, list] = get_all_cif_files(outputs)

    plot_plddt(plddt_plot_input, output_name=output_dir.joinpath("plddt_plot.html"))

    pathway_plots["plddt"] = str(output_dir.joinpath("plddt_plot.html").resolve())

    return pathway_plots


def render_template(in_file_path, out_file_path, **kwargs):
    """
    Templates the given file with the keyword arguments.

    Args:
        in_file_path (Path): The path to the template.
        out_file_path (Path): The path to output the templated file.
        **kwargs (dict): Variables to use in templating.
    """
    env = Environment(
        loader=FileSystemLoader(in_file_path.parent), keep_trailing_newline=True
    )
    template = env.get_template(in_file_path.name)
    output = template.render(**kwargs)
    with open(str(out_file_path), "w") as f:
        f.write(output)


def output_open_html_script(file_out: str, port: int = 8000):
    """
    Make a python script to open the output HTML file in the default web browser

    Args:
        file_out (str): Path to the output script
        port (int): Port to run the server on
    """

    script = f"""
    import http.server
    import socketserver
    import webbrowser
    import sys

    PORT = {port}

    class NoCacheHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
        def end_headers(self):
            self.send_header("Cache-Control",
                            "no-store, no-cache, must-revalidate, max-age=0")
            self.send_header("Pragma", "no-cache")
            self.send_header("Expires", "0")
            super().end_headers()

    try:
        with socketserver.TCPServer(("", PORT),
                                    NoCacheHTTPRequestHandler) as httpd:
            print(
                f"Serving at port {PORT}: http://localhost:{PORT}/index.html"
                )
            print("Press Ctrl+C to stop the server")
            webbrowser.open(f"http://localhost:{PORT}/index.html")
            httpd.serve_forever()
    except KeyboardInterrupt:
        print("Server stopped")
        httpd.server_close()
        sys.exit(0)
    """

    script = textwrap.dedent(script)
    with open(file_out, "w") as f:
        f.write(script)


def get_all_cif_files(outputs) -> Dict[str, list]:
    method_cif_objs: Dict[str, list] = {}

    for output in outputs:
        if isinstance(output, AlphafoldOutput):
            for seed in output.seeds:
                if "Alphafold3" not in method_cif_objs:
                    method_cif_objs["Alphafold3"] = []
                method_cif_objs["Alphafold3"].extend(output.cif_files[seed])
        elif isinstance(output, BoltzOutput):
            for seed in output.seeds:
                if "Boltz" not in method_cif_objs:
                    method_cif_objs["Boltz"] = []
                method_cif_objs["Boltz"].extend(output.cif_files[seed])
        elif isinstance(output, ChaiOutput):
            for seed in output.seeds:
                if "Chai-1" not in method_cif_objs:
                    method_cif_objs["Chai-1"] = []
                method_cif_objs["Chai-1"].extend(output.cif_files[seed])
        elif isinstance(output, ProtenixOutput):
            for seed in output.seeds:
                if "Protenix" not in method_cif_objs:
                    method_cif_objs["Protenix"] = []
                method_cif_objs["Protenix"].extend(output.cif_files[seed])

    return method_cif_objs


def parse_scores(score_file: Union[ConfidenceJsonFile, NpzFile]) -> tuple:
    """
    Parse the scores from the score file.

    Args:
        score_file (Union[ConfidenceJsonFile, NpzFile]): The score file object.

    Returns:
        tuple: A tuple containing ptm_score and iptm_score as floats, or None if invalid
    """
    ptm_score = None
    iptm_score = None

    if isinstance(score_file, ConfidenceJsonFile):
        data = score_file.load_json_file()
    elif isinstance(score_file, NpzFile):
        data = score_file.load_npz_file()
    else:
        return ptm_score, iptm_score
    for key in ("ptm", "iptm"):
        try:
            value = float(data[key])
            if key == "ptm":
                ptm_score = round(value, 2)
            else:
                iptm_score = round(value, 2)
        except (KeyError, TypeError, ValueError):
            continue

    return ptm_score, iptm_score
