import json
import sys
import time
import re
import requests
from datetime import datetime, timedelta

from airflow import DAG
from airflow.models import Variable
from airflow.operators.python_operator import PythonOperator


# from infromatic team, : connect to infromatica (akash)

# Utility to make task_id Airflow-safe
def sanitize_task_id(task_id: str) -> str:
    sanitized = re.sub(r'[^a-zA-Z0-9_]+', '_', task_id)
    if not re.match(r'^[a-zA-Z0-9]', sanitized):
        sanitized = 'task_' + sanitized
    return sanitized

# Fetch parameters from Airflow Variables
iics_username = Variable.get("iics_username")
iics_password = Variable.get("iics_password")
task_type = Variable.get("task_type", default_var="MTT")
base_url = Variable.get("iics_base_url", default_var="")

# Task name
CDI_task_name = "CDI_task"

# Default DAG args
default_args = {
   'owner': 'infa',
   'depends_on_past': False,
   'email': ['airflow@example.com'],
   'email_on_failure': False,
   'email_on_retry': False,
   'retries': 1,
   'retry_delay': timedelta(minutes=1),
   'start_date': datetime.now() - timedelta(seconds=10),
}

# API logic (same as before)
def get_session_id(un, pw):
    data = {'@type': 'login', 'username': un, 'password': pw}
    headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}
    r = requests.post(base_url, data=json.dumps(data), headers=headers)
    if r.status_code == 200:
        return r.json()["icSessionId"], r.json()["serverUrl"]
    else:
        print('API call failed:', r.status_code)
        print(r.text)
        sys.exit(1)

def start_job(session_id, server_url, taskname, taskType):
    job_start_url = server_url + "/api/v2/job"
    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}
    data = {'@type': 'job', 'taskName': taskname, 'taskType': taskType}
    r = requests.post(job_start_url, data=json.dumps(data), headers=headers)
    if r.status_code == 200:
        response_content = r.json()
        print(f"Job {taskname} started successfully")
        return response_content['taskId'], response_content['runId']
    else:
        print('Job failed to start:', r.status_code)
        print(r.text)
        sys.exit(1)

def get_status(server_url, session_id, task_id, run_id):
    job_activity_url = server_url + "/api/v2/activity/activityMonitor"
    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}
    r = requests.get(job_activity_url, headers=headers)
    if r.status_code == 200:
        for obj in r.json():
            if obj['taskId'] == task_id and obj['runId'] == run_id:
                return obj['executionState']
    else:
        print('Failed to get status:', r.status_code)
        print(r.text)
        sys.exit(1)

def execute_task(task_name):
    session_id, server_url = get_session_id(iics_username, iics_password)
    task_id, run_id = start_job(session_id, server_url, task_name, task_type)
    
    log_url = f"{server_url}/api/v2/activity/activityLog/"
    headers = {'Content-Type': 'application/json', 'icSessionId': session_id, 'Accept': 'application/json'}

    while True:
        time.sleep(15)
        status = get_status(server_url, session_id, task_id, run_id)
        print(f"Task status: {status}")
        if status not in {"RUNNING", "INITIALIZED", "STOPPING", "QUEUED"}:
            # Fetch logs on completion
            url = f"{log_url}?taskId={task_id}&runId={run_id}"
            r = requests.get(url, headers=headers)
            logs = r.json()
            for obj in logs:
                log_id = obj['id']
                log_detail = requests.get(f"{log_url}{log_id}/sessionLog", headers=headers)
                print(log_detail.text)
            break

# DAG with no schedule (manual trigger)
dag = DAG(
   'IDMC_Airflow_Test',
   default_args=default_args,
   description='Simplified DAG with one CDI task',
   schedule_interval=None,
   catchup=False
)

safe_task_id = sanitize_task_id(CDI_task_name)

run_cdi_task = PythonOperator(
    task_id=safe_task_id,
    python_callable=execute_task,
    op_kwargs={'task_name': CDI_task_name},
    dag=dag
)