from itertools import product
import json
import uuid
from datetime import date, datetime, timedelta
from decimal import Decimal
from typing import Optional, Any
from bisect import bisect_right
import pandas as pd
import pandas_market_calendars as mcal
from more_itertools import windowed
from sqlalchemy import (
    Connection,
    create_engine,
    MetaData,
    Table,
    select,
    insert,
    update,
    text,
    and_,
    or_,
)
from sqlalchemy.orm import Session
from sqlalchemy.engine import Engine
from sqlalchemy import select, and_, or_, text
from sqlalchemy.engine import Connection
from uuid import uuid4

from io import BytesIO
import base64
import requests
import copy

import re
import numpy as np
from scipy.stats import norm
from enum import Enum
import bisect

from dotenv import load_dotenv
import os
import math

import datetime as datetime
from datetime import datetime as dt

import pandas as pd

from dotenv import load_dotenv
from io import BytesIO



import pandas as pd


import os



from datetime import datetime, date


load_dotenv()
server = os.getenv("SQL_SERVER")
user = os.getenv("SQL_USER")
password = os.getenv("SQL_PASSWORD")
database = os.getenv("SQL_DATABASE")

class SQLHandler:
    def __init__(self, server: str, user: str, password: str, database: str):
        if not all([server, user, password, database]):
            raise ValueError("Database credentials not fully provided.")

        connection_string = (
            f"mssql+pymssql://{user}:{password}@{server}/{database}"
        )

        self.engine: Engine = create_engine(connection_string, future=True)
        self.metadata = MetaData()
        self.metadata.reflect(bind=self.engine)
    
    def upsert_bulk(self, data: list[dict], table, id_column: str) -> list[str]:
        if not data:
            return []

        valid_columns = {c.name for c in table.columns}
        # Columns that DB auto-fills, so omit them from INSERT unless explicitly provided.
        default_constraint_cols = {"new_createdon", "new_statecode"}
        skip_update_cols = {"new_createdby"}
        cleaned_data = []
        
        # Track which columns were actually provided in the input data to preserve selective updates.
        provided_columns = set()

        # Normalize and validate
        for record in data:
            record = record.copy()

            # Map entity_id field to the target id_column.
            if "entity_id" in record:
                record[id_column] = record.pop("entity_id")

            # Extract only valid columns from the record.
            row = {k: v for k, v in record.items() if k in valid_columns}
            
            # Track which columns were provided in this record.
            provided_columns.update(row.keys())
            
            # Generate GUID if ID is missing or empty.
            if id_column not in row or not row[id_column]:
                row[id_column] = str(uuid.uuid4())

            # Fill missing columns with None (except those DB manages with defaults).
            for col in valid_columns:
                if col not in row:
                    if col not in default_constraint_cols:
                        row[col] = None

            cleaned_data.append(row)

        # Extract IDs for return value.
        ids = [str(row[id_column]) for row in cleaned_data]
        # Include only columns that were provided or aren't DB-managed defaults.
        cols = [c for c in valid_columns if c not in default_constraint_cols or c in provided_columns]

        # Format column list for SQL statements.
        insert_cols = ", ".join(f"[{c}]" for c in cols)

        with Session(self.engine) as session:
            connection = session.connection()
            # Create temp table with same schema as target to stage bulk operations.
            temp_table = f"#TempUpsert_{uuid.uuid4().hex[:8]}"
            connection.execute(
                text(f"SELECT TOP 0 * INTO {temp_table} FROM [dbo].[{table.name}]")
            )

            # Bulk insert all records into temp table for efficient processing.
            # insert_stmt = text(f"""
            #     INSERT INTO {temp_table} ({insert_cols})
            #     VALUES ({", ".join(f":{c}" for c in cols)})
            # """)
            # connection.execute(insert_stmt, cleaned_data)
            self._bulk_insert_batched(
                connection=connection,
                table_name=temp_table,
                columns=cols,
                rows=cleaned_data,
            ) 

            # Filter to only columns that user provided and are updatable.
            update_columns = [
                c for c in provided_columns
                if c not in skip_update_cols and c != id_column and c in valid_columns
            ]
            # Update Bulk existing records with provided columns only (selective update pattern).
            if update_columns:
                set_clause = ", ".join(f"T.[{c}] = S.[{c}]" for c in update_columns)
                update_stmt = f"""
                    UPDATE T
                    SET {set_clause}
                    FROM [dbo].[{table.name}] AS T
                    INNER JOIN {temp_table} AS S
                    ON T.[{id_column}] = S.[{id_column}];
                """
                connection.execute(text(update_stmt))

            # Insert Bulk records that don't already exist.
            insert_stmt2 = f"""
                INSERT INTO [dbo].[{table.name}] ({insert_cols})
                SELECT {insert_cols}
                FROM {temp_table} AS S
                WHERE NOT EXISTS (
                    SELECT 1 FROM [dbo].[{table.name}] AS T
                    WHERE T.[{id_column}] = S.[{id_column}]
                );
            """
            connection.execute(text(insert_stmt2))
            # Cleanup temp table.
            connection.execute(text(f"DROP TABLE {temp_table}"))
            session.commit()

        return ids

    def get_entities(self, table, select_columns: str | None = None, filters: dict[str, object] | list[dict[str, object]] | None = None):
        with self.engine.begin() as conn:
            # Columns
            if select_columns:
                cols = []
                for name in (c.strip() for c in select_columns.split(",")):
                    if not name:
                        continue
                    if name not in table.c:
                        raise ValueError(f"Column '{name}' not found in table '{table.name}'")
                    cols.append(table.c[name])
            else:
                cols = list(table.c)

            # No filters → simple select
            if not filters:
                stmt = select(*cols)
                rows = conn.execute(stmt).mappings().all()
                return [{k: self.serialize_value(v) for k, v in row.items()} for row in rows]

            # Normalize to list
            if isinstance(filters, dict):
                filters = [filters]

            return self._execute_with_temp_table(conn, table, cols, filters)

    def _execute_with_temp_table(self, conn, table, cols, filters):

        temp_name = f"#temp_filter_{uuid.uuid4().hex}"

        filter_columns = set()
        for f in filters:
            filter_columns.update(f.keys())
        filter_columns = list(filter_columns)

        # Validate columns
        for col in filter_columns:
            if col not in table.c:
                raise ValueError(f"Filter column '{col}' not found in table '{table.name}'")

        # normalize filters into list of dicts with all combinations for multi-valued entries
        expanded_rows = []
        for group in filters:
            keys = []
            values = []
            for col in filter_columns:
                v = group.get(col, None)

                if isinstance(v, (list, tuple, set)):
                    vals = list(v)
                    if not vals:
                        vals = [None]
                    values.append(vals)
                else:
                    values.append([v])

                keys.append(col)

            for combo in product(*values):
                expanded_rows.append(dict(zip(keys, combo)))

        if not expanded_rows:
            return []
        
        # ----------------------------
        # Create temp table
        # ----------------------------
        column_defs = []
        for col_name in filter_columns:
            col_type = table.c[col_name].type.compile(self.engine.dialect)
            column_defs.append(f"[{col_name}] {col_type}")

        create_sql = f"""
            CREATE TABLE {temp_name} (
                {", ".join(column_defs)}
            )
        """
        conn.execute(text(create_sql))

        # ----------------------------
        # Bulk insert in temp table
        # ----------------------------
        # insert_sql = text(f"""
        #     INSERT INTO {temp_name} ({", ".join(f"[{c}]" for c in filter_columns)})
        #     VALUES ({", ".join(f":{c}" for c in filter_columns)})
        # """)
        # conn.execute(insert_sql, expanded_rows)
        self._bulk_insert_batched(
            connection=conn,
            table_name=temp_name,
            columns=filter_columns,
            rows=expanded_rows,
        ) 

        # ----------------------------
        # NULL-SAFE JOIN CONDITIONS 
        # ----------------------------
        join_conditions = " AND ".join(
            f"(t.[{c}] = f.[{c}] OR (t.[{c}] IS NULL AND f.[{c}] IS NULL))"
            for c in filter_columns
        )

        select_list = ", ".join(f"t.[{c.name}]" for c in cols)
        schema = table.schema or "dbo"

        query_sql = f"""
            SELECT {select_list}
            FROM [{schema}].[{table.name}] t
            INNER JOIN {temp_name} f
                ON {join_conditions}
        """

        rows = conn.execute(text(query_sql)).mappings().all()

        conn.execute(text(f"DROP TABLE {temp_name}"))

        return [{k: self.serialize_value(v) for k, v in row.items()} for row in rows]
    

    def _bulk_insert_batched(self, connection, table_name: str, columns: list[str], rows: list[dict],
        MAX_PARAMS: int = 2000): # stay below SQL Server 2100 limit
        """
        High-performance batched multi-row INSERT.
        """

        if not rows:
            return

        col_count = len(columns)

        # Calculate safe batch size dynamically
        batch_size = max(1, MAX_PARAMS // col_count)

        insert_cols = ", ".join(f"[{c}]" for c in columns)

        for i in range(0, len(rows), batch_size):
            batch = rows[i:i + batch_size]

            values_sql = []
            params = {}

            for row_index, row in enumerate(batch):
                row_params = []
                for col in columns:
                    param_name = f"{col}_{row_index}"
                    row_params.append(f":{param_name}")
                    params[param_name] = row.get(col)
                values_sql.append(f"({', '.join(row_params)})")

            stmt = text(f"""
                INSERT INTO {table_name} ({insert_cols})
                VALUES {', '.join(values_sql)}
            """)

            connection.execute(stmt, params)
    @staticmethod
    def serialize_value(value):
        if isinstance(value, (datetime, date)):
            return value.isoformat()
        elif isinstance(value, Decimal):
            return float(value)
        elif isinstance(value, uuid.UUID):
            return str(value)
        elif isinstance(value, bytes):
            return value.hex() 
        elif isinstance(value, bool):
            return bool(value) 
        elif value is None:
            return None
        elif isinstance(value, (list, dict)):
            return json.dumps(value)
        return value