try:
    from datetime import timedelta
    from datetime import datetime, timezone

    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 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 == 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 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

    print(f"KPI 9 (Open > 3 months): {kpi9_count}")
    print(f"KPI 10 (Open > 6 months): {kpi10_count}")

    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 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
    riskKPIs = ti.xcom_pull(task_ids="getRiskRegisterData")   # KPI 9–10

    kpi1 = float(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)

    kpi9 = float(riskKPIs.get("kpi9_open_gt_3m") or 0)
    kpi10 = float(riskKPIs.get("kpi10_open_gt_6m") or 0)

    # --- DB connection ---
    host           = 'drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'
    dbUserName     = 'dRowAdmin'
    dbUserPassword = 'drowsuper'
    database       = 'drowDateWareHouse'
    port           = "5432"

    conn_string = (
        'postgres://'
        + dbUserName + ':' + dbUserPassword
        + '@' + host + ':' + port
        + '/' + database
    )

    db = create_engine(conn_string)

    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 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 ---
    with db.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,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
    )

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



getDrowToken >> getProgrammeData >> getContractData >> getRiskRegisterData >> writeKpiToPostgres