from airflow.plugins_manager import AirflowPlugin
from airflow.hooks.base import BaseHook
from airflow.configuration import conf
import uuid
import backoff

class YTDLPHook(BaseHook):
    def __init__(self, conn_id='ytdlp_default'):
        super().__init__()
        self.conn_id = conn_id
        self.connection = self.get_connection(conn_id)
        self.timeout = conf.getint('ytdlp', 'timeout', fallback=120)
        self.max_retries = conf.getint('ytdlp', 'max_retries', fallback=3)

    @backoff.on_exception(backoff.expo,
                         Exception,
                         max_tries=3,
                         max_time=300)
    def start_service(self, host, port, service_id, work_dir):
        """Start token service as a long-running process"""
        import subprocess
        import os
        from pathlib import Path

        # Get script path relative to Airflow home
        airflow_home = os.getenv('AIRFLOW_HOME', '')
        script_path = Path(airflow_home).parent / 'ytdlp_ops_server.py'

        # Ensure work directory exists
        os.makedirs(work_dir, exist_ok=True)

        # Start service process
        cmd = [
            'python', str(script_path),
            '--port', str(port),
            '--host', host,
            '--service-id', service_id,
            '--context-dir', work_dir,
            '--script-dir', str(Path(airflow_home) / 'dags' / 'scripts')
        ]

        self.log.info(f"Starting token service: {' '.join(cmd)}")
        
        # Start process detached
        docker_cmd = [
            'docker-compose', '-f', 'docker-compose.yaml',
            'up', '-d', '--build', 'ytdlp-service'
        ]
        subprocess.run(docker_cmd, check=True)

        self.log.info(f"Token service started on {host}:{port}")
        return True

class YTDLPPlugin(AirflowPlugin):
    name = 'ytdlp_plugin'
    hooks = [YTDLPHook]