import os
import zipfile
import logging
import paramiko
import tempfile
import functools
import pandas as pd

from pathlib import Path
from Crypto.Cipher import AES

from typing import Callable

import logging
logger = logging.getLogger(__name__)

def decrypt_file(file_path: Path, key: bytes, iv: bytes, prefix: str="decrypted_") -> Path:
    if file_path.stem.startswith(prefix):
        return file_path
    
    # Initialize the AES cipher with CBC mode
    cipher = AES.new(key, AES.MODE_CBC, iv)
    output_file = file_path.with_name(prefix + file_path.stem + ".zip")
    # Decrypt the input file and write the decrypted content to the output file
    with file_path.open("rb") as f_in, output_file.open("wb") as f_out:
        decrypted_data = cipher.decrypt(f_in.read())
        f_out.write(decrypted_data)
    return output_file

def encrypt_file(file_path: Path, key: bytes, iv: bytes, prefix: str="encrypted_") -> Path:
    if prefix in file_path.stem:
        return file_path
    # Initialize the AES cipher with CBC mode
    cipher = AES.new(key, AES.MODE_CBC, iv)
    output_file = file_path.with_name(prefix + file_path.stem + ".enc")
    # Encrypt the input file and write the encrypted content to the output file
    with file_path.open("rb") as f_in, output_file.open("wb") as f_out:
        while True:
            data = f_in.read(16)
            if len(data) == 0:
                break
            elif len(data) % 16!= 0:
                data += b' ' * (16 - len(data) % 16)
            encrypted_data = cipher.encrypt(data)
            f_out.write(encrypted_data)
    return output_file

def download_decrypt_extract(sftp: paramiko.SFTPClient, remote_file: str, output_path: Path, key: bytes, iv: bytes) -> bool:
    """
    Downloads a file from SFTP, decrypts it using decrypt_file, extracts its contents, and cleans up temporary files.

    Args:
    sftp (paramiko.SFTPClient): An active SFTP client connection.
    remote_file (str): Path to the file on the remote server.
    output_path (Path): Local path where extracted contents should be saved.
    key (bytes): 16-byte key for AES decryption.
    iv (bytes): 16-byte initialization vector for AES decryption.

    Returns:
    bool: True if successful, False otherwise.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        try:
            # Download file
            local_encrypted_path = Path(temp_dir) / Path(remote_file).name
            sftp.get(remote_file, str(local_encrypted_path))

            # Decrypt file using decrypt_file function
            decrypted_path = decrypt_file(local_encrypted_path, key, iv)

            # Extract contents
            with zipfile.ZipFile(decrypted_path, 'r') as zip_ref:
                zip_ref.extractall(output_path)

            logger.debug(f"Successfully processed {remote_file}")
            return True

        except paramiko.SSHException as e:
            logger.error(f"SFTP error while downloading {remote_file}: {str(e)}")
        except ValueError as e:
            logger.error(f"Decryption error for {remote_file}: {str(e)}")
        except zipfile.BadZipFile as e:
            logger.error(f"ZIP extraction error for {remote_file}: {str(e)}")
        except Exception as e:
            logger.error(f"Unexpected error processing {remote_file}: {str(e)}")

        return False

def check_required(required_keys: list[str]):
    def decorator(func: Callable):
        @functools.wraps(func)
        def wrapper(config: dict[str, str], *args, **kwargs):
            missing_keys = [key for key in required_keys if key not in config]
            if missing_keys:
                raise ValueError(f"Missing required keys in config: {', '.join(missing_keys)}")
            return func(config, *args, **kwargs)
        return wrapper
    return decorator

@check_required(['FTP_ADDRESS', 'FTP_USER', 'FTP_PASSWORD', 'AES_KEY', 'AES_IV'])
def download_decrypt_extract_new_files(
    config: dict[str, str], 
    tasks: list[str], 
    local: Path,
    force: bool = False,
    callback: Callable[[str, int, int, str], None] | None = None
) -> list[tuple[str, str]]:
    """
    Downloads, decrypts, and extracts new files from the SFTP server, skipping files that have already been processed.
    Uses a callback function for progress tracking.

    Parameters:
    config (dict[str, str]): Configuration dictionary containing SFTP details, key, and IV.
    tasks (list[str]): List of directory types to process (e.g., ['R15', 'C15']).
    local (Path): The local root path to save extracted files.
    callback (callable, optional): A function to call for progress updates. It should accept the following parameters:
                                   - task_type (str): The current task type being processed.
                                   - total_files (int): Total number of files to process.
                                   - current_file (int): Current file being processed (1-indexed).
                                   - file_name (str): Name of the current file being processed.

    Returns:
    list[tuple[str, str]]: A list of tuples containing (zip_name, task_type) of newly processed files.
    """

    key = bytes.fromhex(config['AES_KEY'])
    iv = bytes.fromhex(config['AES_IV'])

    csv_path = local / "processed_zips.csv"
    if force and csv_path.exists():
        csv_path.unlink()
    
    if csv_path.exists():
        df = pd.read_csv(csv_path)
        processed_zips = set(df['zip_name'])
    else:
        df = pd.DataFrame(columns=['zip_name', 'flux'])
        processed_zips = set()

    transport = paramiko.Transport((config['FTP_ADDRESS'], 22))
    transport.connect(username=config['FTP_USER'], password=config['FTP_PASSWORD'])
    sftp = paramiko.SFTPClient.from_transport(transport)

    newly_processed_files = []

    try:
        for task_type in tasks:
            distant = '/flux_enedis/' + str(config.get(f'FTP_{task_type}_DIR', task_type))
            local_dir = local.joinpath(task_type)
            local_dir.mkdir(parents=True, exist_ok=True)

            try:
                files_to_process = [f for f in sftp.listdir(distant) if f not in processed_zips]
                total_files = len(files_to_process)

                for index, file_name in enumerate(files_to_process, start=1):
                    if callback:
                        callback(task_type, total_files, index, file_name)

                    remote_file_path = os.path.join(distant, file_name)
                    output_path = local_dir / file_name.replace('.zip', '')
                    
                    success = download_decrypt_extract(sftp, remote_file_path, output_path, key, iv)
                    
                    if success:
                        newly_processed_files.append((file_name, task_type))
                        df = pd.concat([df, pd.DataFrame({'zip_name': [file_name], 'flux': [task_type]})], ignore_index=True)

            except Exception as e:
                logger.error(f"Failed to process files from {distant}: {e}")

    finally:
        sftp.close()
        transport.close()

    df.to_csv(csv_path, index=False)

    return newly_processed_files

