import logging
import requests
import pyodbc
import pandas as pd
import sys
import traceback
import argparse
import os
from datetime import datetime
import re
import numpy as np
from decimal import Decimal, ROUND_HALF_UP, InvalidOperation

# -----------------------------
# Config / constants
# -----------------------------
FLOAT_TYPES = {"float", "real"}

# -----------------------------
# Float sanitizer
# -----------------------------
def _sanitize_floats(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    for colmeta in sql_cols:
        if colmeta["type"] not in FLOAT_TYPES:
            continue
        col = colmeta["name"]
        if col not in out.columns:
            continue
        s = pd.to_numeric(out[col], errors="coerce")        # strings -> numeric or NaN
        s = s.where(np.isfinite(s), np.nan)                 # inf/-inf -> NaN
        s = s.astype(object).where(~pd.isna(s), None)       # None for NULL
        out[col] = s
    return out

# -----------------------------
# Schema helpers
# -----------------------------
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"]
        # Max 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))  # e.g., scale 7 -> "1.0000000"

        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
                return d
            except (InvalidOperation, ValueError):
                return None

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

# -----------------------------
# CLI / logging helpers
# -----------------------------
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"))
    return p.parse_args()

def normalize_etl_datetime(value):
    if value is None or value == "":
        return datetime.now().isoformat(timespec="seconds")
    try:
        return pd.to_datetime(value).isoformat()
    except Exception:
        return str(value)

def add_etl_columns(df: pd.DataFrame, meta: dict) -> pd.DataFrame:
    """Ensure the five SSIS/ETL columns exist in df (fill when whole column is null)."""
    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 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 "")
    )

# -----------------------------
# Token
# -----------------------------
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())

# -----------------------------
# Preflight diagnostics
# -----------------------------
def _preflight_numeric_ranges(df: pd.DataFrame, sql_cols: list[dict]) -> None:
    INT_MIN, INT_MAX = -2147483648, 2147483647  # SQL INT range
    for meta in sql_cols:
        name, typ = meta["name"], meta["type"]
        if name not in df.columns:
            continue
        s = df[name]
        if typ == "int":
            v = pd.to_numeric(s, errors="coerce")
            if (v > INT_MAX).any() or (v < INT_MIN).any():
                logging.error(f"Preflight: INT out of range in {name}: min={v.min()} max={v.max()}")
        if typ in FLOAT_TYPES:
            v = pd.to_numeric(s, errors="coerce")
            nonfinite = v[~np.isfinite(v)]
            if len(nonfinite):
                logging.error(f"Preflight: non-finite float in {name} at rows {list(nonfinite.index[:5])}")

# -----------------------------
# Core: page + insert
# -----------------------------
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

                # Identity guard
                for ident in ("ExtractKey", "extractKey"):
                    if ident in temp_df.columns:
                        temp_df.drop(columns=[ident], inplace=True)

                # Add SSIS/ETL meta columns
                temp_df = add_etl_columns(temp_df, meta)

                # Date handling
                for field in date_fields:
                    if field in temp_df.columns:
                        try:
                            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:
                                    temp_df[field] = temp_df[field].dt.tz_convert('Pacific/Auckland')
                                temp_df[field] = temp_df[field].dt.tz_localize(None)
                        except Exception:
                            print(traceback.format_exc())

                # Numeric coercion for typical INT-like columns (incl. audit ids)
                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")

                # Force proper Python ints for int-like columns
                for c in ("INFORMATIONID", "ENTITYID", "PROCESSlogid", "PACKAGElogid", "DATAlogid"):
                    if c in temp_df.columns:
                        s = pd.to_numeric(temp_df[c], errors="coerce").astype("Int64")
                        temp_df[c] = s.where(s.notna(), None).map(lambda v: int(v) if v is not None else None)

                # "" -> NULL; NaN/NaT -> None
                temp_df.replace({"": pd.NA, " ": pd.NA}, inplace=True)
                temp_df = temp_df.where(pd.notnull(temp_df), None)

                # Align with SQL schema and sanitize numerics
                sql_cols = _get_sql_schema(cursor, schema, table_name)
                temp_df = _clamp_and_round_decimals(temp_df, sql_cols)
                temp_df = _sanitize_floats(temp_df, sql_cols)

                # Column order = SQL order (+ extras at end)
                sql_order = [c["name"] for c in sql_cols if c["name"] in temp_df.columns]
                extras = [c for c in temp_df.columns if c not in sql_order]
                temp_df = temp_df[sql_order + extras]

                # Preflight diagnostics (optional but useful)
                _preflight_numeric_ranges(temp_df, sql_cols)

                # Build INSERT
                col_list = list(temp_df.columns)
                for i, c in enumerate(col_list, start=1):
                    logging.info(f"COLMAP {i:03d} -> {c}")
                column_names = ','.join(f'[{c}]' for c in col_list)
                placeholders = ','.join('?' for _ in col_list)

                # Insert (with row-probe on failure)
                try:
                    cursor.fast_executemany = True
                    cursor.executemany(
                        f'INSERT INTO [{schema}].[{table_name}]({column_names}) VALUES({placeholders})',
                        temp_df.values.tolist()
                    )
                except pyodbc.Error as e:
                    logging.error(f"Batch insert failed: {e}. Probing rows now...")
                    for r, row in enumerate(temp_df.itertuples(index=False, name=None), start=1):
                        try:
                            cursor.execute(
                                f'INSERT INTO [{schema}].[{table_name}]({column_names}) VALUES({placeholders})',
                                row
                            )
                        except pyodbc.Error as e2:
                            logging.error(f"Row {r} failed with: {e2}")
                            for i, (c, v) in enumerate(zip(col_list, row), start=1):
                                logging.error(f"  #{i:03d} {c}: {repr(v)} ({type(v).__name__})")
                            # Return so logs persist and job doesn't suppress the details
                            return
                    # If no single row failed, re-raise
                    raise

                conn.commit()
                page += 1

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

# -----------------------------
# Public entry points
# -----------------------------
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):
    # Keep incoming times at UTC (no convert), still strip tz to insert into datetime2
    _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 {})

# -----------------------------
# Script mode
# -----------------------------
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)
