"""
Factor Manager for FactorDB

This module provides the central management class for factor lifecycle operations.
"""

from typing import Optional, List, Dict, Any, Union
from pathlib import Path
import pandas as pd
import logging

from .config import FactorDBConfig, AssetType, FactorType
from .factor import BaseFactor, MinedFactor, ExpressionFactor, FactorMetadata
from ..storage.clickhouse_storage import ClickHouseStorage
from ..parsers.parser_factory import parse_factor_file, get_parser

logger = logging.getLogger(__name__)


class FactorManager:
    """
    Central manager for factor lifecycle operations.

    Provides high-level APIs for:
    - Factor registration and storage
    - Factor retrieval and querying
    - Factor calculation and updates
    - Batch operations for mining results
    """

    def __init__(self, config: Optional[FactorDBConfig] = None):
        self.config = config or FactorDBConfig()
        self.storage = ClickHouseStorage(self.config)
        self._evaluator = None

    @property
    def evaluator(self):
        """Lazy load evaluator to avoid circular imports"""
        if self._evaluator is None:
            from ..evaluators.factor_evaluator import FactorEvaluator
            self._evaluator = FactorEvaluator(self.config)
        return self._evaluator

    def register_factor(
        self,
        factor: BaseFactor,
        calculate_now: bool = False,
        market_data: Optional[pd.DataFrame] = None
    ) -> str:
        """
        Register a factor in the database.

        Args:
            factor: Factor instance to register
            calculate_now: Whether to calculate and save factor values immediately
            market_data: Market data for calculation (required if calculate_now=True)

        Returns:
            Assigned factor ID
        """
        # Create metadata
        metadata = factor.create_metadata(factor_id="")

        # Register in database (ID will be assigned)
        factor_id = self.storage.register_factor(
            metadata,
            factor.asset_type,
            factor.factor_type
        )

        # Update factor with assigned ID
        factor.factor_id = factor_id

        # Calculate and save values if requested
        if calculate_now:
            if market_data is None:
                raise ValueError("market_data is required when calculate_now=True")
            self.calculate_and_save(factor, market_data)

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

    def register_from_expression(
        self,
        expression: str,
        name: Optional[str] = None,
        explanation: Optional[str] = None,
        source_framework: str = "auto",
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY
    ) -> str:
        """
        Register a factor from an expression string.

        Args:
            expression: Factor expression
            name: Optional factor name
            explanation: Optional explanation
            source_framework: Source framework for expression parsing
            asset_type: Asset type
            factor_type: Factor frequency type

        Returns:
            Assigned factor ID
        """
        factor = ExpressionFactor(
            expression=expression,
            name=name,
            explanation=explanation,
            source_framework=source_framework,
            asset_type=asset_type,
            factor_type=factor_type,
            config=self.config
        )
        return self.register_factor(factor)

    def register_mined_factors(
        self,
        file_path: Union[str, Path],
        framework: Optional[str] = None,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        max_factors: Optional[int] = None,
        min_ic: Optional[float] = None
    ) -> List[str]:
        """
        Register factors from a mining result file.

        Args:
            file_path: Path to the mining result JSON file
            framework: Framework name (auto-detected if not provided)
            asset_type: Asset type
            factor_type: Factor frequency type
            max_factors: Maximum number of factors to register
            min_ic: Minimum absolute IC to filter factors

        Returns:
            List of registered factor IDs
        """
        # Parse factors from file
        factors = parse_factor_file(
            file_path,
            framework=framework,
            asset_type=asset_type,
            factor_type=factor_type
        )

        # Filter by IC if specified
        if min_ic is not None:
            factors = [
                f for f in factors
                if abs(f.get_metric("ic_mean", 0) or f.get_metric("mean_ic", 0)) >= min_ic
            ]

        # Limit number of factors
        if max_factors is not None:
            factors = factors[:max_factors]

        # Register each factor
        factor_ids = []
        for factor in factors:
            try:
                factor_id = self.register_factor(factor)
                factor_ids.append(factor_id)
            except Exception as e:
                logger.warning(f"Failed to register factor: {e}")

        logger.info(f"Registered {len(factor_ids)} factors from {file_path}")
        return factor_ids

    def get_factor(
        self,
        factor_id: str,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY
    ) -> Optional[FactorMetadata]:
        """
        Get factor metadata by ID.

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

        Returns:
            FactorMetadata or None if not found
        """
        return self.storage.get_factor_metadata(factor_id, asset_type, factor_type)

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

        Args:
            asset_type: Asset type
            factor_type: Factor frequency type
            limit: Maximum number of factors to return

        Returns:
            DataFrame with factor information
        """
        return self.storage.list_factors(asset_type, factor_type, limit)

    def calculate_and_save(
        self,
        factor: BaseFactor,
        market_data: pd.DataFrame,
        overwrite: bool = False
    ) -> None:
        """
        Calculate factor values and save to database.

        Args:
            factor: Factor instance
            market_data: Market data DataFrame
            overwrite: Whether to overwrite existing values
        """
        if factor.factor_id is None:
            raise ValueError("Factor must be registered first")

        # Check if values already exist
        if not overwrite and self.storage.factor_exists(
            factor.factor_id, factor.asset_type, factor.factor_type
        ):
            existing = self.get_factor_values(
                factor.factor_id, factor.asset_type, factor.factor_type
            )
            if len(existing) > 0:
                logger.info(f"Factor {factor.factor_id} already has values. Use overwrite=True to replace.")
                return

        # Calculate values
        values = factor.calculate(market_data)

        # Convert to DataFrame format
        if isinstance(values, pd.Series):
            if isinstance(values.index, pd.MultiIndex):
                df = values.reset_index()
                df.columns = ['code', 'date', 'factor_value']
            else:
                df = pd.DataFrame({'factor_value': values})
                df['code'] = market_data['code'].values
                df['date'] = market_data['date'].values
        else:
            df = values

        # Save to database
        self.storage.save_factor_values(
            factor.factor_id,
            df,
            factor.asset_type,
            factor.factor_type
        )

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

    def update_factors(
        self,
        factors: List[Union[str, BaseFactor]],
        market_data: pd.DataFrame,
        n_jobs: int = 4
    ) -> Dict[str, str]:
        """
        Update multiple factors to the latest date.

        This method automatically handles:
        1. Checking the last update date of each factor
        2. Slicing new values from the calculation result (to avoid overwriting history)
        3. Parallel processing of multiple factors

        Note on Market Data:
        The provided `market_data` MUST contain sufficient history (lookback window)
        before the expected update start date. The method calculates on the full
        `market_data` but only saves the new portion.

        Args:
            factors: List of factor IDs (str) or factor instances (BaseFactor)
            market_data: Market data DataFrame (must have sufficient history)
            n_jobs: Number of parallel threads (default 4)

        Returns:
            Dictionary mapping factor_id to status ("Updated", "Skipped", "Error")
        """
        import concurrent.futures

        logger.info(f"Updating {len(factors)} factors with {n_jobs} threads")
        results = {}

        def _process_single_factor(factor_input):
            try:
                # Resolve factor instance
                if isinstance(factor_input, str):
                    factor_id = factor_input
                    # Only ExpressionFactor can be loaded from DB automatically
                    metadata = self.get_factor(factor_id)
                    if not metadata:
                        return factor_id, "Error: Not found"
                    if not metadata.expression:
                        return factor_id, "Error: No expression (CustomFactor requires instance)"
                    
                    factor = ExpressionFactor(
                        expression=metadata.expression,
                        name=metadata.name,
                        source_framework=metadata.source_framework or "auto",
                        asset_type=metadata.asset_type,
                        factor_type=metadata.factor_type,
                        config=self.config
                    )
                    factor.factor_id = factor_id
                else:
                    factor = factor_input
                    if not factor.factor_id:
                        return "Unknown", "Error: Factor has no ID"
                    factor_id = factor.factor_id

                # 1. Get last update date
                last_date = self.storage.get_last_factor_date(
                    factor_id, factor.asset_type, factor.factor_type
                )

                # 2. Calculate values on FULL data
                # We calculate on full data to ensure indicators (MA, EMA, etc.) have correct warmup
                full_values = factor.calculate(market_data)
                
                # Convert to DataFrame standard format
                if isinstance(full_values, pd.Series):
                    if isinstance(full_values.index, pd.MultiIndex):
                        df_values = full_values.reset_index()
                        df_values.columns = ['code', 'date', 'factor_value']
                    else:
                        # Should not happen for standard calculate, but safe fallback
                        return factor_id, "Error: Invalid calculation result format"
                else:
                    df_values = full_values

                # 3. Slice new data
                if last_date:
                    # Only keep rows AFTER last_date
                    new_values = df_values[df_values['date'] > last_date]
                    if new_values.empty:
                        return factor_id, "Skipped: Up to date"
                else:
                    # New factor, save all
                    new_values = df_values

                # 4. Save to DB
                # Note: save_factor_values handles its own parallelism if data is large,
                # but here we likely have small daily updates, so we set n_jobs=1 for inner save
                # to avoid exploding thread count.
                self.storage.save_factor_values(
                    factor_id,
                    new_values,
                    factor.asset_type,
                    factor.factor_type,
                    n_jobs=1  # Sequential save because we are parallel in outer loop
                )
                
                return factor_id, f"Updated: {len(new_values)} rows"

            except Exception as e:
                logger.error(f"Failed to update factor {factor_input}: {e}")
                return str(factor_input), f"Error: {str(e)}"

        # Run parallel updates
        # Use ThreadPoolExecutor because:
        # 1. Market data is large (avoid pickling overhead of ProcessPool)
        # 2. Pandas/Numpy release GIL for calculation
        # 3. DB IO releases GIL
        with concurrent.futures.ThreadPoolExecutor(max_workers=n_jobs) as executor:
            futures = {executor.submit(_process_single_factor, f): f for f in factors}
            
            for future in concurrent.futures.as_completed(futures):
                fid, status = future.result()
                results[fid] = status
                logger.info(f"  {fid}: {status}")

        return results

    def calculate_from_expression(
        self,
        factor_id: str,
        market_data: pd.DataFrame,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY
    ) -> None:
        """
        Calculate and save factor values from stored expression.

        Args:
            factor_id: Factor ID
            market_data: Market data DataFrame
            asset_type: Asset type
            factor_type: Factor frequency type
        """
        # Get factor metadata
        metadata = self.get_factor(factor_id, asset_type, factor_type)
        if metadata is None:
            raise ValueError(f"Factor not found: {factor_id}")

        if not metadata.expression:
            raise ValueError(f"Factor {factor_id} has no expression")

        # Create expression factor
        factor = ExpressionFactor(
            expression=metadata.expression,
            name=metadata.name,
            explanation=metadata.explanation,
            asset_type=asset_type,
            factor_type=factor_type,
            config=self.config
        )
        factor.factor_id = factor_id

        # Calculate and save
        self.calculate_and_save(factor, market_data)

    def get_factor_values(
        self,
        factor_id: str,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        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
        """
        return self.storage.get_factor_values(
            factor_id, asset_type, factor_type, codes, start_date, end_date
        )

    def evaluate_factor(
        self,
        factor_id: str,
        returns_data: pd.DataFrame,
        market_cap_data: Optional[pd.DataFrame] = None,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        save_results: bool = True
    ) -> Dict[str, Any]:
        """
        Evaluate a factor and optionally save results.

        Args:
            factor_id: Factor ID
            returns_data: DataFrame with returns (code, date, return columns)
            market_cap_data: Optional DataFrame with market cap data
            asset_type: Asset type
            factor_type: Factor frequency type
            save_results: Whether to save evaluation results to database

        Returns:
            Dictionary of evaluation metrics
        """
        # Get factor values
        factor_values = self.get_factor_values(factor_id, asset_type, factor_type)

        if len(factor_values) == 0:
            raise ValueError(f"No factor values found for {factor_id}")

        # Calculate metrics
        metrics = self.evaluator.evaluate(
            factor_values=factor_values,
            returns_data=returns_data,
            market_cap_data=market_cap_data
        )

        # Save results if requested
        if save_results:
            self.storage.save_factor_evaluation(
                factor_id, metrics, asset_type, factor_type
            )

        return metrics

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

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

        Returns:
            Dictionary of evaluation metrics or None
        """
        return self.storage.get_factor_evaluation(factor_id, asset_type, factor_type)

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

        Args:
            factor_id: Factor ID
            asset_type: Asset type
            factor_type: Factor frequency type
        """
        self.storage.delete_factor(factor_id, asset_type, factor_type)
        logger.info(f"Deleted factor: {factor_id}")

    def batch_evaluate(
        self,
        factor_ids: List[str],
        returns_data: pd.DataFrame,
        market_cap_data: Optional[pd.DataFrame] = None,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        save_results: bool = True
    ) -> Dict[str, Dict[str, Any]]:
        """
        Evaluate multiple factors.

        Args:
            factor_ids: List of factor IDs
            returns_data: DataFrame with returns
            market_cap_data: Optional DataFrame with market cap data
            asset_type: Asset type
            factor_type: Factor frequency type
            save_results: Whether to save evaluation results

        Returns:
            Dictionary mapping factor IDs to their metrics
        """
        results = {}
        for factor_id in factor_ids:
            try:
                metrics = self.evaluate_factor(
                    factor_id, returns_data, market_cap_data,
                    asset_type, factor_type, save_results
                )
                results[factor_id] = metrics
            except Exception as e:
                logger.warning(f"Failed to evaluate {factor_id}: {e}")
                results[factor_id] = {"error": str(e)}

        return results

    def search_factors(
        self,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        min_ic: Optional[float] = None,
        min_icir: Optional[float] = None,
        direction: Optional[int] = None,
        limit: int = 100
    ) -> pd.DataFrame:
        """
        Search factors by evaluation criteria.

        Args:
            asset_type: Asset type
            factor_type: Factor frequency type
            min_ic: Minimum absolute IC mean
            min_icir: Minimum absolute IC IR
            direction: Factor direction (1 for long, -1 for short)
            limit: Maximum number of results

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

        # Build WHERE conditions
        conditions = []
        if min_ic is not None:
            conditions.append(f"abs(ic_mean) >= {min_ic}")
        if min_icir is not None:
            conditions.append(f"abs(ic_ir) >= {min_icir}")
        if direction is not None:
            conditions.append(f"direction = {direction}")

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

        sql = f"""
            SELECT b.factor_id, b.name, b.expression,
                   e.ic_mean, e.rank_ic_mean, e.ic_ir, e.direction
            FROM {db_name}.A1_factor_basic b
            LEFT JOIN {db_name}.A2_factor_evaluate e ON b.factor_id = e.factor_id
            {where_clause}
            ORDER BY abs(e.ic_ir) DESC
            LIMIT {limit}
        """

        return self.storage._fetch(sql)

    def export_factor_library(
        self,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        output_path: Optional[Union[str, Path]] = None
    ) -> pd.DataFrame:
        """
        Export factor library with all metadata and evaluation metrics.

        Args:
            asset_type: Asset type
            factor_type: Factor frequency type
            output_path: Optional path to save CSV

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

        sql = f"""
            SELECT b.*, e.*
            FROM {db_name}.A1_factor_basic b
            LEFT JOIN {db_name}.A2_factor_evaluate e ON b.factor_id = e.factor_id
            ORDER BY b.factor_id
        """

        df = self.storage._fetch(sql)

        if output_path:
            df.to_csv(output_path, index=False)
            logger.info(f"Exported factor library to {output_path}")

        return df
