import argparse
import subprocess
import re
import sys
import time
import json
import ssl
import hashlib
import urllib.request
import time
from datetime import datetime, timedelta
import threading
import asyncio

# Usage Example:
# python3 alamut-batch-wrapper.py --in /foo/bar/my_input.vcf --ann /tmp/my_output.vcf --unann /tmp/my_unann.txt --assbly GFCh37 -v --java_opts="-Dmicronaut.config.files=./application.yaml"


###############################################################################################################
# This section of code, same that appears in sg-upload-v2-wrapper.py will ensure Sophia CLI Jar is up to date #
###############################################################################################################
REMOTE_URL = "https://ddm.sophiagenetics.com/direct/sg/uploaderv2"
UPLOADER_FILENAME = "sg-upload-v2-latest.jar"
UPLOAD_CHECKSUM_FILENAME = "sg-upload-v2-latest.jar.md5"
VERSION = "1.0.3"

# In case of "[SSL: CERTIFICATE_VERIFY_FAILED] certificate verify error please uncomment the following line
# see https://support.sectigo.com/articles/Knowledge/Sectigo-AddTrust-External-CA-Root-Expiring-May-30-2020
ssl._create_default_https_context = ssl._create_unverified_context

def get_remote_checksum():
    """Fetch the checksum of the remote file."""
    url = f"{REMOTE_URL}/{UPLOAD_CHECKSUM_FILENAME}"
    try:
        response = urllib.request.urlopen(url)
        if response.status != 200:
            print(f"Error: Could not find the remote version at {url}")
            sys.exit(1)
        md5sum = response.read().decode('utf-8')
        return md5sum
    except Exception as err:
        print(f"Error: There was a problem connecting to {url}")
        print("If you encounter an [SSL: CERTIFICATE_VERIFY_FAILED] error, please uncomment the following line:")
        print("# ssl._create_default_https_context = ssl._create_unverified_context")
        print(
            "For more information, see https://support.sectigo.com/articles/Knowledge/Sectigo-AddTrust-External-CA-Root-Expiring-May-30-2020")
        print(f"Technical details: {err}")
        return ""


def get_current_checksum():
    """Calculate the checksum of the current file."""
    hash_md5 = hashlib.md5()
    try:
        with open(UPLOADER_FILENAME, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()
    except FileNotFoundError:
        return ""  # return empty checksum if no file found


def download_latest_uploader():
    """Download the latest version of the uploader."""
    update_url = f"{REMOTE_URL}/{UPLOADER_FILENAME}"
    response = urllib.request.urlopen(update_url)
    with open(UPLOADER_FILENAME, 'wb') as f:
        f.write(response.read())

def check_and_update_uploader_if_needed():
    remote_checksum = get_remote_checksum()

    if remote_checksum == "":
        print("WARN. No new version found. Using previous one!")
    else:
        current_checksum = get_current_checksum()
        if current_checksum == "":
            print(f"Downloading latest uploader version. Checksum: {remote_checksum}.")
            download_latest_uploader()
        else:
            remote_checksum = remote_checksum.strip()
            current_checksum = current_checksum.strip()
            if remote_checksum != current_checksum:
                print(f"Current checksum: {current_checksum}")
                print(f"Remote checksum: {remote_checksum}")
                download_latest_uploader()
                print(f"Updated to version {remote_checksum}.")
            else:
                print(f"Script is up-to-date (checksum {remote_checksum})")


last_value = "Submitted"
timeout_duration = timedelta(hours=24)
start_time = datetime.now()
stop_spinner = False
run_id = None

# Will run in separate thread to provide status updates to user
def spinner():
    spinner_chars = ['|', '/', '-', '\\']
    spinner_index = 0
    while stop_spinner == False:
        print(f"\rRun ID: {run_id}, Current Status: {last_value} {spinner_chars[spinner_index]}", end="")
        spinner_index = (spinner_index + 1) % len(spinner_chars)
        time.sleep(0.5)

async def create_run_and_parse_run_id(command):
    polling_task = None
    global run_id

    process = await asyncio.create_subprocess_exec(
        *command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT
    )
    try:
        # Read stdout line by line
        async for line in process.stdout:
            decoded_line = line.decode('utf-8').strip()

            # Check for runId
            success_match = re.search(r"Run successfully created with id (\d+)", decoded_line)
            fail_match = re.search(r"An error occurred: (.*)", decoded_line)
            if success_match and run_id is None:

                run_id = success_match.group(1)

                # Start polling in a separate task
                polling_task = asyncio.create_task(poll_status(run_id))
            elif fail_match:
                print(f"Failed to create run. Error: {fail_match.group(1)}")
                sys.exit(1)

        # Wait for the process to finish
        return_code = await process.wait()

        if return_code != 0:
            if polling_task:
                polling_task.cancel()  # Cancel the polling task if the process failed
            print("\nRun in Error State. Contact SOPHiA Genetics Support For More Information.", file=sys.stderr)
            sys.exit(return_code)

        # Ensure polling completes if it was started
        if polling_task:
            await polling_task

    finally:
        if process.stdout:
            process.stdout.feed_eof()
        if process.stderr:
            process.stderr.feed_eof()


async def poll_status(run_id):
    """
    Poll for the status of the run using the runId.
    """

    # Command to get status of run
    poll_command_template = [
        "java", "-jar", *gen2_eap_java_opts,
        "sg-upload-v2-latest.jar", "status", "-i", run_id
    ]

    spinner_thread = threading.Thread(target=spinner)
    spinner_thread.start()
    global last_value
    global stop_spinner

    max_retries = 5
    retry_count = 0

    try:
        while True:
            elapsed_time = datetime.now() - start_time
            if elapsed_time > timeout_duration:
                stop_spinner = True
                spinner_thread.join()
                print("\nProcess timed out after 24 hours.", file=sys.stderr)
                sys.exit(1)

            try:
                poll_result = subprocess.run(poll_command_template, capture_output=True, text=True, check=True)
                current_value = poll_result.stdout.strip()
                retry_count = 0  # Reset retry count on success

            except subprocess.CalledProcessError as pollingError:
                retry_count += 1
                if retry_count >= max_retries:
                    stop_spinner = True
                    spinner_thread.join()
                    print(f"\nMaximum polling retries reached: {pollingError}. Exiting.", file=sys.stderr)
                    sys.exit(1)
                await asyncio.sleep(30)
                continue

            if current_value != last_value:
                last_value = current_value

            if current_value == "Error":
                stop_spinner = True
                spinner_thread.join()
                print("\nRun in Error State. Contact SOPHiA Genetics Support For More Information.", file=sys.stderr)
                sys.exit(1)

            if current_value == "Finished":
                stop_spinner = True
                spinner_thread.join()
                print(f"\nRun {run_id} completed successfully.")
                break
            # Poll for status every 30 seconds
            await asyncio.sleep(30)

    finally:
        stop_spinner = True
        spinner_thread.join()


check_and_update_uploader_if_needed()

##################################################################################################################
# If we reached here the Sophia CLI Jar is up to date and we can continue with its usage for AlamutBatch Wrapper #
##################################################################################################################


verbose = False
parser = argparse.ArgumentParser(description="Alamut Batch Wrapper to Sophia CLI")
# In is a reserverd word so we ahve to do this remap
parser.add_argument("--in", dest="input_file", required=True, help="Path to the input VCF file")
parser.add_argument("--ann", required=True, help="Path to file to save the output file of annotated variants")
#TODO: implement once pipeline has this file available in output
parser.add_argument("--unann", required=True, help="File path to save unannotated variants (Currently not supported, will produce empty file)")
parser.add_argument("--assbly", required=False, default="GRCh37", help="Genome Assebly, options are GRCh37 or GRCh38. Default is GRCh37")
parser.add_argument("--alltrans", action="store_true", help="No-Op. This is set by default but is included for compatibility with Alamut Batch CLI")
parser.add_argument("--donnsplice", action="store_true", help="No-Op. This is set by default but is included for compatibility with Alamut Batch CLI")
parser.add_argument("--dogenesplicer", action="store_true", help="No-Op. This is set by default but is included for compatibility with Alamut Batch CLI")
parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose mode")
# TODO: Uncomment when we move out of EAP, re-generate commands inclusion in cli docs.
#parser.add_argument("--java_opts", default="", help="Optional Java options to pass to CLI")
args = parser.parse_args()
# TODO: Uncomment when we move out of EAP
#java_opts = args.java_opts


gen2_eap_java_opts = [
    "-Dsg2-api.upload-v1=https://eap.sophiagenetics.com/gen2-eap/api/upload/v1",
    "-Dsg2-api.bar-v1=https://gen2-eap.eap.sophiagenetics.com/api/bar/v1",
    "-Dsg2-api.base-url=https://eap.sophiagenetics.com/gen2-eap",
    "-Dsg2-api.file-svc-host=https://haproxy-staging.sophiagenetics.com/gen2-eap/api/file",
    "-Dsg2-ddm.host=https://gen2-eap.eap.sophiagenetics.com/",
    "-Diam.host=https://gen2-eap.eap.sophiagenetics.com/iam",
    "-Dmicronaut.http.services.iam.url=https://gen2-eap.eap.sophiagenetics.com/iam"
]


login_command = [
    "java", "-jar", *gen2_eap_java_opts,
    "sg-upload-v2-latest.jar", "sso"
]
if args.verbose:
    print("Running command:", " ".join(login_command))
subprocess.run(login_command)


if args.assbly.lower() == "grch37":
    pip_id = "7211"
    ref_genome = "hg19"
elif args.assbly.lower() == "grch38":
    pip_id = "7212"
    ref_genome = "hg38"
else:
    print("Error: Invalid genome assembly. Accepted values are GFCh37 or GFCh38 (case insensitive).")
    exit(1)

##################################
##          Create Run          ##
##################################

print(f"Starting run creation with input file {args.input_file} and assbly option {args.assbly}")
command = [
    "java", "-jar", *gen2_eap_java_opts,
    "sg-upload-v2-latest.jar", "new", "-p", pip_id,
    "-f", args.input_file,
    "--ref", "alamut_batch", "--sampletype", "126000","--pipelineParameters", f"reference.genome.ucsc_id={ref_genome}", "-fp", "--upload"
]
if args.verbose:
    print("Running command:", " ".join(command))

asyncio.run(create_run_and_parse_run_id(command))


##################################
##      Get Run Output Files    ##
##################################

# If we have reached this point, then the run is done and we can start getting output file information
file_list_command = [
    "java", "-jar", *gen2_eap_java_opts,
    "sg-upload-v2-latest.jar", "file", "--list", "--run-id", run_id
]

file_list_result = subprocess.run(file_list_command, capture_output=True, text=True)

file_to_download="full_variant_table_for_batch.txt"
try:
    file_list = json.loads(file_list_result.stdout)
    file_id = next((item["id"] for item in file_list if item["name"] == file_to_download), None)
    
    if file_id is None:
        print(f"Error: {file_to_download} not found in the file list.", file=sys.stderr)
        sys.exit(1)

except json.JSONDecodeError:
    print("Error: Failed to parse JSON output.", file=sys.stderr)
    sys.exit(1)


##################################
##    Download Files We Need    ##
##################################

download_command = [
    "java", "-jar", *gen2_eap_java_opts,
    "sg-upload-v2-latest.jar", "file", "--download",
    "--file-id", str(file_id),
    "--file-out", args.ann
]
download_result = subprocess.run(download_command, capture_output=True, text=True)

if download_result.returncode == 0:
    print(f"File downloaded successfully to {args.ann}")
else:
    print("Error during file download:", download_result.stderr, file=sys.stderr)
    sys.exit(1)

# TODO: When pipeline produces unannotated variants file, use same pattern as above to save file
try:
    with open(args.unann, 'w') as unann_file:
        pass  # Creating an empty file
except IOError as e:
    print(f"Error creating file at {args.unann}: {e}", file=sys.stderr)
    sys.exit(1)