import subprocess
import os
import re
from .bridrive import upload_to_bridrive

def run_cmd(cmd):
    proc = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return proc.returncode == 0, proc.stdout, proc.stderr

def list_part_files(hdfs_dir):
    success, output, error = run_cmd(f'hadoop fs -ls "{hdfs_dir}"')
    if not success:
        raise RuntimeError(f"Error listing HDFS directory:\n{error}")
    lines = output.strip().split('\n')
    part_files = []
    for line in lines:
        if 'part-' in line:
            match = re.search(r'(/.*part-[^\s]*)', line)
            if match:
                part_files.append(match.group(1))
    return part_files

def download_and_rename_files(part_files, base_file_name):
    local_dir = os.environ['PATH_DATA']
    local_files = []

    if len(part_files) == 1:
        hdfs_file = part_files[0]
        local_file = os.path.join(local_dir, f"{base_file_name}.csv")

        if os.path.exists(local_file):
            print(f"Deleting existing file {local_file}")
            os.remove(local_file)
            # Confirm deletion
            if os.path.exists(local_file):
                raise RuntimeError(f"Failed to delete existing file {local_file}")

        success, _, err = run_cmd(f'hadoop fs -get "{hdfs_file}" "{local_file}"')
        if not success:
            raise RuntimeError(f"Failed to copy {hdfs_file}:\n{err}")
        local_files.append(local_file)

    else:
        for i, hdfs_file in enumerate(part_files):
            local_file = os.path.join(local_dir, f"{base_file_name}_part_{i}.csv")

            if os.path.exists(local_file):
                print(f"Deleting existing file {local_file}")
                os.remove(local_file)
                if os.path.exists(local_file):
                    raise RuntimeError(f"Failed to delete existing file {local_file}")

            success, _, err = run_cmd(f'hadoop fs -get "{hdfs_file}" "{local_file}"')
            if not success:
                raise RuntimeError(f"Failed to copy {hdfs_file}:\n{err}")
            local_files.append(local_file)
    return local_files

def delete_hdfs_dir(hdfs_dir):
    success, _, err = run_cmd(f'hadoop fs -rm -r -skipTrash "{hdfs_dir}"')
    if not success:
        raise RuntimeError(f"Failed to delete HDFS directory:\n{err}")

def download_hdfs_parts_and_cleanup(file_name):
    hdfs_dir = f"/user/{os.getenv('HADOOP_USER_NAME')}/{file_name}"
    part_files = list_part_files(hdfs_dir)
    if not part_files:
        print("No part files found in HDFS directory.")
        return []
    try:
        files = download_and_rename_files(part_files, file_name)
        delete_hdfs_dir(hdfs_dir)
        return files
    except RuntimeError as e:
        print(str(e))
        print("Operation aborted. HDFS directory NOT deleted.")
        return []
        
def dataframe_to_bridrive(df, link_bridrive, file_name, repartition_number = 1, delete=True):
  """
  Write a Spark DataFrame to HDFS as CSV files, download the parts locally and upload them to BriDrive.

  Parameters:
  -----------
  df : pyspark.sql.DataFrame
      The Spark DataFrame to be saved and uploaded.
  link_bridrive : str
      The destination upload link for bridrive.
  base_file_name : str
      The base name used for naming downloaded CSV files.
  repartition_number : int, optional (default=1)
      The number of partitions to repartition the DataFrame before writing CSV files.

  Returns:
  --------
  None

  Example:
  --------
  >>> dataframe_to_bridrive(df, "http://xxxx", "tes")
  >>> dataframe_to_bridrive(df, "http://xxxx", "brimo_daily", repartition_number=2)
  """
  if " " in file_name:
    raise ValueError(f"Invalid file_name: contains space. '{file_name}'")
  print("Generate file in HDFS.....")
  df\
    .repartition(repartition_number)\
    .write\
    .option("delimiter", ",")\
    .csv(file_name, header=True, mode='overwrite')
  print("csv file success generated")
  files = download_hdfs_parts_and_cleanup(file_name)
  if files:
      upload_to_bridrive(link_bridrive, files, delete)
  else:
      print("No files to upload to bridrive.")
    
