import logging
import requests
import pyodbc
import pandas as pd
import sys
import traceback
import argparse
import os
from datetime import datetime
import re


# ---------- NEW: arg parsing & helpers ----------
from decimal import Decimal, ROUND_HALF_UP, InvalidOperation


def _get_sql_schema(cursor, schema, table_name):
    cols = []
    for row in cursor.columns(table=table_name, schema=schema):
        # row.type_name, row.column_name, row.column_size, row.decimal_digits
        cols.append({
            "name": row.column_name,
            "type": (row.type_name or "").lower(),
            "size": int(row.column_size or 0),
            "scale": int(row.decimal_digits or 0),
        })
    return cols


def _clamp_and_round_decimals(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    """Round DECIMAL/NUMERIC columns to their scale and NULL-out out-of-range values."""
    out = df.copy()
    for colmeta in sql_cols:
        if colmeta["type"] not in {"decimal", "numeric"}:
            continue
        col = colmeta["name"]
        if col not in out.columns:
            continue
        prec = colmeta["size"]
        scale = colmeta["scale"]
        # Maximum absolute value that fits: (10^(prec-scale)) - 10^-scale
        max_abs = (Decimal(10) ** (prec - scale)) - (Decimal(1) / (Decimal(10) ** scale))
        q = Decimal("1." + ("0" * scale))  # quantize target, e.g. 7 zeros for scale 7

        def fix(v):
            if v is None or (isinstance(v, float) and pd.isna(v)):
                return None
            try:
                d = Decimal(str(v)).quantize(q, rounding=ROUND_HALF_UP)
                if d.copy_abs() > max_abs:
                    return None  # too large -> NULL it rather than blow up the batch
                return d
            except (InvalidOperation, ValueError):
                return None

        out[col] = out[col].map(fix).astype(object)
    return out


def parse_args():
    p = argparse.ArgumentParser(description="API to SQL migrator with SSIS metadata")
    p.add_argument("--process-logid", type=int, default=os.getenv("PROCESSlogid"))
    p.add_argument("--package-logid", type=int, default=os.getenv("PACKAGElogid"))
    p.add_argument("--data-logid", type=int, default=os.getenv("DATAlogid"))
    p.add_argument("--etl-datetime", default=os.getenv("ETLdatetime"))  # ISO8601 or SQL-convertible
    p.add_argument("--etl-auditname", default=os.getenv("ETLauditname"))
    # your existing params can also be added here if you prefer CLI over hardcoding
    return p.parse_args()


def normalize_etl_datetime(value):
    if value is None or value == "":
        # fall back to "now" if not provided
        return datetime.now().isoformat(timespec="seconds")
    try:
        # accept a lot of common shapes
        return pd.to_datetime(value).isoformat()
    except Exception:
        return str(value)


def add_etl_columns(df: pd.DataFrame, meta: dict) -> pd.DataFrame:
    """
    Ensures the five SSIS/ETL columns exist in df as constants.
    If any already exist from the API, we don't overwrite them unless they're all-null.
    """
    # normalize ETLdatetime once
    meta = meta.copy()
    meta["ETLdatetime"] = normalize_etl_datetime(meta.get("ETLdatetime"))

    for col in ("PROCESSlogid", "PACKAGElogid", "DATAlogid", "ETLdatetime", "ETLauditname"):
        val = meta.get(col)
        if col not in df.columns:
            df[col] = val
        else:
            # if present but empty, fill with the meta value
            if df[col].isna().all():
                df[col] = val
    return df


def configure_logging(audit_name: str | None):
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s" + (f" [audit={audit_name}]" if audit_name else "")
    )


# ---------- existing functions (minor edits only) ----------

def get_token(auth_server_url: str, client_id: str, secret: str):
    try:
        token_req = {'client_id': client_id,
                     'client_secret': secret,
                     'grant_type': 'client_credentials'}
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        try:
            token_response = requests.post(auth_server_url, data=token_req, headers=headers)
            token_response.raise_for_status()
        except requests.exceptions.RequestException:
            print(traceback.format_exc())
            sys.exit(1)
        tokens = token_response.json()
        return tokens.get('access_token')
    except:
        print(traceback.format_exc())


def _page_and_insert(api_url: str, schema: str, table_name: str, token: str, conn_string: str,
                     date_fields: tuple[str, ...], convert_to_nz: bool, meta: dict):
    try:
        api_call_headers = {'Authorization': f'Bearer {token}'}
        page = 1
        try:
            conn = pyodbc.connect(conn_string)
            cursor = conn.cursor()
        except:
            print(traceback.format_exc())
            sys.exit(1)

        try:
            cursor.execute(f'TRUNCATE TABLE [{schema}].[{table_name}]')
            while True:
                params = {'pageSize': 10000, 'page': page}
                try:
                    api_call_response = requests.get(api_url, headers=api_call_headers, params=params)
                except:
                    print(traceback.format_exc())
                    sys.exit(1)

                if api_call_response.status_code != 200:
                    logging.error(f'Return code: {api_call_response.status_code}. Exiting script.')
                    sys.exit(1)

                response_data = api_call_response.json()
                if 'DataSet' not in response_data:
                    logging.error('No data in response.')
                    sys.exit(1)

                temp_df = pd.DataFrame(response_data['DataSet'])
                if temp_df.empty:
                    break

                # ---- NEW: add the SSIS/ETL metadata columns
                temp_df = add_etl_columns(temp_df, meta)

                # date handling
                for field in date_fields:
                    if field in temp_df.columns:
                        try:
                            # Parse as UTC to be safe (time-only strings like "14:40:39" become today's date UTC)
                            temp_df[field] = pd.to_datetime(temp_df[field], errors='coerce', utc=True)

                            if pd.api.types.is_datetime64_any_dtype(temp_df[field]):
                                if convert_to_nz:
                                    # Convert from UTC→Pacific/Auckland when asked
                                    temp_df[field] = temp_df[field].dt.tz_convert('Pacific/Auckland')

                                # ALWAYS drop timezone (pyodbc needs naive datetimes for datetime2)
                                temp_df[field] = temp_df[field].dt.tz_localize(None)
                        except Exception:
                            print(traceback.format_exc())
                numeric_like = [c for c in temp_df.columns if re.match(r"^NUMBER\d+$", c, re.I)]
                for extra in ("INFORMATIONID", "ENTITYID", "PROCESSlogid", "PACKAGElogid", "DATAlogid"):
                    if extra in temp_df.columns:
                        numeric_like.append(extra)

                if numeric_like:
                    temp_df[numeric_like] = temp_df[numeric_like].apply(pd.to_numeric, errors="coerce")

                # ---- turn empty strings into NULL, and replace NaN/NaT with None (pyodbc-friendly)
                temp_df.replace({"": pd.NA, " ": pd.NA}, inplace=True)
                temp_df = temp_df.where(pd.notnull(temp_df), None)
                sql_cols = _get_sql_schema(cursor, schema, table_name)
                temp_df = _clamp_and_round_decimals(temp_df, sql_cols)
                col_list = list(temp_df.columns)
                column_names = ','.join(f'[{c}]' for c in col_list)
                placeholders = ','.join('?' for _ in col_list)

                try:
                    cursor.fast_executemany = True
                    cursor.executemany(
                        f'INSERT INTO [{schema}].[{table_name}]({column_names}) VALUES({placeholders})',
                        temp_df.values.tolist()
                    )
                except pyodbc.DataError as e:
                    logging.error(f"Batch insert failed with {e}; probing rows…")
                    for i, row in enumerate(temp_df.itertuples(index=False, name=None)):
                        try:
                            cursor.execute(
                                f'INSERT INTO [{schema}].[{table_name}]({column_names}) VALUES({placeholders})',
                                row
                            )
                        except pyodbc.DataError as e2:
                            types = {c: type(v).__name__ for c, v in zip(col_list, row)}
                            logging.error(f"Row {i} failed: {e2}; types={types}; values={row}")
                            raise
                conn.commit()
                page += 1

            logging.info('Closing connection to db')
            conn.close()
        except:
            print(traceback.format_exc())
    except:
        print(traceback.format_exc())


def migrate_utc(api_url: str, schema: str, table_name: str, token: str, conn_string: str, *date_field: str,
                meta: dict | None = None):
    _page_and_insert(api_url, schema, table_name, token, conn_string, date_field, True, meta or {})


def migrate_nzst(api_url: str, schema: str, table_name: str, token: str, conn_string: str, *date_field: str,
                 meta: dict | None = None):
    # your original "NZST" version left incoming times at UTC; here we mirror that behavior (no tz convert)
    _page_and_insert(api_url, schema, table_name, token, conn_string, date_field, False, meta or {})


def migrate(api_url: str, schema: str, table_name: str, token: str, conn_string: str, *date_field: str,
            meta: dict | None = None):
    print('migrate will be deprecated in future releases of toapigdc, please use migrate_utc or migrate_nzst')
    _page_and_insert(api_url, schema, table_name, token, conn_string, date_field, True, meta or {})


# ---------- entrypoint ----------
if __name__ == '__main__':
    args = parse_args()
    configure_logging(args.etl_auditname)

    etl_meta = {'PROCESSlogid': args.process_logid, 'PACKAGElogid': args.package_logid, 'DATAlogid': args.data_logid,
                'ETLdatetime': args.etl_datetime, 'ETLauditname': args.etl_auditname}

    # Example call (replace with your real values or wire these up to args too)
    # token = get_token(AUTH_URL, CLIENT_ID, CLIENT_SECRET)
    # migrate_utc(API_URL, "dbo", "TargetTable", token, CONN_STR, "SomeDateField", meta=etl_meta)
