import streamlit as st
import pandas as pd
import boto3
from io import BytesIO
from datetime import datetime
import os

# =========================================================
# PAGE CONFIG
# =========================================================

st.set_page_config(
    page_title="SCD2 Incremental Pipeline",
    page_icon="📊",
    layout="wide"
)

st.title("📊 S3 Incremental SCD2 Pipeline")

# =========================================================
# AWS CONFIG
# =========================================================

AWS_REGION = "ap-south-1"

BUCKET_NAME = "trial-bucket-treat-1"

# =========================================================
# S3 CLIENT
# =========================================================

s3_client = boto3.client(
    "s3",
    region_name=AWS_REGION
)

# =========================================================
# TABLE CONFIG
# =========================================================

TABLES = {

    "customer": {

        "business_key": "cust_id",

        "columns": [
            "cust_id",
            "Phone",
            "First Name",
            "Last Name",
            "Address",
            "Pincode"
        ],

        "s3_key":
        "output/parquet-data/customer.parquet"
    },

    "loan": {

        "business_key": "cust_id",

        "columns": [
            "cust_id",
            "loan_id",
            "Sanction_Loan",
            "Cibil_Score",
            "Start_Date",
            "Disbursed_Date",
            "Location",
            "Pincode",
            "Loan_Type",
            "Unit"
        ],

        "s3_key":
        "output/parquet-data/loan.parquet"
    },

    "document": {

        "business_key": "cust_id",

        "columns": [
            "cust_id",
            "KYC_Status",
            "Annual_Income",
            "Nationality",
            "Application_Type"
        ],

        "s3_key":
        "output/parquet-data/document.parquet"
    }
}

# =========================================================
# SELECT TABLE
# =========================================================

selected_table = st.selectbox(
    "Select Table",
    list(TABLES.keys())
)

table_config = TABLES[selected_table]

BUSINESS_KEY = table_config["business_key"]
COLUMNS = table_config["columns"]
S3_KEY = table_config["s3_key"]

# =========================================================
# INPUT FORM
# =========================================================

st.subheader("Enter New Record")

new_record = {}

with st.form("data_form"):

    for column in COLUMNS:

        value = st.text_input(column)

        new_record[column] = value

    submit_button = st.form_submit_button(
        "Submit Data"
    )

# =========================================================
# PIPELINE LOGIC
# =========================================================

if submit_button:

    st.info("Running incremental SCD2 pipeline...")

    # -----------------------------------------------------
    # TYPE CONVERSION
    # -----------------------------------------------------

    int_columns = [
        "loan_id",
        "Pincode",
        "Annual_Income",
        "Cibil_Score",
        "Sanction_Loan",
        "Phone"
    ]

    for col in int_columns:

        if col in new_record and new_record[col] != "":

            new_record[col] = int(new_record[col])

    # -----------------------------------------------------
    # READ PARQUET
    # -----------------------------------------------------

    try:

        response = s3_client.get_object(
            Bucket=BUCKET_NAME,
            Key=S3_KEY
        )

        parquet_data = BytesIO(
            response["Body"].read()
        )

        target_df = pd.read_parquet(parquet_data)

        st.success("Parquet loaded from S3")

    except Exception:

        st.warning("Parquet not found. Creating new dataset")

        target_df = pd.DataFrame(columns=COLUMNS)

    # -----------------------------------------------------
    # ADD SCD2 COLUMNS
    # -----------------------------------------------------

    if "start_date" not in target_df.columns:
        target_df["start_date"] = datetime.now()

    if "end_date" not in target_df.columns:
        target_df["end_date"] = pd.NaT

    if "is_current" not in target_df.columns:
        target_df["is_current"] = "Y"

    if "etl_timestamp" not in target_df.columns:
        target_df["etl_timestamp"] = datetime.now()

    # -----------------------------------------------------
    # FIND ACTIVE RECORD
    # -----------------------------------------------------

    existing_rows = target_df[
        (
            target_df[BUSINESS_KEY]
            ==
            new_record[BUSINESS_KEY]
        )
        &
        (
            target_df["is_current"] == "Y"
        )
    ]

    # -----------------------------------------------------
    # ADD METADATA
    # -----------------------------------------------------

    new_record["start_date"] = datetime.now()

    new_record["end_date"] = pd.NaT

    new_record["is_current"] = "Y"

    new_record["etl_timestamp"] = datetime.now()

    # -----------------------------------------------------
    # INSERT OR UPDATE
    # -----------------------------------------------------

    if existing_rows.empty:

        st.success("New record inserted")

        new_df = pd.DataFrame([new_record])

        final_df = pd.concat(
            [target_df, new_df],
            ignore_index=True
        )

    else:

        current_record = existing_rows.iloc[0]

        compare_columns = [
            col for col in COLUMNS
            if col != BUSINESS_KEY
        ]

        is_same = True

        for col in compare_columns:

            if str(new_record[col]) != str(current_record[col]):

                is_same = False
                break

        if is_same:

            st.warning("No changes detected")

            final_df = target_df

        else:

            st.success("Changes detected — applying SCD2")

            target_df.loc[
                (
                    (
                        target_df[BUSINESS_KEY]
                        ==
                        new_record[BUSINESS_KEY]
                    )
                    &
                    (
                        target_df["is_current"] == "Y"
                    )
                ),
                "is_current"
            ] = "N"

            target_df.loc[
                (
                    (
                        target_df[BUSINESS_KEY]
                        ==
                        new_record[BUSINESS_KEY]
                    )
                    &
                    (
                        target_df["is_current"] == "N"
                    )
                ),
                "end_date"
            ] = datetime.now()

            new_df = pd.DataFrame([new_record])

            final_df = pd.concat(
                [target_df, new_df],
                ignore_index=True
            )

    # -----------------------------------------------------
    # REMOVE DUPLICATES
    # -----------------------------------------------------

    final_df = final_df.drop_duplicates()

    # -----------------------------------------------------
    # WRITE BACK TO S3
    # -----------------------------------------------------

    parquet_buffer = BytesIO()

    final_df.to_parquet(
        parquet_buffer,
        index=False,
        engine="pyarrow"
    )

    s3_client.put_object(
        Bucket=BUCKET_NAME,
        Key=S3_KEY,
        Body=parquet_buffer.getvalue()
    )

    st.success("Updated parquet uploaded to S3")

    # -----------------------------------------------------
    # SHOW DATA
    # -----------------------------------------------------

    st.subheader("Updated Dataset")

    st.dataframe(final_df.tail(20))