"""
ClickHouse storage backend for FactorDB
"""

from typing import Optional, Dict, Any, List, Tuple
from datetime import datetime, date
import pandas as pd
import numpy as np
import logging
import concurrent.futures
import multiprocessing

from ..core.config import FactorDBConfig, AssetType, FactorType
from ..core.factor import FactorMetadata
from .schema import (
    FACTOR_BASIC_SCHEMA,
    FACTOR_EVALUATE_SCHEMA,
    FACTOR_UPGRADE_SCHEMA,
    FACTOR_VALUE_SCHEMA,
    FACTOR_VALUE_HF_SCHEMA,
    get_create_table_sql,
    get_evaluation_metric_columns,
)

logger = logging.getLogger(__name__)


def _insert_worker(
    db_config: Dict[str, Any],
    db_name: str,
    table_name: str,
    data_dict: Dict[str, List[Any]],
    factor_type: FactorType,
    asset_type: AssetType
) -> int:
    """
    Worker function for parallel insertion.
    Running in a separate process, so it creates its own DB connection.
    """
    try:
        from quantchdb import ClickHouseDatabase
        
        # Reconstruct DataFrame from dict (efficient serialization)
        # Using dict is often safer/faster than pickling heavy DataFrame across processes
        df = pd.DataFrame(data_dict)
        
        if df.empty:
            return 0
            
        # Connect to database using context manager to avoid socket leak
        with ClickHouseDatabase(
            config=db_config,
            terminal_log=False,
            file_log=False
        ) as db:
            # Vectorized string construction for performance
            # This replaces the slow row-by-row iteration
            codes = df['code'].values
            dates = df['date'].values
            values = df['factor_value'].values.astype(str)

            if factor_type == FactorType.HIGH_FREQ:
                datetimes = df['datetime'].values
                # Format: ('code', 'date', 'datetime', value)
                # vectorized string concat
                # Note: We assume inputs are already strings/clean
                # We use list comprehension with zip for safety if numpy string casting is tricky with quote escaping
                # But here we can assume simple codes/dates

                # Fast vectorized approach:
                # Create a single array of formatted strings
                rows = [
                    f"('{c}', '{d}', '{dt}', {v})"
                    for c, d, dt, v in zip(codes, dates, datetimes, values)
                ]
                columns = "code, date, datetime, factor_value"
            else:
                # Format: ('code', 'date', value)
                rows = [
                    f"('{c}', '{d}', {v})"
                    for c, d, v in zip(codes, dates, values)
                ]
                columns = "code, date, factor_value"

            # Join all value tuples
            values_str = ", ".join(rows)

            sql = f"INSERT INTO {db_name}.{table_name} ({columns}) VALUES {values_str}"

            db.execute(sql)
            return len(df)
        
    except Exception as e:
        # We can't log to the main logger easily from here, print might be captured
        print(f"Error in insert worker: {e}")
        raise e

class ClickHouseStorage:
    """
    ClickHouse storage backend for factor data.

    Handles all database operations including:
    - Table creation and schema management
    - Factor registration and metadata storage
    - Factor value storage and retrieval
    - Factor evaluation metrics storage
    """

    def __init__(self, config: Optional[FactorDBConfig] = None):
        self.config = config or FactorDBConfig()
        self._db_cache: Dict[str, Any] = {}  # Cache connections by database name
        self._initialized_dbs = set()

    def close(self) -> None:
        """Close all cached database connections."""
        for db_name, db in self._db_cache.items():
            try:
                db.close()
                logger.debug(f"Closed connection to {db_name}")
            except Exception as e:
                logger.warning(f"Failed to close connection to {db_name}: {e}")
        self._db_cache.clear()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False

    def __del__(self):
        self.close()

    def _get_db(
        self,
        asset_type: Optional[AssetType] = None,
        factor_type: Optional[FactorType] = None,
        database: Optional[str] = None
    ):
        """
        Get or create database connection for specific database.

        Args:
            asset_type: Asset type (STOCK or ETF)
            factor_type: Factor frequency type (DAILY or HIGH_FREQ)
            database: Explicit database name (overrides asset_type/factor_type)

        Returns:
            ClickHouseDatabase instance
        """
        # Determine target database name
        if database is not None:
            db_name = database
        elif asset_type is not None and factor_type is not None:
            db_name = self.config.get_factor_db_name(asset_type, factor_type)
        else:
            db_name = self.config.stock_factor_db

        # Return cached connection if exists
        if db_name in self._db_cache:
            return self._db_cache[db_name]

        # Create new connection
        try:
            from quantchdb import ClickHouseDatabase
        except ImportError as e:
            raise ImportError(
                f"quantchdb is required for database operations. "
                f"Please install it first. Original error: {e}"
            )

        try:
            db = ClickHouseDatabase(
                config=self.config.get_db_config(
                    asset_type=asset_type,
                    factor_type=factor_type,
                    database=database
                ),
                terminal_log=False,
                file_log=False
            )
            self._db_cache[db_name] = db
            return db
        except Exception as e:
            raise RuntimeError(f"Failed to connect to database {db_name}: {e}")

    def _execute(
        self,
        sql: str,
        asset_type: Optional[AssetType] = None,
        factor_type: Optional[FactorType] = None,
        database: Optional[str] = None
    ) -> None:
        """Execute a SQL statement"""
        db = self._get_db(asset_type, factor_type, database)
        db.execute(sql)

    def _fetch(
        self,
        sql: str,
        asset_type: Optional[AssetType] = None,
        factor_type: Optional[FactorType] = None,
        database: Optional[str] = None
    ) -> pd.DataFrame:
        """Execute a query and return results as DataFrame"""
        db = self._get_db(asset_type, factor_type, database)
        return db.fetch(sql)

    def initialize_database(
        self,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> None:
        """
        Initialize database tables for a specific factor type.

        Creates A1_factor_basic, A2_factor_evaluate, A3_factor_upgrade tables
        if they don't exist.
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)
        db_key = f"{asset_type.value}_{factor_type.value}"

        if db_key in self._initialized_dbs:
            return

        # Create database if not exists
        self._execute(f"CREATE DATABASE IF NOT EXISTS {db_name}", asset_type, factor_type)

        # Create basic info table
        basic_sql = get_create_table_sql(FACTOR_BASIC_SCHEMA, database=db_name)
        self._execute(basic_sql, asset_type, factor_type)

        # Create evaluation table
        evaluate_sql = get_create_table_sql(FACTOR_EVALUATE_SCHEMA, database=db_name)
        self._execute(evaluate_sql, asset_type, factor_type)

        # Create upgrade tracking table
        upgrade_sql = get_create_table_sql(FACTOR_UPGRADE_SCHEMA, database=db_name)
        self._execute(upgrade_sql, asset_type, factor_type)

        self._initialized_dbs.add(db_key)
        logger.info(f"Initialized database tables for {db_name}")

    def _get_next_factor_id(
        self,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> str:
        """
        Get the next available factor ID.

        Factor ID format:
        - Stock daily: F_stk_000001
        - Stock high-freq: HF_stk_000001
        - ETF daily: F_etf_000001
        - ETF high-freq: HF_etf_000001
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)
        prefix = self.config.get_factor_id_prefix(asset_type, factor_type)

        # Query max existing ID
        sql = f"""
            SELECT factor_id
            FROM {db_name}.A1_factor_basic
            WHERE factor_id LIKE '{prefix}_%'
            ORDER BY factor_id DESC
            LIMIT 1
        """

        try:
            result = self._fetch(sql, asset_type, factor_type)
            if len(result) > 0:
                last_id = result.iloc[0]['factor_id']
                # Extract number from ID (e.g., "F_stk_000001" -> 1)
                num_str = last_id.split('_')[-1]
                next_num = int(num_str) + 1
            else:
                next_num = 1
        except Exception:
            next_num = 1

        # Format as 6-digit number
        return f"{prefix}_{next_num:06d}"

    def register_factor(
        self,
        metadata: FactorMetadata,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> str:
        """
        Register a new factor in the database.

        Args:
            metadata: Factor metadata (factor_id will be assigned if empty)
            asset_type: Asset type
            factor_type: Factor frequency type

        Returns:
            Assigned factor ID
        """
        self.initialize_database(asset_type, factor_type)
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Generate factor ID if not provided
        if not metadata.factor_id:
            metadata.factor_id = self._get_next_factor_id(asset_type, factor_type)

        # Escape string values
        def escape(val):
            if val is None:
                return "NULL"
            if isinstance(val, str):
                return f"'{val.replace(chr(39), chr(39)+chr(39))}'"
            return str(val)

        # Insert into A1_factor_basic
        sql = f"""
            INSERT INTO {db_name}.A1_factor_basic
            (factor_id, name, expression, explanation, other_info, register_datetime)
            VALUES (
                {escape(metadata.factor_id)},
                {escape(metadata.name)},
                {escape(metadata.expression)},
                {escape(metadata.explanation)},
                {escape(str(metadata.other_info) if metadata.other_info else None)},
                now()
            )
        """
        self._execute(sql, asset_type, factor_type)

        logger.info(f"Registered factor: {metadata.factor_id}")
        return metadata.factor_id

    def get_factor_metadata(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> Optional[FactorMetadata]:
        """Get factor metadata by ID"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        sql = f"""
            SELECT factor_id, name, expression, explanation, other_info, register_datetime
            FROM {db_name}.A1_factor_basic
            WHERE factor_id = '{factor_id}'
        """

        result = self._fetch(sql, asset_type, factor_type)
        if len(result) == 0:
            return None

        row = result.iloc[0]
        return FactorMetadata(
            factor_id=row['factor_id'],
            name=row['name'],
            expression=row['expression'],
            explanation=row['explanation'],
            other_info=eval(row['other_info']) if row['other_info'] else None,
            register_datetime=row['register_datetime'],
            asset_type=asset_type,
            factor_type=factor_type,
        )

    def list_factors(
        self,
        asset_type: AssetType,
        factor_type: FactorType,
        limit: int = 1000
    ) -> pd.DataFrame:
        """List all registered factors"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        sql = f"""
            SELECT factor_id, name, expression, register_datetime
            FROM {db_name}.A1_factor_basic
            ORDER BY factor_id
            LIMIT {limit}
        """

        return self._fetch(sql, asset_type, factor_type)

    def create_factor_value_table(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> None:
        """Create a table for storing factor values"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        if factor_type == FactorType.HIGH_FREQ:
            schema = FACTOR_VALUE_HF_SCHEMA
        else:
            schema = FACTOR_VALUE_SCHEMA

        sql = get_create_table_sql(schema, table_name=factor_id, database=db_name)
        self._execute(sql, asset_type, factor_type)
        logger.info(f"Created factor value table: {db_name}.{factor_id}")

    def save_factor_values(
        self,
        factor_id: str,
        values: pd.DataFrame,
        asset_type: AssetType,
        factor_type: FactorType,
        n_jobs: int = 16
    ) -> None:
        """
        Save factor values to database.

        Args:
            factor_id: Factor ID
            values: DataFrame with columns:
                - Daily: code, date, factor_value
                - High-freq: code, date, datetime, factor_value
            asset_type: Asset type
            factor_type: Factor frequency type
            n_jobs: Number of parallel workers (default 16)
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Ensure table exists
        self.create_factor_value_table(factor_id, asset_type, factor_type)

        # Prepare data
        df = values.copy()
        df['factor_value'] = df['factor_value'].astype(np.float32)

        # Convert date to string format
        if 'date' in df.columns:
            if not pd.api.types.is_string_dtype(df['date']):
                try:
                    # Optimized: Use vectorized .dt accessor
                    df['date'] = pd.to_datetime(df['date']).dt.strftime('%Y-%m-%d')
                except Exception:
                    df['date'] = df['date'].astype(str)

        if factor_type == FactorType.HIGH_FREQ and 'datetime' in df.columns:
            if not pd.api.types.is_string_dtype(df['datetime']):
                try:
                    # Optimized: Use vectorized .dt accessor
                    df['datetime'] = pd.to_datetime(df['datetime']).dt.strftime('%Y-%m-%d %H:%M:%S')
                except Exception:
                    df['datetime'] = df['datetime'].astype(str)

        # Parallel insertion
        if n_jobs > 1 and len(df) > 5000000:
            logger.info(f"Inserting data with {n_jobs} workers...")
            
            # Get DB config (thread/process safe dict)
            db_config = self.config.get_db_config(asset_type, factor_type)
            
            # Split data by date to optimize caching and partitioning in ClickHouse
            unique_dates = df['date'].unique()
            
            # If too few dates, just split the dataframe
            if len(unique_dates) < n_jobs:
                chunks = np.array_split(df, n_jobs)
            else:
                # Divide dates into chunks
                date_chunks = np.array_split(unique_dates, n_jobs)
                chunks = []
                for date_chunk in date_chunks:
                    if len(date_chunk) > 0:
                        chunk_df = df[df['date'].isin(date_chunk)]
                        if not chunk_df.empty:
                            chunks.append(chunk_df)

            tasks = []
            with concurrent.futures.ProcessPoolExecutor(max_workers=n_jobs) as executor:
                for chunk in chunks:
                    if len(chunk) == 0:
                        continue
                        
                    # Convert to dict for safer serialization
                    chunk_dict = chunk.to_dict(orient='list')
                    
                    tasks.append(
                        executor.submit(
                            _insert_worker,
                            db_config,
                            db_name,
                            factor_id,
                            chunk_dict,
                            factor_type,
                            asset_type
                        )
                    )
                
                # Wait for completion
                total_inserted = 0
                for future in concurrent.futures.as_completed(tasks):
                    try:
                        total_inserted += future.result()
                    except Exception as e:
                        logger.error(f"Worker failed: {e}")
            
            logger.info(f"Parallel insertion complete. Total rows: {total_inserted}")

        else:
            # Serial insertion using optimized batch method
            # Insert data in batches
            batch_size = 5000000
            for i in range(0, len(df), batch_size):
                batch = df.iloc[i:i + batch_size]
                self._insert_batch(db_name, factor_id, batch, factor_type, asset_type)

        # Update upgrade tracking
        self._update_upgrade_tracking(factor_id, df, asset_type, factor_type)

        logger.info(f"Saved {len(df)} factor values for {factor_id}")

    def _insert_batch(
        self,
        db_name: str,
        table_name: str,
        batch: pd.DataFrame,
        factor_type: FactorType,
        asset_type: AssetType
    ) -> None:
        """Insert a batch of factor values"""
        if factor_type == FactorType.HIGH_FREQ:
            columns = "code, date, datetime, factor_value"
            values_template = "('{}', '{}', '{}', {})"
            values_list = [
                values_template.format(
                    row['code'], row['date'], row['datetime'], row['factor_value']
                )
                for _, row in batch.iterrows()
            ]
        else:
            columns = "code, date, factor_value"
            values_template = "('{}', '{}', {})"
            values_list = [
                values_template.format(row['code'], row['date'], row['factor_value'])
                for _, row in batch.iterrows()
            ]

        sql = f"""
            INSERT INTO {db_name}.{table_name} ({columns})
            VALUES {', '.join(values_list)}
        """
        self._execute(sql, asset_type, factor_type)

    def _update_upgrade_tracking(
        self,
        factor_id: str,
        values: pd.DataFrame,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> None:
        """Update the upgrade tracking table"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Get latest date/datetime
        if factor_type == FactorType.HIGH_FREQ and 'datetime' in values.columns:
            latest = values['datetime'].max()
            date_col = "latest_datetime"
            date_val = f"'{latest}'"
            date_col2 = "latest_date"
            date_val2 = "NULL"
        else:
            latest = values['date'].max()
            date_col = "latest_date"
            date_val = f"'{latest}'"
            date_col2 = "latest_datetime"
            date_val2 = "NULL"

        sql = f"""
            INSERT INTO {db_name}.A3_factor_upgrade
            (factor_id, {date_col}, {date_col2}, upgrade_datetime)
            VALUES ('{factor_id}', {date_val}, {date_val2}, now())
        """
        self._execute(sql, asset_type, factor_type)

    def get_last_factor_date(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> Optional[str]:
        """
        Get the last available date for a factor.

        Args:
            factor_id: Factor ID
            asset_type: Asset type
            factor_type: Factor frequency type

        Returns:
            Last date string (YYYY-MM-DD or YYYY-MM-DD HH:MM:SS) or None
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)
        
        # Determine column based on factor type
        date_col = "latest_datetime" if factor_type == FactorType.HIGH_FREQ else "latest_date"
        
        sql = f"""
            SELECT {date_col}
            FROM {db_name}.A3_factor_upgrade
            WHERE factor_id = '{factor_id}'
            ORDER BY upgrade_datetime DESC
            LIMIT 1
        """
        
        try:
            result = self._fetch(sql, asset_type, factor_type)
            if len(result) > 0:
                val = result.iloc[0][date_col]
                # Handle NaT or None
                if pd.isna(val) or val is None:
                    return None
                
                # Convert to string
                if isinstance(val, (pd.Timestamp, datetime, date)):
                    fmt = '%Y-%m-%d %H:%M:%S' if factor_type == FactorType.HIGH_FREQ else '%Y-%m-%d'
                    return val.strftime(fmt)
                return str(val)
        except Exception:
            pass
            
        return None

    def get_factor_values(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType,
        codes: Optional[List[str]] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
    ) -> pd.DataFrame:
        """
        Retrieve factor values from database.

        Args:
            factor_id: Factor ID
            asset_type: Asset type
            factor_type: Factor frequency type
            codes: Optional list of codes to filter
            start_date: Optional start date (YYYY-MM-DD)
            end_date: Optional end date (YYYY-MM-DD)

        Returns:
            DataFrame with factor values
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Build WHERE clause
        conditions = []
        if codes:
            codes_str = ', '.join(f"'{c}'" for c in codes)
            conditions.append(f"code IN ({codes_str})")
        if start_date:
            conditions.append(f"date >= '{start_date}'")
        if end_date:
            conditions.append(f"date <= '{end_date}'")

        where_clause = ""
        if conditions:
            where_clause = "WHERE " + " AND ".join(conditions)

        sql = f"""
            SELECT * FROM {db_name}.{factor_id}
            {where_clause}
            ORDER BY code, date
        """

        return self._fetch(sql, asset_type, factor_type)

    def has_factor_values(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> bool:
        """Check if factor value table exists and has data."""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)
        try:
            # Check table existence via system.tables first to avoid
            # noisy error logs from quantchdb when table doesn't exist
            tables = self._fetch(
                f"SELECT name FROM system.tables "
                f"WHERE database = '{db_name}' AND name = '{factor_id}'",
                asset_type, factor_type
            )
            if len(tables) == 0:
                return False

            # Table exists, check if it has data
            result = self._fetch(
                f"SELECT 1 FROM {db_name}.{factor_id} LIMIT 1",
                asset_type, factor_type
            )
            return len(result) > 0
        except Exception:
            return False

    def save_factor_evaluation(
        self,
        factor_id: str,
        metrics: Dict[str, Any],
        asset_type: AssetType,
        factor_type: FactorType
    ) -> None:
        """
        Save factor evaluation metrics.

        Args:
            factor_id: Factor ID
            metrics: Dictionary of evaluation metrics
            asset_type: Asset type
            factor_type: Factor frequency type
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Get all metric columns
        all_columns = get_evaluation_metric_columns()

        # Build column and value lists
        columns = ["factor_id", "upgrade_datetime"]
        values = [f"'{factor_id}'", "now()"]

        for col in all_columns:
            columns.append(col)
            val = metrics.get(col)
            if val is None:
                values.append("NULL")
            elif isinstance(val, (int, float)):
                values.append(str(val))
            else:
                values.append(f"'{val}'")

        sql = f"""
            INSERT INTO {db_name}.A2_factor_evaluate
            ({', '.join(columns)})
            VALUES ({', '.join(values)})
        """
        self._execute(sql, asset_type, factor_type)
        logger.info(f"Saved evaluation metrics for {factor_id}")

    def get_factor_evaluation(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> Optional[Dict[str, Any]]:
        """Get the latest evaluation metrics for a factor"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        sql = f"""
            SELECT * FROM {db_name}.A2_factor_evaluate
            WHERE factor_id = '{factor_id}'
            ORDER BY upgrade_datetime DESC
            LIMIT 1
        """

        result = self._fetch(sql, asset_type, factor_type)
        if len(result) == 0:
            return None

        return result.iloc[0].to_dict()

    def delete_factor(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> None:
        """
        Delete a factor and all its data.

        Args:
            factor_id: Factor ID to delete
            asset_type: Asset type
            factor_type: Factor frequency type
        """
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        # Delete from basic table
        self._execute(
            f"DELETE FROM {db_name}.A1_factor_basic WHERE factor_id = '{factor_id}'",
            asset_type, factor_type
        )

        # Delete from evaluate table
        self._execute(
            f"DELETE FROM {db_name}.A2_factor_evaluate WHERE factor_id = '{factor_id}'",
            asset_type, factor_type
        )

        # Delete from upgrade table
        self._execute(
            f"DELETE FROM {db_name}.A3_factor_upgrade WHERE factor_id = '{factor_id}'",
            asset_type, factor_type
        )

        # Drop factor value table
        try:
            self._execute(
                f"DROP TABLE IF EXISTS {db_name}.{factor_id}",
                asset_type, factor_type
            )
        except Exception as e:
            logger.warning(f"Failed to drop factor value table: {e}")

        logger.info(f"Deleted factor: {factor_id}")

    def factor_exists(
        self,
        factor_id: str,
        asset_type: AssetType,
        factor_type: FactorType
    ) -> bool:
        """Check if a factor exists in the database"""
        db_name = self.config.get_factor_db_name(asset_type, factor_type)

        sql = f"""
            SELECT count() as cnt FROM {db_name}.A1_factor_basic
            WHERE factor_id = '{factor_id}'
        """

        result = self._fetch(sql, asset_type, factor_type)
        return result.iloc[0]['cnt'] > 0
