"""DuckDB table management and incremental sync utilities.

This module provides utilities for:
- Automatic table creation with schema inference
- Composite index creation for (trade_date, ts_code)
- Incremental sync strategies (by date or by stock)
- Table statistics and metadata
"""

import pandas as pd
from typing import Optional, List, Dict, Any, Callable, Tuple, Literal
from datetime import datetime, timedelta
from collections import defaultdict
from src.data.storage import Storage, ThreadSafeStorage, DEFAULT_TYPE_MAPPING


class TableManager:
    """Manages DuckDB table creation and schema."""

    def __init__(self, storage: Optional[Storage] = None):
        """Initialize table manager.

        Args:
            storage: Storage instance (creates new if None)
        """
        self.storage = storage or Storage()

    def create_table_from_dataframe(
        self,
        table_name: str,
        data: pd.DataFrame,
        primary_keys: Optional[List[str]] = None,
        create_index: bool = True,
    ) -> bool:
        """Create table from DataFrame schema with automatic type inference.

        Automatically creates composite index on (trade_date, ts_code) if both exist.

        Args:
            table_name: Name of the table to create
            data: DataFrame to infer schema from
            primary_keys: List of columns for primary key (default: auto-detect)
            create_index: Whether to create composite index

        Returns:
            True if table created successfully
        """
        if data.empty:
            print(
                f"[TableManager] Cannot create table {table_name} from empty DataFrame"
            )
            return False

        try:
            # Build column definitions
            columns = []
            for col in data.columns:
                if col in DEFAULT_TYPE_MAPPING:
                    col_type = DEFAULT_TYPE_MAPPING[col]
                else:
                    # Infer type from pandas dtype
                    dtype = str(data[col].dtype)
                    if "int" in dtype:
                        col_type = "INTEGER"
                    elif "float" in dtype:
                        col_type = "DOUBLE"
                    elif "bool" in dtype:
                        col_type = "BOOLEAN"
                    elif "datetime" in dtype:
                        col_type = "TIMESTAMP"
                    else:
                        col_type = "VARCHAR"
                columns.append(f'"{col}" {col_type}')

            # Determine primary key
            pk_constraint = ""
            if primary_keys:
                pk_cols = ", ".join([f'"{k}"' for k in primary_keys])
                pk_constraint = f", PRIMARY KEY ({pk_cols})"
            elif "ts_code" in data.columns and "trade_date" in data.columns:
                pk_constraint = ', PRIMARY KEY ("ts_code", "trade_date")'

            # Create table
            columns_sql = ", ".join(columns)
            create_sql = f'CREATE TABLE IF NOT EXISTS "{table_name}" ({columns_sql}{pk_constraint})'

            self.storage._connection.execute(create_sql)
            print(
                f"[TableManager] Created table '{table_name}' with {len(data.columns)} columns"
            )

            # Create composite index if requested and columns exist
            if (
                create_index
                and "trade_date" in data.columns
                and "ts_code" in data.columns
            ):
                index_name = f"idx_{table_name}_date_code"
                self.storage._connection.execute(f"""
                    CREATE INDEX IF NOT EXISTS "{index_name}" ON "{table_name}"("trade_date", "ts_code")
                """)
                print(
                    f"[TableManager] Created composite index on '{table_name}'(trade_date, ts_code)"
                )

            return True

        except Exception as e:
            print(f"[TableManager] Error creating table {table_name}: {e}")
            return False

    def ensure_table_exists(
        self,
        table_name: str,
        sample_data: Optional[pd.DataFrame] = None,
    ) -> bool:
        """Ensure table exists, create if it doesn't.

        Args:
            table_name: Name of the table
            sample_data: Sample DataFrame to infer schema (required if table doesn't exist)

        Returns:
            True if table exists or was created successfully
        """
        if self.storage.exists(table_name):
            return True

        if sample_data is None or sample_data.empty:
            print(
                f"[TableManager] Table '{table_name}' doesn't exist and no sample data provided"
            )
            return False

        return self.create_table_from_dataframe(table_name, sample_data)


class IncrementalSync:
    """Handles incremental synchronization strategies."""

    # Sync strategy types
    SYNC_BY_DATE = "by_date"  # Sync all stocks for date range
    SYNC_BY_STOCK = "by_stock"  # Sync specific stocks for full date range

    def __init__(self, storage: Optional[Storage] = None):
        """Initialize incremental sync manager.

        Args:
            storage: Storage instance (creates new if None)
        """
        self.storage = storage or Storage()
        self.table_manager = TableManager(self.storage)

    def get_sync_strategy(
        self,
        table_name: str,
        start_date: str,
        end_date: str,
        stock_codes: Optional[List[str]] = None,
    ) -> Tuple[str, Optional[str], Optional[str], Optional[List[str]]]:
        """Determine the best sync strategy based on existing data.

        Logic:
        1. If table doesn't exist: full sync by date
        2. If table exists and has data:
           - If stock_codes provided: sync by stock (update specific stocks)
           - Otherwise: sync by date from last_date + 1

        Args:
            table_name: Name of the table to sync
            start_date: Requested start date (YYYYMMDD)
            end_date: Requested end date (YYYYMMDD)
            stock_codes: Optional list of specific stocks to sync

        Returns:
            Tuple of (strategy, sync_start, sync_end, stocks_to_sync)
            - strategy: 'by_date' or 'by_stock' or 'none'
            - sync_start: Start date for sync (None if no sync needed)
            - sync_end: End date for sync (None if no sync needed)
            - stocks_to_sync: List of stocks to sync (None for all)
        """
        # Check if table exists
        if not self.storage.exists(table_name):
            print(
                f"[IncrementalSync] Table '{table_name}' doesn't exist, will create and do full sync"
            )
            return (self.SYNC_BY_DATE, start_date, end_date, None)

        # Get table stats
        stats = self.get_table_stats(table_name)

        if stats["row_count"] == 0:
            print(f"[IncrementalSync] Table '{table_name}' is empty, doing full sync")
            return (self.SYNC_BY_DATE, start_date, end_date, None)

        # If specific stocks requested, sync by stock
        if stock_codes:
            existing_stocks = set(self.storage.get_distinct_stocks(table_name))
            requested_stocks = set(stock_codes)
            missing_stocks = requested_stocks - existing_stocks

            if not missing_stocks:
                print(
                    f"[IncrementalSync] All requested stocks already exist in '{table_name}'"
                )
                return ("none", None, None, None)

            print(
                f"[IncrementalSync] Syncing {len(missing_stocks)} missing stocks by stock strategy"
            )
            return (self.SYNC_BY_STOCK, start_date, end_date, list(missing_stocks))

        # Check if we need date-based sync
        table_last_date = stats.get("max_date")

        if table_last_date is None:
            return (self.SYNC_BY_DATE, start_date, end_date, None)

        # Compare dates
        table_last = int(table_last_date)
        requested_end = int(end_date)

        if table_last >= requested_end:
            print(
                f"[IncrementalSync] Table '{table_name}' is up-to-date (last: {table_last_date})"
            )
            return ("none", None, None, None)

        # Incremental sync from next day after last_date
        next_date = self._get_next_date(table_last_date)
        print(f"[IncrementalSync] Incremental sync needed: {next_date} to {end_date}")
        return (self.SYNC_BY_DATE, next_date, end_date, None)

    def get_table_stats(self, table_name: str) -> Dict[str, Any]:
        """Get statistics about a table.

        Returns:
            Dict with exists, row_count, min_date, max_date, unique_stocks
        """
        stats = {
            "exists": False,
            "row_count": 0,
            "min_date": None,
            "max_date": None,
            "unique_stocks": 0,
        }

        if not self.storage.exists(table_name):
            return stats

        try:
            conn = self.storage._connection

            # Row count
            row_count = conn.execute(f'SELECT COUNT(*) FROM "{table_name}"').fetchone()[
                0
            ]
            stats["row_count"] = row_count
            stats["exists"] = True

            # Get column names
            columns_result = conn.execute(
                """
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = ?
            """,
                [table_name],
            ).fetchall()
            columns = [row[0] for row in columns_result]

            # Date range
            if "trade_date" in columns:
                date_result = conn.execute(f'''
                    SELECT MIN("trade_date"), MAX("trade_date") FROM "{table_name}"
                ''').fetchone()
                if date_result[0]:
                    stats["min_date"] = (
                        date_result[0].strftime("%Y%m%d")
                        if hasattr(date_result[0], "strftime")
                        else str(date_result[0])
                    )
                if date_result[1]:
                    stats["max_date"] = (
                        date_result[1].strftime("%Y%m%d")
                        if hasattr(date_result[1], "strftime")
                        else str(date_result[1])
                    )

            # Unique stocks
            if "ts_code" in columns:
                unique_count = conn.execute(f'''
                    SELECT COUNT(DISTINCT "ts_code") FROM "{table_name}"
                ''').fetchone()[0]
                stats["unique_stocks"] = unique_count

        except Exception as e:
            print(f"[IncrementalSync] Error getting stats for {table_name}: {e}")

        return stats

    def sync_data(
        self,
        table_name: str,
        data: pd.DataFrame,
        strategy: Literal["by_date", "by_stock", "replace"] = "by_date",
    ) -> Dict[str, Any]:
        """Sync data to table using specified strategy.

        Args:
            table_name: Target table name
            data: DataFrame to sync
            strategy: Sync strategy
                - 'by_date': UPSERT based on primary key (ts_code, trade_date)
                - 'by_stock': Replace data for specific stocks
                - 'replace': Full replace of table

        Returns:
            Dict with status, rows_inserted, rows_updated
        """
        if data.empty:
            return {"status": "skipped", "rows_inserted": 0, "rows_updated": 0}

        # Ensure table exists
        if not self.table_manager.ensure_table_exists(table_name, data):
            return {"status": "error", "error": "Failed to create table"}

        try:
            if strategy == "replace":
                # Full replace
                result = self.storage.save(table_name, data, mode="replace")
                return {
                    "status": result["status"],
                    "rows_inserted": result.get("rows", 0),
                    "rows_updated": 0,
                }

            elif strategy == "by_stock":
                # Delete existing data for these stocks, then insert
                if "ts_code" in data.columns:
                    stocks = data["ts_code"].unique().tolist()
                    placeholders = ", ".join(["?"] * len(stocks))
                    self.storage._connection.execute(
                        f'''
                        DELETE FROM "{table_name}" WHERE "ts_code" IN ({placeholders})
                    ''',
                        stocks,
                    )
                    print(
                        f"[IncrementalSync] Deleted existing data for {len(stocks)} stocks"
                    )

                result = self.storage.save(table_name, data, mode="append")
                return {
                    "status": result["status"],
                    "rows_inserted": result.get("rows", 0),
                    "rows_updated": 0,
                }

            else:  # by_date (default)
                # UPSERT using INSERT OR REPLACE
                result = self.storage.save(table_name, data, mode="append")
                return {
                    "status": result["status"],
                    "rows_inserted": result.get("rows", 0),
                    "rows_updated": 0,  # DuckDB doesn't distinguish in UPSERT
                }

        except Exception as e:
            print(f"[IncrementalSync] Error syncing data to {table_name}: {e}")
            return {"status": "error", "error": str(e)}

    def _get_next_date(self, date_str: str) -> str:
        """Get the next day after given date.

        Args:
            date_str: Date in YYYYMMDD format

        Returns:
            Next date in YYYYMMDD format
        """
        dt = datetime.strptime(date_str, "%Y%m%d")
        next_dt = dt + timedelta(days=1)
        return next_dt.strftime("%Y%m%d")


class SyncManager:
    """High-level sync manager that coordinates table creation and incremental updates."""

    def __init__(self, storage: Optional[Storage] = None):
        """Initialize sync manager.

        Args:
            storage: Storage instance (creates new if None)
        """
        self.storage = storage or Storage()
        self.table_manager = TableManager(self.storage)
        self.incremental_sync = IncrementalSync(self.storage)

    def sync(
        self,
        table_name: str,
        fetch_func: Callable[..., pd.DataFrame],
        start_date: str,
        end_date: str,
        stock_codes: Optional[List[str]] = None,
        **fetch_kwargs,
    ) -> Dict[str, Any]:
        """Main sync method - handles full logic of table creation and incremental sync.

        This is the recommended way to sync data:
        1. Checks if table exists, creates if not
        2. Determines best sync strategy
        3. Fetches data using provided function
        4. Applies incremental update

        Args:
            table_name: Target table name
            fetch_func: Function to fetch data (should return DataFrame)
            start_date: Start date for sync (YYYYMMDD)
            end_date: End date for sync (YYYYMMDD)
            stock_codes: Optional list of stocks to sync (None = all)
            **fetch_kwargs: Additional arguments to pass to fetch_func

        Returns:
            Dict with sync results
        """
        print(f"\n[SyncManager] Starting sync for table '{table_name}'")
        print(f"[SyncManager] Date range: {start_date} to {end_date}")

        # Determine sync strategy
        strategy, sync_start, sync_end, stocks_to_sync = (
            self.incremental_sync.get_sync_strategy(
                table_name, start_date, end_date, stock_codes
            )
        )

        if strategy == "none":
            print(f"[SyncManager] No sync needed for '{table_name}'")
            return {
                "status": "skipped",
                "table": table_name,
                "reason": "up-to-date",
            }

        # Fetch data
        print(f"[SyncManager] Fetching data with strategy '{strategy}'...")
        try:
            if stocks_to_sync:
                # Fetch specific stocks
                data_list = []
                for ts_code in stocks_to_sync:
                    df = fetch_func(
                        ts_code=ts_code,
                        start_date=sync_start,
                        end_date=sync_end,
                        **fetch_kwargs,
                    )
                    if not df.empty:
                        data_list.append(df)

                if data_list:
                    data = pd.concat(data_list, ignore_index=True)
                else:
                    data = pd.DataFrame()
            else:
                # Fetch all data at once
                data = fetch_func(
                    start_date=sync_start, end_date=sync_end, **fetch_kwargs
                )
        except Exception as e:
            print(f"[SyncManager] Error fetching data: {e}")
            return {
                "status": "error",
                "table": table_name,
                "error": str(e),
            }

        if data.empty:
            print(f"[SyncManager] No data fetched")
            return {
                "status": "no_data",
                "table": table_name,
            }

        print(f"[SyncManager] Fetched {len(data)} rows")

        # Ensure table exists
        if not self.table_manager.ensure_table_exists(table_name, data):
            return {
                "status": "error",
                "table": table_name,
                "error": "Failed to create table",
            }

        # Apply sync
        result = self.incremental_sync.sync_data(table_name, data, strategy)

        print(f"[SyncManager] Sync complete: {result}")
        return {
            "status": result["status"],
            "table": table_name,
            "strategy": strategy,
            "rows": result.get("rows_inserted", 0),
            "date_range": f"{sync_start} to {sync_end}"
            if sync_start and sync_end
            else None,
        }


# Convenience functions


def ensure_table(
    table_name: str,
    sample_data: pd.DataFrame,
    storage: Optional[Storage] = None,
) -> bool:
    """Ensure a table exists, creating it if necessary.

    Args:
        table_name: Name of the table
        sample_data: Sample DataFrame to define schema
        storage: Optional Storage instance

    Returns:
        True if table exists or was created
    """
    manager = TableManager(storage)
    return manager.ensure_table_exists(table_name, sample_data)


def get_table_info(
    table_name: str, storage: Optional[Storage] = None
) -> Dict[str, Any]:
    """Get information about a table.

    Args:
        table_name: Name of the table
        storage: Optional Storage instance

    Returns:
        Dict with table statistics
    """
    sync = IncrementalSync(storage)
    return sync.get_table_stats(table_name)


def sync_table(
    table_name: str,
    fetch_func: Callable[..., pd.DataFrame],
    start_date: str,
    end_date: str,
    stock_codes: Optional[List[str]] = None,
    storage: Optional[Storage] = None,
    **fetch_kwargs,
) -> Dict[str, Any]:
    """Sync data to a table with automatic table creation and incremental updates.

    This is the main entry point for syncing data to DuckDB.

    Args:
        table_name: Target table name
        fetch_func: Function to fetch data
        start_date: Start date (YYYYMMDD)
        end_date: End date (YYYYMMDD)
        stock_codes: Optional list of specific stocks
        storage: Optional Storage instance
        **fetch_kwargs: Additional arguments for fetch_func

    Returns:
        Dict with sync results

    Example:
        >>> from src.data.api_wrappers import get_daily
        >>> result = sync_table(
        ...     "daily",
        ...     get_daily,
        ...     "20240101",
        ...     "20240131",
        ...     stock_codes=["000001.SZ", "600000.SH"]
        ... )
    """
    manager = SyncManager(storage)
    return manager.sync(
        table_name=table_name,
        fetch_func=fetch_func,
        start_date=start_date,
        end_date=end_date,
        stock_codes=stock_codes,
        **fetch_kwargs,
    )
