import os
import sys
import orjson
import gzip
import json
import tempfile
from typing import TextIO
from argparse import ArgumentParser

from eltetrado.analysis import eltetrado
from eltetrado.model import generate_dto
from eltetrado.structure import read_3d_structure, read_2d_structure

import drawtetrado.structure as structure
import drawtetrado.svg_painter as svg_painter

def Draw(struct, output_file, config = svg_painter.Config(1.0)):
    if len(struct.tetrads) == 0:
        print("No tetrads available in the processed structure!")
    for idx in range(len(struct.tetrads)):
        quadruplex = structure.Quadruplex(struct, idx)
        svg_maker = svg_painter.SvgMaker(config, output_file + "_" + str(idx) + ".svg", \
                quadruplex)

        # OPTIMIZE, Takes argument "optimizer" with location to the optimizer
        # binary. Default is "./svg_optimizer"
        quadruplex.Optimize()

        # Prepare + Draw
        svg_maker.DrawAll()

        # Save
        print("Helix " + str(idx) + " full: " + output_file + "_" + str(idx) + ".svg")
        svg_maker.svg.save(pretty=True)

        # Also draw single tetrads from quadruplex structure.
        if (len(struct.single_tetrads[idx]) > 1):
            for tetrad_idx in range(len(struct.single_tetrads[idx])):
                quadruplex = structure.Quadruplex(struct, idx, tetrad_idx)
                svg_maker = svg_painter.SvgMaker(config, output_file + "_" + str(idx) + \
                        "_" + str(tetrad_idx) + ".svg", quadruplex)

                # OPTIMIZE, Takes argument "optimizer" with location to the optimizer
                # binary. Default is "./svg_optimizer"
                quadruplex.Optimize()

                # Prepare + Draw
                svg_maker.DrawAll()

                # Save
                print("Helix " + str(idx) + ", Quadruplex " + str(tetrad_idx) + ": " + output_file + "_" + str(idx) + ".svg")
                svg_maker.svg.save(pretty=True)

def DrawFromString(json, output_file, config = svg_painter.Config(1.0)):
    Draw(structure.Structure().fromString(json), output_file, config)


def DrawFromFile(filename_json, output_file, config = svg_painter.Config(1.0)):
    Draw(structure.Structure().fromFile(filename_json), output_file, config)

# For ElTetrado
def handle_input_file(path) -> TextIO:
    root, ext = os.path.splitext(path)

    if ext == '.gz':
        root, ext = os.path.splitext(root)
        file = tempfile.NamedTemporaryFile('w+', suffix=ext)
        with gzip.open(path, 'rt') as f:
            file.write(f.read())
            file.seek(0)
    else:
        file = tempfile.NamedTemporaryFile('w+', suffix=ext)
        with open(path) as f:
            file.write(f.read())
            file.seek(0)
    return file

def IsFileJson(path):
    with open(path) as file:
        try:
            data = json.load(file)
            return True
        except:
            print("Provided file is not a .json. Processing file using ElTetrado.")
            return False
    return False

def main():
    parser = ArgumentParser('drawtetrado',
        epilog='The output path is a template. Program will generate drawings for each '
        'individual quadruplex in each helix as well as each full helix.'
        'If we have file with 2 helices with 2 quadruplexes each, if no output_template value '
        'is provided, path and basename of the input file will be used. If it is provided, output_template=/tmp/out, '
        'the resulting files will be /tmp/out_0.svg, with full helix 0, /tmp/out_0_0.svg '
        'and /tmp/out_0_1.svg for each quadruplex in helice 0. Similar files will be created '
        'for helix 1 with /tmp/out_1.svg and /tmp/out_1_0.svg, /tmp/out_1_1.svg.')
    parser.add_argument('-i', '--input', help='path to input PDB, PDBx/mmCIF file or JSON generated by ElTetrado. '
            'If PDB or PDBx/mmCIF file is provided, it will be first analyzed using ElTetrado.')
    parser.add_argument('-o', '--output-template', help='(optional) path to output SVG file template '
            '[default=input file path and basename]')
    parser.add_argument('--config',
            help='(optional) JSON config file containing all parameter changes and individual nucleotide '
            'coloring overrides',
            default=None)
    # ElTetrado options.
    parser.add_argument('-m', '--model', help='(optional, ElTetrado) model number to process', default=1, type=int)
    parser.add_argument('--stacking-mismatch',
                        help='(optional, ElTetrado) a perfect tetrad stacking covers 4 nucleotides; this option can be used with value 1 or '
                             '2 to allow this number of nucleotides to be non-stacked with otherwise well aligned '
                             'tetrad [default=2]',
                        default=2,
                        type=int)
    parser.add_argument('--strict',
                        action='store_true',
                        help='(optional, ElTetrado) nucleotides in tetrad are found when linked only by cWH pairing')
    parser.add_argument('--no-reorder',
                        action='store_true',
                        help='(optional, ElTetrado) chains of bi- and tetramolecular quadruplexes should be reordered to be able to have '
                             'them classified; when this is set, chains will be processed in original order, which for '
                             'bi-/tetramolecular means that they will likely be misclassified; use with care!')

    args = parser.parse_args()

    config = svg_painter.Config(1.0, args.config)

    if not args.input:
        print(parser.print_help())
        sys.exit(1)

    root, ext = os.path.splitext(args.input)
    if IsFileJson(args.input):
        if  not args.output_template:
            DrawFromFile(args.input, root, config)
        else:
            DrawFromFile(args.input, args.output_template, config)
    else:
        cif_or_pdb = handle_input_file(args.input)
        structure3d = read_3d_structure(cif_or_pdb, args.model)
        structure2d = read_2d_structure(cif_or_pdb, args.model)
        analysis = eltetrado(structure2d, structure3d, args.strict, args.no_reorder, args.stacking_mismatch)
        dto = generate_dto(analysis)
        if  not args.output_template:
            DrawFromString(orjson.dumps(dto), root, config)
        else:
            DrawFromString(orjson.dumps(dto), args.output_template, config)

