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

# -----------------------------
# Type sets
# -----------------------------
FLOAT_TYPES   = {"float", "real"}
DEC_TYPES     = {"decimal", "numeric"}
INT_TYPES     = {"int", "bigint", "smallint", "tinyint"}
DT_TYPES      = {"datetime", "datetime2", "smalldatetime", "date", "time"}  # (exclude datetimeoffset for now)
CHAR_TYPES    = {"varchar", "nvarchar", "char", "nchar", "text", "ntext"}
BIT_TYPES     = {"bit"}

# -----------------------------
# 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"))
    p.add_argument("--etl-auditname", default=os.getenv("ETLauditname"))
    return p.parse_args()

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 "")
    )

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, allowed_cols: set[str]) -> pd.DataFrame:
    """Add ETL columns only if the target table actually has them; fill when entire column null."""
    meta = dict(meta or {})
    meta["ETLdatetime"] = normalize_etl_datetime(meta.get("ETLdatetime"))
    for col in ("PROCESSlogid", "PACKAGElogid", "DATAlogid", "ETLdatetime", "ETLauditname"):
        if col not in allowed_cols:
            continue
        val = meta.get(col)
        if col not in df.columns:
            df[col] = val
        else:
            if df[col].isna().all():
                df[col] = val
    return df

# -----------------------------
# OAuth 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'}
        resp = requests.post(auth_server_url, data=token_req, headers=headers)
        resp.raise_for_status()
        return resp.json().get('access_token')
    except Exception:
        print(traceback.format_exc()); sys.exit(1)

# -----------------------------
# Schema helpers
# -----------------------------
def _get_sql_schema(cursor, schema, table_name):
    cols = []
    for row in cursor.columns(table=table_name, schema=schema):
        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),
            "nullable": bool(getattr(row, "nullable", True)),
        })
    return cols

def _get_identity_columns(cursor, schema, table_name) -> set[str]:
    q = """
    SELECT c.name
    FROM sys.columns c
    JOIN sys.tables t   ON c.object_id=t.object_id
    JOIN sys.schemas s  ON s.schema_id=t.schema_id
    WHERE s.name=? AND t.name=? AND c.is_identity=1;
    """
    ids = set()
    for r in cursor.execute(q, (schema, table_name)):
        ids.add(r[0])
    return ids

# -----------------------------
# Coercers
# -----------------------------
def _coerce_datetimes(df: pd.DataFrame, sql_cols: list[dict], convert_to_nz: bool) -> pd.DataFrame:
    out = df.copy()
    for meta in sql_cols:
        col, typ = meta["name"], meta["type"]
        if col not in out.columns or typ not in DT_TYPES:
            continue
        try:
            if typ == "date":
                s = pd.to_datetime(out[col], errors="coerce", utc=True, format="%Y-%m-%d")
            elif typ == "time":
                # parse as time (date will be epoch); still OK when we drop tz
                s = pd.to_datetime(out[col], errors="coerce", utc=True, format="%H:%M:%S")
            else:
                # datetime / smalldatetime / datetime2
                s = pd.to_datetime(out[col], errors="coerce", utc=True)
            if pd.api.types.is_datetime64_any_dtype(s):
                if convert_to_nz:
                    s = s.dt.tz_convert('Pacific/Auckland')
                s = s.dt.tz_localize(None)  # SQL datetime2 likes naive
            out[col] = s
        except Exception:
            # fall back to best-effort parse
            s = pd.to_datetime(out[col], errors="coerce", utc=True)
            if convert_to_nz and pd.api.types.is_datetime64_any_dtype(s):
                s = s.dt.tz_convert('Pacific/Auckland')
            if pd.api.types.is_datetime64_any_dtype(s):
                s = s.dt.tz_localize(None)
            out[col] = s
    return out

def _clamp_and_round_decimals(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    for meta in sql_cols:
        if meta["type"] not in DEC_TYPES:
            continue
        col, prec, scale = meta["name"], meta["size"], meta["scale"]
        if col not in out.columns:
            continue
        max_abs = (Decimal(10) ** (prec - scale)) - (Decimal(1) / (Decimal(10) ** scale))
        q = Decimal("1." + ("0" * scale)) if scale > 0 else Decimal("1")
        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 scale > 0 else Decimal(str(v)).to_integral_value(rounding=ROUND_HALF_UP)
                return None if d.copy_abs() > max_abs else d
            except (InvalidOperation, ValueError): return None
        out[col] = out[col].map(fix).astype(object)
    return out

def _sanitize_floats(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    for meta in sql_cols:
        if meta["type"] not in FLOAT_TYPES: continue
        col = meta["name"]
        if col not in out.columns: continue
        s = pd.to_numeric(out[col], errors="coerce")
        s = s.where(np.isfinite(s), np.nan)
        out[col] = s.astype(object).where(~pd.isna(s), None)
    return out

def _coerce_bits(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    truthy = {"1","true","t","y","yes","on"}
    falsy  = {"0","false","f","n","no","off"}
    for meta in sql_cols:
        if meta["type"] not in BIT_TYPES: continue
        col = meta["name"]
        if col not in out.columns: continue
        def to_bit(v):
            if v is None or (isinstance(v, float) and pd.isna(v)): return None
            if isinstance(v, (int, bool)): return int(bool(v))
            s = str(v).strip().lower()
            if s in truthy: return 1
            if s in falsy: return 0
            try:
                return 1 if float(s) != 0.0 else 0
            except Exception:
                return None
        out[col] = out[col].map(to_bit)
    return out

def _coerce_ints(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    ranges = {
        "tinyint": (0, 255),
        "smallint": (-32768, 32767),
        "int": (-2147483648, 2147483647),
        "bigint": (-9223372036854775808, 9223372036854775807),
    }
    for meta in sql_cols:
        typ, col = meta["type"], meta["name"]
        if typ not in INT_TYPES or col not in out.columns:
            continue
        s = pd.to_numeric(out[col], errors="coerce")
        lo, hi = ranges["bigint" if typ not in ranges else typ]
        def to_int(v):
            if pd.isna(v): return None
            try:
                iv = int(float(v))  # tolerate "123.0"
                return iv if lo <= iv <= hi else None
            except Exception:
                return None
        out[col] = s.apply(to_int)
    return out

def _truncate_strings(df: pd.DataFrame, sql_cols: list[dict]) -> pd.DataFrame:
    out = df.copy()
    for meta in sql_cols:
        if meta["type"] not in CHAR_TYPES: continue
        col, size = meta["name"], meta["size"]
        if col not in out.columns or size <= 0: continue
        out[col] = out[col].map(lambda v: None if v is None or (isinstance(v, float) and pd.isna(v))
                                else (str(v) if len(str(v)) <= size else str(v)[:size]))
    return out

# -----------------------------
# 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:
        headers = {'Authorization': f'Bearer {token}'}
        page = 1
        try:
            conn = pyodbc.connect(conn_string)
            cursor = conn.cursor()
        except Exception:
            print(traceback.format_exc()); sys.exit(1)

        try:
            # Prepare schema details once
            sql_cols   = _get_sql_schema(cursor, schema, table_name)
            sql_names  = [c["name"] for c in sql_cols]
            sql_types  = {c["name"]: c["type"] for c in sql_cols}
            sql_name_set = set(sql_names)
            identity_cols = _get_identity_columns(cursor, schema, table_name)

            cursor.execute(f'TRUNCATE TABLE [{schema}].[{table_name}]')

            while True:
                params = {'pageSize': 10000, 'page': page}
                try:
                    resp = requests.get(api_url, headers=headers, params=params)
                except Exception:
                    print(traceback.format_exc()); sys.exit(1)

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

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

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

                # Drop identity cols if they appear in payload
                drop_id = [c for c in identity_cols if c in temp_df.columns]
                if drop_id:
                    temp_df.drop(columns=drop_id, inplace=True)

                # Add ETL columns only if target table has them
                temp_df = add_etl_columns(temp_df, meta, allowed_cols=sql_name_set)

                # Coerce by SQL types
                #   1) Datetimes (by SQL type, not by hard-coded names)
                temp_df = _coerce_datetimes(temp_df, sql_cols, convert_to_nz)

                #   2) Decimal / Numeric precision & scale
                temp_df = _clamp_and_round_decimals(temp_df, sql_cols)

                #   3) Float/Real -> clean Python float/None
                temp_df = _sanitize_floats(temp_df, sql_cols)

                #   4) Bit
                temp_df = _coerce_bits(temp_df, sql_cols)

                #   5) Int family (range-checked)
                temp_df = _coerce_ints(temp_df, sql_cols)

                # Strings: normalize empty to NULL and truncate to column width
                temp_df.replace({"": pd.NA, " ": pd.NA}, inplace=True)
                temp_df = temp_df.where(pd.notnull(temp_df), None)
                temp_df = _truncate_strings(temp_df, sql_cols)

                # Keep only columns that exist in the target table (ignore extras from API)
                keep_cols = [c for c in sql_names if c in temp_df.columns and c not in identity_cols]
                if not keep_cols:
                    logging.info("No matching columns to insert for this page; skipping.")
                    page += 1
                    continue
                temp_df = temp_df[keep_cols]

                # Build INSERT
                column_names = ','.join(f'[{c}]' for c in keep_cols)
                placeholders = ','.join('?' for _ in keep_cols)

                # Optional: log param map once per page
                for i, c in enumerate(keep_cols, start=1):
                    logging.info(f"COLMAP {i:03d} -> {c} ({sql_types.get(c)})")

                # Insert with row-probe fallback
                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(keep_cols, row), start=1):
                                logging.error(f"  #{i:03d} {c}: {repr(v)} ({type(v).__name__})")
                            return  # keep logs
                    raise

                conn.commit()
                page += 1

            logging.info('Closing connection to db')
            conn.close()
        except Exception:
            print(traceback.format_exc())
    except Exception:
        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):
    # convert UTC -> Pacific/Auckland for datetime* targets
    _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 UTC (no conversion); still stores naive datetime for 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 (only parses args & logging)
# -----------------------------
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: wire up in your own launcher (left commented intentionally)
    # token = get_token(AUTH_URL, CLIENT_ID, CLIENT_SECRET)
    # migrate_utc(API_URL, "dbo", "TargetTable", token, CONN_STR, meta=etl_meta)
