from datetime import datetime, timedelta
from airflow.decorators import dag, task, task_group, sensor
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.models import XCom
from airflow.utils.session import create_session
from airflow.exceptions import AirflowSensorTimeout
from typing import Dict, List
import yaml
import time

ssh_hook = SSHHook(ssh_conn_id='linux_server')

def load_config(config_path: str) -> dict:
   with open(config_path, 'r') as f:
       return yaml.safe_load(f)

def _send_missing_files_alert(file_type: str, missing_files: list, **context):
   task_instance = context['task_instance']
   last_alert_time = task_instance.xcom_pull(key=f'last_alert_{file_type}', default=0)
   current_time = time.time()
   
   if current_time - last_alert_time >= 60:
       # Implement alert logic (email/slack)
       print(f"Alert for {file_type}: {missing_files}")  # Replace with actual alert
       task_instance.xcom_push(key=f'last_alert_{file_type}', value=current_time)

@task_group
def process_files(config: dict, file_type: str):
   @sensor(poke_interval=180, timeout=3600, soft_fail=True)
   def file_sensor(**context):
       with ssh_hook.get_conn() as ssh_client:
           missing_zip_files = []
           missing_contents = {}
           
           for item in config:
               cmd = f"test -f /path/to/{file_type}/files/{item['zip_file']} && echo 1 || echo 0"
               stdin, stdout, stderr = ssh_client.exec_command(cmd)
               if stdout.read().decode().strip() == "0":
                   missing_zip_files.append(item['zip_file'])
                   continue
                   
               cmd = f"unzip -Z1 /path/to/{file_type}/files/{item['zip_file']}"
               stdin, stdout, stderr = ssh_client.exec_command(cmd)
               zip_contents = set(stdout.read().decode().splitlines())
               missing = set(item['expected_files']) - zip_contents
               if missing:
                   missing_contents[item['zip_file']] = list(missing)

           if missing_zip_files:
               _send_missing_files_alert(f"{file_type}_missing_zips", missing_zip_files, **context)
               return False
               
           if missing_contents:
               missing_files_info = []
               for zip_file, missing_files in missing_contents.items():
                   missing_files_info.append(f"Zip: {zip_file}, Missing files: {', '.join(missing_files)}")
               _send_missing_files_alert(f"{file_type}_missing_contents", missing_files_info, **context)
               cleanup_xcom()
               raise AirflowSensorTimeout(f"Missing expected files in {file_type} zips")
               
           return True

   @task
   def cleanup_xcom():
       with create_session() as session:
           session.query(XCom).filter(
               XCom.execution_date <= datetime.now() - timedelta(days=30)
           ).delete()

   @task(task_id=f'unzip_{file_type}')
   def unzip_files():
       with ssh_hook.get_conn() as ssh_client:
           cmds = []
           for item in config:
               cmds.extend([
                   f"mkdir -p /path/to/{file_type}/temp",
                   f"cd /path/to/{file_type}/files",
                   f"unzip -o {item['zip_file']} -d /path/to/{file_type}/temp"
               ])
           ssh_client.exec_command(" && ".join(cmds))

   @task(task_id=f'move_{file_type}_to_prod')
   def move_to_prod():
       with ssh_hook.get_conn() as ssh_client:
           cmds = [
               f"mkdir -p /path/to/{file_type}/prod",
               f"mv /path/to/{file_type}/temp/* /path/to/{file_type}/prod/",
               f"rm -rf /path/to/{file_type}/temp"
           ]
           ssh_client.exec_command(" && ".join(cmds))

   @task(task_id=f'cleanup_{file_type}_work_dir')
   def cleanup_work_dir():
       with ssh_hook.get_conn() as ssh_client:
           ssh_client.exec_command(f"rm -rf /path/to/{file_type}/temp")

   sensor_task = file_sensor()
   unzip = unzip_files()
   move = move_to_prod()
   cleanup = cleanup_work_dir()
   
   sensor_task >> unzip >> move
   sensor_task.on_failure_callback = cleanup

   return move

@dag(
   dag_id='file_processing_workflow',
   schedule='@daily',
   start_date=datetime(2025, 1, 1),
   catchup=False,
   tags=['file_processing'],
   default_args={
       'owner': 'airflow',
       'retries': 1,
       'retry_delay': timedelta(minutes=5),
       'email_on_failure': True
   }
)
def file_processing_workflow():
   vector_config = load_config('vector_config.yml')
   control_config = load_config('control_config.yml')

   vector_move = process_files(vector_config, 'vector')
   control_move = process_files(control_config, 'control')

   @task
   def archive_files():
       with ssh_hook.get_conn() as ssh_client:
           date_suffix = "$(date +%Y%m%d)"
           cmds = [
               "cd /path/to/vector/prod",
               f"tar -czf /path/to/archive/vector_{date_suffix}.tar.gz *",
               "cd /path/to/control/prod",
               f"tar -czf /path/to/archive/control_{date_suffix}.tar.gz *"
           ]
           ssh_client.exec_command(" && ".join(cmds))

   archive = archive_files()
   [vector_move, control_move] >> archive

dag = file_processing_workflow()