# -*- coding: utf-8 -*-
#
# Copyright © 2024 rl
#
# Distributed under terms of the MIT license.

"""
Maintenance DAG for managing the lifecycle of ytdlp-ops accounts.
This DAG is responsible for:
- Un-banning accounts whose ban duration has expired.
- Transitioning accounts from RESTING to ACTIVE after their cooldown period.
- Transitioning accounts from ACTIVE to RESTING after their active duration.
This logic was previously handled inside the ytdlp-ops-server and has been
moved here to give the orchestrator full control over account state.
"""
from __future__ import annotations

import logging
import time
from datetime import datetime

from airflow.decorators import task
from airflow.models import Variable
from airflow.models.dag import DAG
from airflow.utils.dates import days_ago

# Import utility functions and Thrift modules
from utils.redis_utils import _get_redis_client
from pangramia.yt.tokens_ops import YTTokenOpService
from thrift.protocol import TBinaryProtocol
from thrift.transport import TSocket, TTransport

# Configure logging
logger = logging.getLogger(__name__)

# Default settings from Airflow Variables or hardcoded fallbacks
DEFAULT_REDIS_CONN_ID = 'redis_default'
DEFAULT_YT_AUTH_SERVICE_IP = Variable.get("YT_AUTH_SERVICE_IP", default_var="172.17.0.1")
DEFAULT_YT_AUTH_SERVICE_PORT = Variable.get("YT_AUTH_SERVICE_PORT", default_var=9080)

DEFAULT_ARGS = {
    'owner': 'airflow',
    'retries': 1,
    'retry_delay': 30,
    'queue': 'maintenance',
}


# --- Helper Functions ---

def _get_thrift_client(host, port, timeout=60):
    """Helper to create and connect a Thrift client."""
    transport = TSocket.TSocket(host, port)
    transport.setTimeout(timeout * 1000)
    transport = TTransport.TFramedTransport(transport)
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    client = YTTokenOpService.Client(protocol)
    transport.open()
    logger.info(f"Connected to Thrift server at {host}:{port}")
    return client, transport


@task
def manage_account_states():
    """
    Fetches all account statuses and performs necessary state transitions.
    """
    host = DEFAULT_YT_AUTH_SERVICE_IP
    port = int(DEFAULT_YT_AUTH_SERVICE_PORT)
    redis_conn_id = DEFAULT_REDIS_CONN_ID

    client, transport = None, None
    try:
        client, transport = _get_thrift_client(host, port)
        redis_client = _get_redis_client(redis_conn_id)
        
        logger.info("Fetching all account statuses from the service...")
        all_accounts = client.getAccountStatus(accountPrefix=None)
        logger.info(f"Found {len(all_accounts)} accounts to process.")

        accounts_to_unban = []
        accounts_to_activate = []
        accounts_to_rest = []

        for acc in all_accounts:
            if acc.status == "BANNED (expired)":
                accounts_to_unban.append(acc.accountId)
            elif acc.status == "RESTING (expired)":
                accounts_to_activate.append(acc.accountId)
            elif acc.status == "ACTIVE (should be resting)":
                accounts_to_rest.append(acc.accountId)

        # --- Perform State Transitions ---
        
        # 1. Un-ban accounts via Thrift call
        if accounts_to_unban:
            logger.info(f"Un-banning {len(accounts_to_unban)} accounts: {accounts_to_unban}")
            for acc_id in accounts_to_unban:
                try:
                    client.unbanAccount(acc_id, "Automatic un-ban by Airflow maintenance DAG.")
                    logger.info(f"Successfully un-banned account '{acc_id}'.")
                except Exception as e:
                    logger.error(f"Failed to un-ban account '{acc_id}': {e}")
        
        # 2. Activate resting accounts via direct Redis write
        if accounts_to_activate:
            logger.info(f"Activating {len(accounts_to_activate)} accounts: {accounts_to_activate}")
            now_ts = int(time.time())
            with redis_client.pipeline() as pipe:
                for acc_id in accounts_to_activate:
                    key = f"account_status:{acc_id}"
                    pipe.hset(key, "status", "ACTIVE")
                    pipe.hset(key, "active_since_timestamp", now_ts)
                    pipe.hset(key, "status_changed_timestamp", now_ts)
                pipe.execute()
            logger.info("Finished activating accounts.")

        # 3. Rest active accounts via direct Redis write
        if accounts_to_rest:
            logger.info(f"Putting {len(accounts_to_rest)} accounts to rest: {accounts_to_rest}")
            now_ts = int(time.time())
            with redis_client.pipeline() as pipe:
                for acc_id in accounts_to_rest:
                    key = f"account_status:{acc_id}"
                    pipe.hset(key, "status", "RESTING")
                    pipe.hset(key, "status_changed_timestamp", now_ts)
                pipe.execute()
            logger.info("Finished putting accounts to rest.")

    finally:
        if transport and transport.isOpen():
            transport.close()


with DAG(
    dag_id='ytdlp_ops_account_maintenance',
    default_args=DEFAULT_ARGS,
    schedule='*/5 * * * *',  # Run every 5 minutes
    start_date=days_ago(1),
    catchup=False,
    tags=['ytdlp', 'maintenance'],
    doc_md=__doc__,
) as dag:
    manage_account_states()