try:
    from datetime import timedelta
    from datetime import datetime, timezone
    from dateutil.relativedelta import relativedelta

    from airflow import DAG
    from airflow.operators.python_operator import PythonOperator
    from airflow.operators.http_operator import SimpleHttpOperator
    from airflow.operators.postgres_operator import PostgresOperator

    import pandas as pd
    from pandas.io.json import json_normalize
    from pandas.tseries.offsets import CustomBusinessDay
    
    import json
    import requests
    import numpy as np
    import psycopg2

    from sqlalchemy import create_engine, text



except Exception as e:
    print(f"An error occurred: {e}")

dRoW_api_end_url = "https://drow.cloud"

settings = {"mid_status_name":"Review Record by", "end_status_name":"End Case "}

def parse_any_dt(x):
    if not x:
        return None
    # common dRoW patterns you’ve seen
    for fmt in ("%Y-%m-%dT%H:%M:%S.%f%z",
                "%Y-%m-%dT%H:%M:%S%z",
                "%Y-%m-%dT%H:%M:%S.%fZ",
                "%Y-%m-%dT%H:%M:%SZ"):
        try:
            dt = datetime.strptime(x, fmt)
            # if ends with Z, fmt returns naive sometimes; ensure tz-aware
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt.astimezone(timezone.utc)
        except Exception:
            pass

    # last resort: try python dateutil (if you want)
    try:
        from dateutil.parser import parse
        dt = parse(x)
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt.astimezone(timezone.utc)
    except Exception:
        return None

def getDrowToken(**context):
    print("let get drow token")
    response = requests.post(
        url=f"{dRoW_api_end_url}/api/auth/authenticate",
        data={
            "username": "icwp2@drow.cloud",
            "password": "dGVzdDAxQHRlc3QuY29t"
        }
    ).json()
    context["ti"].xcom_push(key="token", value=response['token'])
    print("we got drow token")

def evaluate_kpi_condition(kpi_no: int, value: float):
    """
    Returns (condition, action) for a given KPI number and numeric value.
    condition: "Green" | "Yellow" | "Red"
    action: "" | "Escalate to D1" | "Escalate to D2"
    """
    cond = "Green"
    action = ""

    if kpi_no == 1:
        # KPI 1: days between acceptance & reporting date
        if value <= 30:
            cond, action = "Green", ""
        elif 30 < value < 90:
            cond, action = "Yellow", "Escalate to D1"
        else:  # >= 90
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 2:
        # KPI 2: PWDD % for target cost contract
        if value <= 30:
            cond, action = "Green", ""
        elif 30 < value < 50:
            cond, action = "Yellow", "Escalate to D1"
        else:  # >= 50
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 3:
        # KPI 3: Difference of total of the Prices (%)
        if value <= 5:
            cond, action = "Green", ""
        elif 5 < value < 10:
            cond, action = "Yellow", "Escalate to D1"
        else:  # >= 10
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 4:
        # KPI 4: % of Pain Share / total of Prices (can be negative)
        if value >= 5:
            cond, action = "Green", ""
        elif -10 < value < 5:
            cond, action = "Yellow", "Escalate to D1"
        else:  # <= -10
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 5:
        if value == 0:
            cond, action = "Green", ""
        elif value > 0:
            cond, action = "Yellow", "Escalate to D1"

    elif kpi_no == 6:
        if value == 0:
            cond, action = "Green", ""
        elif value > 0:
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 7:
        if value == 0:
            cond, action = "Green", ""
        elif value > 0:
            cond, action = "Yellow", "Escalate to D1"

    elif kpi_no == 8:
        if value == 0:
            cond, action = "Green", ""
        elif value > 0:
            cond, action = "Red", "Escalate to D2"

    elif kpi_no == 9:
        # KPI 9: No. of Open EW > 3 months
        if value == 0:
            cond, action = "Green", ""
        else:  # >0
            cond, action = "Yellow", "Escalate to D1"

    elif kpi_no == 10:
        # KPI 10: No. of Open EW > 6 months
        if value == 0:
            cond, action = "Green", ""
        else:  # >0
            cond, action = "Yellow", "Escalate to D2"

    return cond, action

def get_dw_engine():
    host = "drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com"
    user = "dRowAdmin"
    pwd = "drowsuper"
    db = "drowDateWareHouse"
    port = "5432"
    return create_engine(f"postgresql://{user}:{pwd}@{host}:{port}/{db}", pool_pre_ping=True)

def getdrowPSQLConnectionString():
    host = 'drowdatewarehose.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'
    dbUserName = 'dRowAdmin'
    dbUserPassword = 'drowsuper'
    database = 'drowDateWareHouse'
    charSet = 'utf8mb4'
    port = '5432'
    conn_string = ('postgress://' +dbUserName+':'+dbUserPassword+'@'+host+':'+port+'/'+database)
    return conn_string

def parse_iso_date(date_str):
    """Safely parse ISO 8601 with or without milliseconds."""
    try:
        return datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%fZ")
    except:
        return datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%SZ")

def getProgrammeData(**context):
    token = context.get("ti").xcom_pull(key="token")

    response = requests.get(
        url=f"{dRoW_api_end_url}/api/sheets/63fd698ed6779f0c607a677c?with_records=true&fields=",
        headers={"x-access-token": f"Bearer {token}"}
    )

    sheet = json.loads(response.text)
    records = sheet["record"]

    target_col = "Programme Approval Elapsed Time (days)"

    latest_id = -1
    latest_elapsed_time = None

    for record in records:
        # Build column map for easy access
        col_map = {v["colName"]: v.get("value") for v in record.get("values", [])}

        raw_id = col_map.get("id")
        if raw_id is None:
            continue

        try:
            numeric_id = int(raw_id)
        except ValueError:
            continue

        # If this is the largest id so far → update answer
        if numeric_id > latest_id:
            latest_id = numeric_id
            latest_elapsed_time = col_map.get(target_col)

    print(f"Largest Programme ID = {latest_id}")
    print(f"Elapsed Time = {latest_elapsed_time}")

    context["ti"].xcom_push(key="kpi_1", value=latest_elapsed_time)

    return latest_elapsed_time

def getContractData(**context):
    print("going to get Contract data now!")
    ti = context.get("ti")
    token = ti.xcom_pull(key="token")

    response = requests.get(
        url=f"{dRoW_api_end_url}/api/sheets/63fd68e49f48080c646e7f32?with_records=true&fields=",
        headers={"x-access-token": f"Bearer {token}"}
    )

    sheet = response.json()
    records = sheet["record"]

    # Column names in dRoW – confirm spelling with your sheet headers
    col_id = "id"
    col_cumulative = "Cumulative PWDD"
    col_forecast_pwdd = "Forecast of the final Prices for the Work Done to Date (PWDD)"

    col_latest_est_contractor = "Latest Estimate by the Contractor"
    col_latest_forecast_total = "Latest Forecast Total of the Prices"

    col_total_pain = "Total Pain"
    col_contractor_pain_share = "Contractor Pain Share"  # might be 0.5 or "50%"
    col_forecast_final_total = "Forecast of the final total of the Prices"

    latest_id = -1
    latest_map = None  # will hold the col_map for the max id row

    def to_float(val):
        """Convert dRoW string/number/percent to float safely."""
        if val in (None, "", "NA"):
            return None
        if isinstance(val, (int, float)):
            return float(val)

        s = str(val).strip()
        if s.endswith("%"):
            try:
                return float(s[:-1].replace(",", "").strip()) / 100.0
            except ValueError:
                return None

        try:
            return float(s.replace(",", ""))
        except ValueError:
            return None

    for record in records:
        col_map = {v["colName"]: v.get("value") for v in record.get("values", [])}

        raw_id = col_map.get(col_id)
        if raw_id is None or raw_id == "":
            continue

        try:
            numeric_id = int(raw_id)
        except ValueError:
            continue

        if numeric_id > latest_id:
            latest_id = numeric_id
            latest_map = col_map

    if latest_map is None:
        print("No valid contract record found (no numeric id).")
        return None

    print(f"Using contract record with largest id = {latest_id}")

    # 2) Extract values from that latest row
    # ----- KPI 2 -----
    cumulative_pwdd = to_float(latest_map.get(col_cumulative)) or 0
    forecast_pwdd = to_float(latest_map.get(col_forecast_pwdd)) or 0

    if forecast_pwdd != 0:
        kpi2 = round((cumulative_pwdd / forecast_pwdd) * 100, 2)
    else:
        kpi2 = 0

    # ----- KPI 3 -----
    latest_est_contractor = to_float(latest_map.get(col_latest_est_contractor)) or 0
    latest_forecast_total = to_float(latest_map.get(col_latest_forecast_total)) or 0

    avg = (latest_est_contractor + latest_forecast_total) / 2.0
    print(f"check average: {avg}")

    if avg != 0:
        kpi3 = round(((latest_est_contractor - latest_forecast_total) / avg) * 100, 2)
    else:
        kpi3 = 0

    # ----- KPI 4 -----
    total_pain = to_float(latest_map.get(col_total_pain)) or 0
    contractor_pain_share = to_float(latest_map.get(col_contractor_pain_share)) or 0
    forecast_final_total = to_float(latest_map.get(col_forecast_final_total)) or 0

    if forecast_final_total != 0:
        kpi4 = round(((total_pain * contractor_pain_share) / forecast_final_total) * 100, 2)
    else:
        kpi4 = 0


    print("KPI2 (PWDD ratio %) =", kpi2)
    print("KPI3 (Estimate vs Forecast %) =", kpi3)
    print("KPI4 (Pain share %) =", kpi4)

    kpi_result = {
        "kpi2_ratio_pwdd": kpi2,
        "kpi3_estimate_vs_forecast": kpi3,
        "kpi4_pain_share": kpi4,
    }

    # If you want to push explicitly:
    ti.xcom_push(key="CONTRACT_KPIS", value=kpi_result)

    return kpi_result

def getRiskRegisterData(**context):
    ti = context.get("ti")
    token = ti.xcom_pull(key="token")

    response = requests.get(
        url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/68f72cb8ee99729ef014ca38?export_type=0",
        headers={
            "x-access-token": f"Bearer {token}",
            "ICWPxAccessKey": "5WSD21ICWP_[1AG:4UdI){n=b~"
        }
    )

    risk_register_data = response.json()
    print(f"check rr data: {risk_register_data}")

    # thresholds: 3 months ≈ 90 days, 6 months ≈ 180 days
    THRESHOLD_3M = 90
    THRESHOLD_6M = 180

    today = datetime.now(timezone.utc).date()

    kpi9_count = 0  # Open records with elapsed_days > 90
    kpi10_count = 0 # Open records with elapsed_days > 180

    for item in risk_register_data:
        rec = item.get("data", {}) or {}

        status = rec.get("Status")
        if status != "Open":
            continue  # only consider open records

        ew_str = rec.get("Date of Early Warning")
        if not ew_str:
            continue  # cannot compute elapsed time

        try:
            ew_date = parse_iso_date(ew_str).date()
        except ValueError:
            print(f"Unable to parse Date of Early Warning: {ew_str}")
            continue

        elapsed_days = (today - ew_date).days

        if elapsed_days > THRESHOLD_3M:
            kpi9_count += 1
        if elapsed_days > THRESHOLD_6M:
            kpi10_count += 1

    result = {
        "kpi9_open_gt_3m": kpi9_count,
        "kpi10_open_gt_6m": kpi10_count,
    }

    # If you want to use XCom:
    # ti.xcom_push(key="RISK_REGISTER_KPIS", value=result)

    return result

def getCEImplementationData(**context):
    ti = context["ti"]
    engine = get_dw_engine()

    q_kpi7 = """
    SELECT COALESCE(SUM(total_events), 0) AS v
    FROM c5_health_check_report_ce_merged_metrics
    WHERE metric = 'ce_not_yet_implemented'
      AND label IN (
        '3 - 6 months',
        '6 - 9 months',
        '9 - 12 months',
        '12 - 24 months',
        '> 24 months'
      );
    """

    q_kpi8 = """
    SELECT COALESCE(SUM(total_events), 0) AS v
    FROM c5_health_check_report_ce_merged_metrics
    WHERE metric = 'ce_not_yet_implemented'
      AND label IN (
        '6 - 9 months',
        '9 - 12 months',
        '12 - 24 months',
        '> 24 months'
      );
    """

    with engine.begin() as conn:
        kpi7 = conn.execute(text(q_kpi7)).scalar() or 0
        kpi8 = conn.execute(text(q_kpi8)).scalar() or 0


    kpi7 = int(kpi7)
    kpi8 = int(kpi8)

    print(f"KPI 7 > 3 months: {kpi7}")
    print(f"KPI 7 > 6 months: {kpi8}")

    ti.xcom_push(key="kpi7", value=kpi7)
    ti.xcom_push(key="kpi8", value=kpi8)

    return {"kpi7": kpi7, "kpi8": kpi8} 

def getCNCEData(**context):
    ti = context["ti"]
    token = ti.xcom_pull(key="token")

    response = requests.get(
    url=f"{dRoW_api_end_url}/api/module/document-export/airflow/workflow/637c7d22b38f8ca02f5c49ab?export_type=0",
    headers={
        "x-access-token": f"Bearer {token}",
        "ICWPxAccessKey": "5WSD21ICWP_[1AG:4UdI){n=b~"
        }
    )

    cnceData = response.json()
    print(f"check cnce data: {cnceData}")

    today = datetime.now(timezone.utc)

    rows = []
    for item in cnceData:
        d = (item.get("data") or {})
        event_no = (d.get("NEC Event No.") or "").strip()
        doc_type = (d.get("NEC Doc Type") or "").strip()
        doc_date_raw = d.get("Doc Date")  # boss wants this
        receive_date_raw = d.get("Receive Date")

        # pick a best-effort "doc_datetime"
        doc_dt = parse_any_dt(doc_date_raw) or parse_any_dt(receive_date_raw)

        if not event_no or not doc_type or not doc_dt:
            continue

        rows.append({
            "event_no": event_no,
            "doc_type": doc_type.strip().upper(),
            "doc_dt": doc_dt,
        })

    # 2) group by event_no and find latest doc
    latest_by_event = {}
    for r in rows:
        key = r["event_no"]
        cur = latest_by_event.get(key)
        if (cur is None) or (r["doc_dt"] > cur["doc_dt"]):
            latest_by_event[key] = r

    # 3) unreplied NCE = latest doc is NCE
    unreplied_nce = []
    for event_no, latest in latest_by_event.items():
        if latest["doc_type"].startswith("NCE"):
            unreplied_nce.append(latest)  # keep latest doc row (its doc_dt is the NCE date)

    # 4) KPI rules (calendar months)
    kpi5_cutoff = today - relativedelta(months=3)
    kpi6_cutoff = today - relativedelta(months=6)

    kpi5 = sum(1 for r in unreplied_nce if r["doc_dt"] < kpi5_cutoff)
    kpi6 = sum(1 for r in unreplied_nce if r["doc_dt"] < kpi6_cutoff)

    print("unreplied NCE total:", len(unreplied_nce))
    print("KPI 5:", kpi5)
    print("KPI 6:", kpi6)

    # push into XCom so writeKpiToPostgres can store
    ti.xcom_push(key="kpi5", value=int(kpi5))
    ti.xcom_push(key="kpi6", value=int(kpi6))

    return {"kpi5": int(kpi5), "kpi6": int(kpi6)}

def writeKpiToPostgres(**context):
    ti = context["ti"]

    # --- Pull KPI values from previous tasks ---
    elapsed_time = ti.xcom_pull(task_ids="getProgrammeData")  # KPI 1 (days)
    contractKPIs = ti.xcom_pull(task_ids="getContractData")   # KPI 2–4
    cnceKPIs = ti.xcom_pull(task_ids="getCNCEData") # KPI 5-6
    ceImplementationKPIs = ti.xcom_pull(task_ids="getCEImplementationData") # KPI 7-8
    riskKPIs = ti.xcom_pull(task_ids="getRiskRegisterData")   # KPI 9–10

    kpi1 = int(elapsed_time or 0)
    kpi2 = float(contractKPIs.get("kpi2_ratio_pwdd") or 0)
    kpi3 = float(contractKPIs.get("kpi3_estimate_vs_forecast") or 0)
    kpi4 = float(contractKPIs.get("kpi4_pain_share") or 0)
    kpi5 = int(cnceKPIs.get("kpi5") or 0)
    kpi6 = int(cnceKPIs.get("kpi6") or 0)
    kpi7 = int(ceImplementationKPIs.get("kpi7") or 0)
    kpi8 = int(ceImplementationKPIs.get("kpi8") or 0)
    kpi9 = int(riskKPIs.get("kpi9_open_gt_3m") or 0)
    kpi10 = int(riskKPIs.get("kpi10_open_gt_6m") or 0)

    print(f"kpi5: {kpi5}")
    print(f"kpi5: {kpi6}")

    delete_table_sql_query = """DROP TABLE IF EXISTS nd201905_kpi"""

    create_table_sql_query = """
    CREATE TABLE IF NOT EXISTS nd201905_kpi (
        id SERIAL PRIMARY KEY,
        kpi_no INT NOT NULL,
        kpi_name VARCHAR(50) NOT NULL,
        subject TEXT,
        description TEXT,
        value VARCHAR(100) NOT NULL,
        unit VARCHAR(50),
        condition VARCHAR(10),
        action TEXT,
        created_at TIMESTAMP DEFAULT NOW(),
        updated_at TIMESTAMP DEFAULT NOW()
    );
    """

    # --- Build KPI rows ---
    rows = []

    # KPI 1
    cond, action = evaluate_kpi_condition(1, kpi1)
    rows.append({
        "kpi_no": 1,
        "kpi_name": "KPI No. 1",
        "subject": "Programme",
        "description": "No. of days between the acceptance of programme and the reporting date",
        "value": str(kpi1),
        "unit": "days",
        "condition": cond,
        "action": action,
    })

    # KPI 2
    cond, action = evaluate_kpi_condition(2, kpi2)
    rows.append({
        "kpi_no": 2,
        "kpi_name": "KPI No. 2",
        "subject": "Prices",
        "description": "PWDD % for target cost contract",
        "value": f"{kpi2:.2f}",
        "unit": "%",
        "condition": cond,
        "action": action,
    })

    # KPI 3
    cond, action = evaluate_kpi_condition(3, kpi3)
    rows.append({
        "kpi_no": 3,
        "kpi_name": "KPI No. 3",
        "subject": "Prices",
        "description": "Difference of total of the Prices between PM and Contractor",
        "value": f"{kpi3:.2f}",
        "unit": "%",
        "condition": cond,
        "action": action,
    })

    # KPI 4
    cond, action = evaluate_kpi_condition(4, kpi4)
    rows.append({
        "kpi_no": 4,
        "kpi_name": "KPI No. 4",
        "subject": "Prices",
        "description": "% of Pain Share / total of the Prices (adjusted by fee percentage)",
        "value": f"{kpi4:.2f}",
        "unit": "%",
        "condition": cond,
        "action": action,
    })

    # KPI 5
    cond, action = evaluate_kpi_condition(5,kpi5)
    rows.append({
        "kpi_no": 5,
        "kpi_name": "KPI No. 5",
        "subject": "CNCE",
        "description": "No. > 3 months",
        "value": f"{kpi5}",
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    # KPI 6
    cond, action = evaluate_kpi_condition(6,kpi6)
    rows.append({
        "kpi_no": 6,
        "kpi_name": "KPI No. 6",
        "subject": "CNCE",
        "description": "No. > 6 months",
        "value": f"{kpi6}",
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    # KPI 7
    cond, action = evaluate_kpi_condition(7, kpi7)
    rows.append({
        "kpi_no": 7,
        "kpi_name": "KPI No. 7",
        "subject": "CE Implementation",
        "description": "No. > 3 months",
        "value": f"{kpi7}",
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    # KPI 8
    cond, action = evaluate_kpi_condition(8, kpi8)
    rows.append({
        "kpi_no": 8,
        "kpi_name": "KPI No. 8",
        "subject": "CE Implementation",
        "description": "No. > 6 months",
        "value": f"{kpi8}",
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    # KPI 9
    cond, action = evaluate_kpi_condition(9, kpi9)
    rows.append({
        "kpi_no": 9,
        "kpi_name": "KPI No. 9",
        "subject": "Early Warning",
        "description": "No. of Early Warnings with Elapse time > 3 months",
        "value": str(int(kpi9)),
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    # KPI 10
    cond, action = evaluate_kpi_condition(10, kpi10)
    rows.append({
        "kpi_no": 10,
        "kpi_name": "KPI No. 10",
        "subject": "Early Warning",
        "description": "No. of Early Warnings with Elapse time > 6 months",
        "value": str(int(kpi10)),
        "unit": "no.",
        "condition": cond,
        "action": action,
    })

    kpi_df = pd.DataFrame(rows)

    print(f"check dataframe: {kpi_df}")

    # --- Drop, create, insert ---
    engine = get_dw_engine()
    with engine.begin() as conn:
        conn.execute(delete_table_sql_query)
        conn.execute(create_table_sql_query)

        kpi_df.to_sql(
            "nd201905_kpi",
            con=conn,
            if_exists="append",
            index=False,
        )

    print("nd201905_kpi table refreshed with KPI 1,2,3,4,5,6,7,8.9,10.")

with DAG(
    dag_id="nd201905_kpi",
    schedule_interval="0 7,15 * * *",
    default_args={
        "owner": "airflow",
        "retries": 1,
        "retry_delay": timedelta(minutes=5),
        "start_date": datetime(2023, 1, 17)
    },
    catchup=False) as f:

    getDrowToken = PythonOperator(
        task_id="getDrowToken",
        python_callable=getDrowToken,
        provide_context=True,
        # op_kwargs={"name": "Dylan"}
    )

    getProgrammeData = PythonOperator(
        task_id="getProgrammeData",
        python_callable=getProgrammeData,
        provide_context=True,
        # op_kwargs={"name": "Dylan"},
    )

    getContractData = PythonOperator(
        task_id="getContractData",
        python_callable=getContractData,
        provide_context=True
    )

    getRiskRegisterData = PythonOperator(
        task_id="getRiskRegisterData",
        python_callable=getRiskRegisterData,
        provide_context=True
    )

    getCEImplementationData = PythonOperator(
        task_id="getCEImplementationData",
        python_callable=getCEImplementationData,
        provide_context=True
    )

    getCNCEData = PythonOperator(
        task_id="getCNCEData",
        python_callable=getCNCEData,
        provide_context=True
    )

    writeKpiToPostgres = PythonOperator(
        task_id="writeKpiToPostgres",
        python_callable=writeKpiToPostgres,
        provide_context=True
    )



getDrowToken >> getProgrammeData >> getContractData >> getRiskRegisterData >> getCEImplementationData >> getCNCEData >> writeKpiToPostgres