#! /usr/bin/python

import urllib.request
import urllib.parse
import re
import os
import subprocess
import ssl
import sys
import argparse
import tempfile

try:
  from tqdm import tqdm
  TQDM_AVAILABLE = True
except ImportError:
  TQDM_AVAILABLE = False

try:
  import certifi
  CERTIFI_AVAILABLE = True
except ImportError:
  CERTIFI_AVAILABLE = False

class DownloadProgressBar:
  def __init__(self, total=0, unit='iB', unit_scale=True, unit_divisor=1024):
    self.pbar = None
    self.total = total
    self.unit = unit
    self.unit_scale = unit_scale
    self.unit_divisor = unit_divisor

  def __call__(self, downloaded, block_size, total_size):
    if not self.pbar:
      self.total = total_size
      self.pbar = tqdm(total=total_size, unit=self.unit, unit_scale=self.unit_scale, unit_divisor=self.unit_divisor)

    self.pbar.update(downloaded - self.pbar.n)

    if downloaded >= total_size:
      self.pbar.close()

def simple_progress_bar(downloaded, block_size, total_size):
  progress = min(round(downloaded / total_size * 100, 2), 100)
  sys.stdout.write(f"\rDownload progress: {progress:6.2f}%")
  sys.stdout.flush()

def create_ssl_context():
  """
  Build an SSL context that works on macOS, Linux, and Windows.

  macOS's Python does not use the system keychain by default.  We try
  certifi first (most reliable), then the system store, then a glob search
  for certifi's cacert.pem left by the official Python installer.
  """
  if CERTIFI_AVAILABLE:
    return ssl.create_default_context(cafile=certifi.where())

  import glob
  ctx = ssl.create_default_context()

  if sys.platform == "darwin" and not ctx.get_ca_certs():
    for pattern in [
      "/Library/Frameworks/Python.framework/Versions/*/lib/python*/site-packages/certifi/cacert.pem",
      "/usr/local/lib/python*/site-packages/certifi/cacert.pem",
      "/opt/homebrew/lib/python*/site-packages/certifi/cacert.pem",
    ]:
      matches = glob.glob(pattern)
      if matches:
        try:
          return ssl.create_default_context(cafile=matches[0])
        except Exception:
          pass

  return ctx

def download_large_file(file_id, output_filename, ssl_context):
  url = f"https://drive.google.com/uc?export=download&id={file_id}"
  req = urllib.request.Request(url, method="GET")

  try:
      with urllib.request.urlopen(req, context=ssl_context) as response:
          response_content = response.read().decode('utf-8')
  except urllib.error.URLError as e:
      print(f"Error accessing Google Drive: {e}")
      return False

  form_action = re.search(r'<form id="download-form" action="([^"]*)"', response_content)
  confirm = re.search(r'name="confirm" value="([^"]*)"', response_content)
  uuid_value = re.search(r'name="uuid" value="([^"]*)"', response_content)

  if not all([form_action, confirm, uuid_value]):
      print("Failed to extract form details. The download process might have changed.")
      return False

  download_url = (f"{form_action.group(1)}?"
                  f"id={file_id}&"
                  f"export=download&"
                  f"confirm={confirm.group(1)}&"
                  f"uuid={uuid_value.group(1)}")

  print(f"Downloading file to {output_filename}...")
  try:
    with urllib.request.urlopen(download_url, context=ssl_context) as response, open(output_filename, 'wb') as out_file:
      file_size = int(response.headers.get('Content-Length', 0))
      print(file_size)
      block_size = 8192

      if TQDM_AVAILABLE:
        progress_bar = tqdm(total=file_size, unit='iB', unit_scale=True)
      else:
        progress_bar = simple_progress_bar

      size = 0
      while True:
        buffer = response.read(block_size)
        if not buffer:
          break
        size += len(buffer)
        out_file.write(buffer)
        if TQDM_AVAILABLE:
          progress_bar.update(len(buffer))
        else:
          progress_bar(size, block_size, file_size)

      if TQDM_AVAILABLE:
        progress_bar.close()
      else:
        print()  # New line after progress bar

    print("Download completed.")
    return True
  except urllib.error.URLError as e:
    print(f"\nError downloading file: {e}")
    return False

def extract_file(filename, extract_path):
  print(f"Extracting \033[94m{filename}\033[0m to \033[96m{extract_path}\033[0m...")
  try:
    subprocess.run(["7z", "x", filename, f"-o{extract_path}"], check=True)
    print("Extraction completed.")
    return True
  except subprocess.CalledProcessError:
    print("Extraction failed. Make sure 7z is installed and in your PATH.")
    return False

def check_writable(path):
  try:
    testfile = tempfile.TemporaryFile(dir=path)
    testfile.close()
  except OSError:
    return False
  return True

def print_usage():
  print("Usage: python3 download_script.py [--datadir PATH]")
  print("Options:")
  print("  --datadir PATH, -d PATH    Specify the path to the Stealth data directory")
  print("                             (default: ~/.StealthCoin)")

def main():
  parser = argparse.ArgumentParser(description="Download and extract Stealth bootstrap file.")
  parser.add_argument('--datadir', '-d', type=str, help='Path to the Stealth data directory')
  args = parser.parse_args()

  if not check_writable('.'):
    print("Error: Current working directory is not writable.")
    sys.exit(1)

  if args.datadir:
    data_dir = os.path.abspath(args.datadir)
  else:
    data_dir = os.path.expanduser("~/.StealthCoin")

  if not os.path.exists(data_dir):
    print(f"Error: Cannot find the Stealth data directory at {data_dir}")
    print_usage()
    sys.exit(1)

  file_id = "1WNmH1DiO4-_D-UYupoyFXM-rKMlGVTet"
  output_filename = "xst-bootstrap.7z"

  print("This script will download a large file (>3.0GB) from Google Drive.")
  print("Google Drive can't scan this file for viruses due to its size.")

  ssl_context = create_ssl_context()
  if CERTIFI_AVAILABLE:
    print("SSL: using certifi certificate bundle.")
  else:
    print("SSL: using system certificate store.")
    if sys.platform == "darwin":
      print("     (If SSL errors occur, run: \033[1mpip3 install certifi\033[0m)")

  user_input = input("\033[1mDo you want to proceed with the download? (y/n):\033[0m ").lower()

  if user_input == 'y':
    if download_large_file(file_id, output_filename, ssl_context):
      user_input = input("\033[1mDo you want to extract the downloaded file? (y/n):\033[0m ").lower()
      if user_input == 'y':
        extract_file(output_filename, data_dir)
  else:
    print("\033[31mDownload cancelled.\033[0m")

if __name__ == "__main__":
  main()
