"""Bulk VARIANT insert via staged JSON for Snowflake.

Replaces row-by-row PARSE_JSON(%s) inserts with PUT → COPY INTO,
reducing 55+ minute metadata loads to under 1 minute.
"""

from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional


class SnowflakeBulkLoader:
    """Load VARIANT data into Snowflake via staged NDJSON files.

    Instead of row-by-row ``INSERT ... PARSE_JSON(%s)`` calls, this class:
    1. Writes records as newline-delimited JSON to a temp file
    2. ``PUT`` the file to ``@~/databridge_stage``
    3. ``COPY INTO`` with ``FILE_FORMAT=(TYPE='JSON')``
    4. Removes the staged file
    """

    STAGE_NAME = "@~/databridge_bulk_stage"

    def bulk_insert_variant(
        self,
        cursor: Any,
        table: str,
        id_col: str,
        variant_col: str,
        records: List[Dict[str, Any]],
        extra_cols: Optional[Dict[str, str]] = None,
    ) -> int:
        """Insert records as VARIANT rows via staged JSON.

        Args:
            cursor: Active Snowflake cursor.
            table: Fully qualified target table name.
            id_col: Name of the ID / key column.
            variant_col: Name of the VARIANT column.
            records: List of dicts — each becomes one row.
                     Each dict must have ``id_col`` and the remaining keys
                     become the VARIANT payload.
            extra_cols: Optional mapping of column_name → literal SQL expr
                        to include as additional columns (e.g. CURRENT_TIMESTAMP()).

        Returns:
            Number of rows loaded.
        """
        if not records:
            return 0

        # Write NDJSON temp file
        tmp = tempfile.NamedTemporaryFile(
            mode="w", suffix=".json", delete=False, encoding="utf-8"
        )
        try:
            for rec in records:
                tmp.write(json.dumps(rec, default=str) + "\n")
            tmp.close()

            filename = Path(tmp.name).name
            # PUT to internal stage
            cursor.execute(
                f"PUT 'file://{tmp.name}' {self.STAGE_NAME}/ AUTO_COMPRESS=TRUE OVERWRITE=TRUE"
            )

            # Build column list and SELECT expressions
            cols = [id_col, variant_col]
            select_exprs = [
                f"$1:{id_col}::STRING",
                "$1",
            ]
            if extra_cols:
                for col_name, expr in extra_cols.items():
                    cols.append(col_name)
                    select_exprs.append(expr)

            col_list = ", ".join(cols)
            select_list = ", ".join(select_exprs)

            cursor.execute(
                f"COPY INTO {table} ({col_list}) "
                f"FROM (SELECT {select_list} FROM {self.STAGE_NAME}/{filename}.gz) "
                f"FILE_FORMAT=(TYPE='JSON' STRIP_OUTER_ARRAY=FALSE) "
                f"PURGE=TRUE"
            )

            result = cursor.fetchone()
            return result[0] if result else len(records)
        finally:
            if os.path.exists(tmp.name):
                os.unlink(tmp.name)

    def bulk_merge_variant(
        self,
        cursor: Any,
        table: str,
        id_col: str,
        variant_col: str,
        records: List[Dict[str, Any]],
    ) -> int:
        """Upsert records via staged JSON with MERGE semantics.

        Args:
            cursor: Active Snowflake cursor.
            table: Fully qualified target table name.
            id_col: Name of the merge key column.
            variant_col: Name of the VARIANT column.
            records: List of dicts — each dict must contain the id_col key.

        Returns:
            Number of rows merged.
        """
        if not records:
            return 0

        tmp = tempfile.NamedTemporaryFile(
            mode="w", suffix=".json", delete=False, encoding="utf-8"
        )
        try:
            for rec in records:
                tmp.write(json.dumps(rec, default=str) + "\n")
            tmp.close()

            filename = Path(tmp.name).name
            cursor.execute(
                f"PUT 'file://{tmp.name}' {self.STAGE_NAME}/ AUTO_COMPRESS=TRUE OVERWRITE=TRUE"
            )

            cursor.execute(
                f"MERGE INTO {table} t "
                f"USING ("
                f"  SELECT $1:{id_col}::STRING AS {id_col}, $1 AS {variant_col} "
                f"  FROM {self.STAGE_NAME}/{filename}.gz "
                f"  (FILE_FORMAT => (TYPE='JSON' STRIP_OUTER_ARRAY=FALSE))"
                f") s "
                f"ON t.{id_col} = s.{id_col} "
                f"WHEN MATCHED THEN UPDATE SET {variant_col} = s.{variant_col} "
                f"WHEN NOT MATCHED THEN INSERT ({id_col}, {variant_col}) "
                f"VALUES (s.{id_col}, s.{variant_col})"
            )

            return len(records)
        finally:
            if os.path.exists(tmp.name):
                os.unlink(tmp.name)
