import requests
import json
import time
import logging
import os
import datetime
from http import HTTPStatus
import urllib3
import sys
import tarfile
import argparse
import gc
try:
    from colorama import Style, init as colorama_init
    colorama_init(autoreset=True)
    BOLD = Style.BRIGHT
    RESET = Style.RESET_ALL
except ImportError:
    # Fallback if colorama is not installed
    BOLD = ''
    RESET = ''


class ImportScriptHandler:
    def __init__(self, tenant_aapi_url: str, aapi_token: str, export_files_folder: str):
        self.INPUT_FOLDER = export_files_folder
        self.AAPI_TOKEN = aapi_token
        self.AAPI_URL = tenant_aapi_url
        urllib3.disable_warnings()
        self.DEFAULT_INJECT_ID = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        self.LOG_DIR = 'log'
        log_file_name = f"{self.LOG_DIR}{os.sep}archivedb_import_{self.DEFAULT_INJECT_ID}.log"
        if not os.path.exists(self.LOG_DIR):
            os.makedirs(self.LOG_DIR)
        logging.basicConfig(
            filename=log_file_name,
            format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
            level=logging.DEBUG)
        logging.log(level=logging.INFO, msg="Starting Import Procedure.")

        self.PROGRESS_FILE = 'progress-import.json'
        self.RULES_FILE = 'rules.json'

        if self.AAPI_URL is None or len(self.AAPI_URL) == 0:
            self.AAPI_URL = input("Type your SaaS Control-M Automation API endpoint: \n")
        if self.AAPI_TOKEN is None or len(self.AAPI_TOKEN) == 0:
            self.AAPI_TOKEN = input("Type your SaaS Control-M Automation API Token: \n")
        if self.INPUT_FOLDER is None or len(self.INPUT_FOLDER) == 0:
            self.INPUT_FOLDER = input("Type the full path to the local or network directory where the compressed Workload Archiving data files were exported: \n")

        # Remove trailing slash from AAPI URL to avoid double slashes
        self.AAPI_URL = self.AAPI_URL.rstrip('/')
        self.is_saas = '.com' in self.AAPI_URL
        if self.is_saas:
            self.AAPI_URL = self.AAPI_URL + "/automation-api"
        self.s3_pre_sign_url = self.AAPI_URL + "/config/data-platform"
        self.tenant_current_rules = self.AAPI_URL + "/config/data-platform/lifecycle/rules"
        if self.is_saas:
            self.headers = {
                'x-api-key': self.AAPI_TOKEN,
                'Content-Type': 'application/json'
            }
        else:
            self.headers = {
                'Authorization': f'Bearer {self.AAPI_TOKEN}',
                'Content-Type': 'application/json'
            }

    def run(self):
        print(f'Uploading TAR files from the following directory: {self.INPUT_FOLDER}. \nInject-Id is {self.DEFAULT_INJECT_ID}')
        logging.log(level=logging.INFO, msg=f'Inject-Id {self.DEFAULT_INJECT_ID}')
        progress = self.load_progress_from_file()
        files_to_upload = False
        with os.scandir(self.INPUT_FOLDER) as it:
            for entry in it:
                if entry.name.endswith(".tar") and entry.is_file():
                    try:
                        if entry.path in progress.keys() and progress[entry.path] == 'DONE':
                            print("Ignoring " + entry.path + " since it is marked as uploaded (Delete or modify the " + self.PROGRESS_FILE + " file if you want to re-upload)")
                            continue
                        files_to_upload = True
                        if not self.verify_rules(entry):
                            sys.exit(2)
                        http_response = requests.get(
                            url=self.s3_pre_sign_url + "/inject/signed-url?injectSource=archive&injectId=" + self.DEFAULT_INJECT_ID + "&fileName=" + entry.name,
                            headers=self.headers, verify=False)
                        http_response_text_as_json = json.loads(http_response.text)
                        logging.log(level=logging.INFO, msg='Rules http_response: ' + str(http_response_text_as_json))
                        logging.log(level=logging.INFO, msg=entry.name + '. injectId=' + http_response_text_as_json['injectId'])
                        
                        # Try upload with retry logic
                        upload_success = False
                        try:
                            with open(entry.path, 'rb') as file:
                                print("Uploading " + entry.path)
                                http_response = self.upload_with_retry(url=http_response_text_as_json['signedURL'], file=file, filepath=entry.path)
                                logging.log(level=logging.INFO, msg=str(http_response))
                                if http_response and http_response.status_code == HTTPStatus.OK:
                                    print("Successfully uploaded: " + entry.path)
                                    progress[entry.path] = 'DONE'
                                    upload_success = True
                                else:
                                    print("Error: " + str(http_response) + " when uploading " + entry.path)
                                    logging.log(level=logging.ERROR, msg="Error: " + str(http_response))
                        except Exception as upload_error:
                            logging.log(level=logging.ERROR, msg=f'Upload exception: {str(upload_error)}')
                            #print(f"Upload failed: {str(upload_error)}")
                        
                        # If upload failed, abort the entire process
                        if not upload_success:
                            progress[entry.path] = 'FAILED'
                            # Save progress before exiting
                            with open(self.PROGRESS_FILE, 'w') as f:
                                json.dump(progress, f)
                            print(f"ERROR: Upload failed for {entry.path}. Aborting import process.")
                            logging.log(level=logging.ERROR, msg=f"Upload failed for {entry.path}. Aborting import process.")
                            sys.exit(2)
                            
                    except Exception as e:
                        logging.log(level=logging.ERROR, msg='Exception: ' + str(e))
                        print("Error: " + str(e) + " when uploading " + entry.path)
                        progress[entry.path] = 'FAILED'
                        # Save progress before exiting
                        with open(self.PROGRESS_FILE, 'w') as f:
                            json.dump(progress, f)
                        sys.exit(2)
                else:
                    logging.log(level=logging.INFO, msg='Ignoring file: ' + entry.name)
                logging.log(level=logging.INFO, msg="progress: " + str(progress))
                with open(self.PROGRESS_FILE, 'w') as f:
                    json.dump(progress, f)
        if not files_to_upload:
            print(f"ERROR: No .tar files found in the directory: {self.INPUT_FOLDER}")
            print("Please verify the export_files_folder path contains valid .tar export files.")
            logging.log(level=logging.ERROR, msg=f"No .tar files found in {self.INPUT_FOLDER}")
            sys.exit(2)
        print(f'Finished uploading all compressed, self-hosted Workload Archiving export files from {self.INPUT_FOLDER}. \n'
              f'You can now run the verification script ({BOLD}run_archive_migration_report{RESET}) to verify that the exported data has been imported into Control-M SaaS.')
        sys.exit(0)

    def upload_with_curl_fallback(self, url, filepath):
        """Fallback upload using curl command when Python requests fails"""
        import subprocess
        try:
            logging.log(level=logging.INFO, msg="Attempting upload using curl as fallback")
            
            # Use curl with verbose SSL options for debugging
            cmd = [
                'curl',
                '-L',  # Follow redirects (important for S3 pre-signed URLs)
                '-X', 'PUT',
                '-H', 'Content-Type: application/octet-stream',
                '--data-binary', f'@{filepath}',
                '--insecure',  # Skip SSL verification
                '--ssl-no-revoke',  # Skip certificate revocation check
                '--max-time', '7200',  # 2 hour timeout
                '--write-out', '%{http_code}',  # Output HTTP status code
                '--silent',  # Silent mode
                '--show-error',  # Show errors
                '-o', 'nul',  # Discard response body (Windows nul)
                url
            ]
            
            logging.log(level=logging.INFO, msg=f"curl command: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=700)
            
            # Check HTTP status code from stdout
            http_code = result.stdout.strip() if result.stdout else ''
            logging.log(level=logging.INFO, msg=f"curl HTTP status: {http_code}, returncode: {result.returncode}")
            
            if result.returncode == 0 and http_code == '200':
                logging.log(level=logging.INFO, msg="curl upload succeeded with HTTP 200")
                return True
            else:
                logging.log(level=logging.ERROR, msg=f"curl upload failed - HTTP {http_code}, stderr: {result.stderr}")
                if result.stderr:
                    logging.log(level=logging.ERROR, msg=f"curl stderr: {result.stderr}")
                return False
                
        except FileNotFoundError:
            logging.log(level=logging.INFO, msg="curl not found in PATH")
            return False
        except subprocess.TimeoutExpired:
            logging.log(level=logging.ERROR, msg="curl upload timed out")
            return False
        except Exception as e:
            logging.log(level=logging.ERROR, msg=f"curl fallback error: {e}")
            return False
    
    def upload_with_retry(self, url, file, filepath, retries=5, backoff: int = 15):
        import ssl
        from requests.adapters import HTTPAdapter
        from urllib3.util.retry import Retry
        
        # Custom adapter that completely disables urllib3 retries
        class NoRetryAdapter(HTTPAdapter):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                
            def init_poolmanager(self, *args, **kwargs):
                # Create SSL context with broad compatibility
                ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
                ctx.check_hostname = False
                ctx.verify_mode = ssl.CERT_NONE
                ctx.set_ciphers('DEFAULT@SECLEVEL=1')  # Lower security level for compatibility
                # Apply SSL context
                kwargs['ssl_context'] = ctx
                # Disable all urllib3 retries - we handle retries ourselves
                kwargs['retries'] = Retry(0, read=False, connect=False, redirect=False)
                return super().init_poolmanager(*args, **kwargs)
        
        for attempt in range(1, retries + 1):
            session = None
            try:
                # Create fresh session for each attempt
                session = requests.Session()
                session.mount('https://', NoRetryAdapter())
                
                file.seek(0)  # Reset file pointer
                logging.log(level=logging.INFO, msg=f"Upload attempt {attempt}/{retries}")
                #print(f"Uploading... (attempt {attempt}/{retries})")
                
                resp = session.put(
                    url,
                    data=file,
                    headers={"Content-Type": "application/octet-stream"},
                    timeout=(10, 120),  # Reduced: (connect timeout 10s, read timeout 120s)
                    verify=False
                )
                resp.raise_for_status()
                logging.log(level=logging.INFO, msg=f"Upload succeeded on attempt {attempt}")
                #print(f"Upload completed successfully!")
                return resp
                
            except requests.exceptions.Timeout as e:
                error_msg = str(e)
                logging.log(level=logging.ERROR, msg=f"Attempt {attempt}/{retries} timeout: {error_msg}")
                
                if attempt < retries:
                    wait_time = backoff * attempt
                    logging.log(level=logging.INFO, msg=f"Connection timeout on attempt {attempt}/{retries}")
                    logging.log(level=logging.INFO, msg=f"Retrying in {wait_time} seconds...")
                    time.sleep(wait_time)
                else:
                    logging.log(level=logging.ERROR, msg=f"Upload timed out after {retries} attempts")
                    logging.log(level=logging.INFO, msg=f"Attempting curl fallback...")
                    
                    # Try curl as fallback
                    if self.upload_with_curl_fallback(url, filepath):
                        # Create a mock response object for compatibility
                        class MockResponse:
                            status_code = HTTPStatus.OK
                        return MockResponse()
                    
                    # Curl also failed
                    logging.log(level=logging.ERROR, msg="Timeout issues suggest network connectivity problems")
                    raise
                
            except (requests.exceptions.SSLError, ssl.SSLError, requests.exceptions.ConnectionError) as e:
                error_msg = str(e)
                logging.log(level=logging.ERROR, msg=f"Attempt {attempt}/{retries} SSL/connection error: {error_msg}")
                
                if attempt < retries:
                    wait_time = backoff * attempt  # Exponential backoff: 15s, 30s, 45s, 60s, 75s
                    logging.log(level=logging.ERROR, msg=f"SSL/TLS error on attempt {attempt}/{retries}")
                    logging.log(level=logging.INFO, msg=f"This may be due to network instability or firewall interference")
                    logging.log(level=logging.INFO, msg=f"Waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    logging.log(level=logging.ERROR, msg=f"Upload failed after {retries} attempts due to persistent SSL/TLS connection issues")
                    logging.log(level=logging.INFO, msg=f"Attempting curl fallback...")
                    
                    # Try curl as fallback
                    if self.upload_with_curl_fallback(url, filepath):
                        # Create a mock response object for compatibility
                        class MockResponse:
                            status_code = HTTPStatus.OK
                        return MockResponse()
                    
                    # Curl also failed
                    logging.log(level=logging.ERROR, msg="Possible causes: Network firewall blocking AWS S3 connections, Corporate proxy interfering with SSL connections, Network instability or routing issues, AWS S3 endpoint temporarily unavailable")
                    logging.log(level=logging.ERROR, msg="Please contact your network administrator or try again later")
                    raise
                    
            except Exception as e:
                logging.log(level=logging.ERROR, msg=f"Attempt {attempt}/{retries} unexpected error: {e}")
                
                if attempt < retries:
                    wait_time = backoff * attempt
                    logging.log(level=logging.ERROR, msg=f"Upload error on attempt {attempt}/{retries}: {type(e).__name__}")
                    logging.log(level=logging.INFO, msg=f"Retrying in {wait_time} seconds...")
                    time.sleep(wait_time)
                else:
                    logging.log(level=logging.ERROR, msg=f"Upload failed after {retries} attempts")
                    raise
                    
            finally:
                if session:
                    session.close()

    def verify_rules(self, tar):
        logging.log(level=logging.INFO,
                    msg="calling : " + str(self.tenant_current_rules) + " with headers " + str(self.headers))
        http_response = requests.get(url=self.tenant_current_rules, headers=self.headers, verify=False)
        logging.log(level=logging.INFO, msg=self.tenant_current_rules + ' response=' + str(http_response))
        if not '200' in str(http_response):
            logging.log(level=logging.ERROR,
                        msg=f"'{str(http_response)}' str(http_response) failure.")
            error_detail = ""
            try:
                error_detail = " - " + http_response.text
            except:
                pass
            print("Failed to retrieve the Control-M SaaS Workload Archiving policies: " + str(http_response) + error_detail)
            return False
        if not 'rules' in str(http_response.text):
            logging.log(level=logging.ERROR,
                        msg=f"'{http_response.text}' http_response.text failure - verify aapi URL.")
            error_detail = ""
            try:
                error_detail = " - " + http_response.text
            except:
                pass
            print("Failed to retrieve the Control-M SaaS Workload Archiving policies: " + str(http_response) + error_detail)
            return False
        http_response_text_as_json = json.loads(http_response.text)
        logging.log(level=logging.INFO, msg="accepted tenant rules response : " + str(http_response_text_as_json))
        tenant_current_rules = [rule['name'] for rule in http_response_text_as_json['rules']]
        logging.log(level=logging.INFO, msg=tar.name + '. tenant current rules=' + str(tenant_current_rules))
        if len(tenant_current_rules) == 0:
            logging.log(level=logging.ERROR,
                        msg=f"'{str(tenant_current_rules)}' no archive rules defined in the data platform.")
            print("No Control-M SaaS Workload Archiving policies are defined.")
            #return False
        for item in tenant_current_rules:
            logging.log(level=logging.INFO, msg=item)
        if not tenant_current_rules:
            logging.log(level=logging.ERROR,
                        msg=f"'{tenant_current_rules}' tenant_current_rules failure.")
            print("Failed to retrieve the Control-M SaaS Workload Archiving policies")
            return False
        tar_rules = self.get_tar_rules(tar.path, self.RULES_FILE)
        if not tar_rules:
            logging.log(level=logging.ERROR, msg=f"'{tar.path}' is not valid. The rules file is missing or empty.")
            print(f"'{tar.path}' is not valid. The rules file is missing.")
            return False
        logging.log(level=logging.INFO, msg=tar.name + '. tar rules=' + str(tar_rules))
        for tar_rule in tar_rules:
            if tar_rule not in tenant_current_rules:
                logging.log(level=logging.INFO, msg=tar_rule + ' is missing in SaaS')
                print(
                    tar.name + ': The following Control-M Self-Hosted and SaaS Workload Archiving policies do not match. You must define the \'' + tar_rule + '\' Workload Archiving Policy in Control-M SaaS exactly as it is defined in Control-M Self-Hosted.')
                return False
        return True

    def get_tar_rules(self, tar_path, file_name):
        # Open the tar file in read mode
        with tarfile.open(tar_path, 'r') as tar:
            # Iterate through the files in the tar
            for member in tar.getnames():
                # Check if the file name exists in the member's name
                if file_name in member:
                    # Extract the file-like object without extracting the file to disk
                    file_obj = tar.extractfile(member)

                    # Read and print the file's content
                    if file_obj:
                        content = file_obj.read().decode('utf-8')  # Decode assuming the file is text (UTF-8)
                        logging.log(level=logging.INFO, msg=tar_path + '. rules json content=' + str(content))
                        json_content = content.strip('"').replace('\\"', "\"")
                        logging.log(level=logging.INFO, msg=tar_path + '. rules json converted=' + str(json_content))
                        values = json.loads(str(json_content))
                        logging.log(level=logging.INFO, msg=tar_path + '. rules json array=' + str(values))
                        return values
                    break  # Exit after printing the first match (remove this if you want all matches)

    def load_progress_from_file(self):
        progress = {}
        try:
            with open(self.PROGRESS_FILE) as f:
                progress = json.load(f)
                logging.log(level=logging.INFO, msg=self.PROGRESS_FILE + " content loaded: " + str(progress))
        except Exception as e:
            logging.log(level=logging.INFO,
                        msg=self.PROGRESS_FILE + " not exists in workspace. running import from scrach")
        return progress


if __name__ == '__main__':
    try:
        # Required command line arguments for the import script:
        parser = argparse.ArgumentParser(description="Control-M Workload Archiving data migration"
                                                     " into Control-M SaaS export script's args parser")
        parser.add_argument('--tenant_aapi_url', required=True, type=str, help="Control-M SaaS Tenant's AAPI url")
        parser.add_argument('--aapi_token', required=True, type=str, help="Control-M SaaS Tenant AAPI token")
        parser.add_argument('--export_files_folder', required=True, type=str,
                            help="The file system location of the export files to be import")
        parser.add_argument('--python_exe', required=False, default='python3', type=str,
                            help="The python interpreter name")

        # Parse the command-line arguments:
        args = parser.parse_args()
        tenant_aapi_url = args.tenant_aapi_url
        aapi_token = args.aapi_token
        export_files_folder = args.export_files_folder

        import_script_handler = ImportScriptHandler(tenant_aapi_url, aapi_token, export_files_folder)
        import_script_handler.run()
        gc.collect()
        exit(0)
    except Exception as e:
        err_msg = "An import exception has occurred, \n" \
                  f"the exception message is: {str(e)}"
        logging.log(level=logging.INFO, msg=err_msg)
        print(err_msg)
        gc.collect()
        exit(2)

