"""Background tasks for Celery worker."""

import asyncio
import os
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional, List
from celery import shared_task
from celery.utils.log import get_task_logger

from src.autopilot.strategy_selector import get_strategy_selector
from src.core.logger import get_logger

# Use Celery logger
logger = get_task_logger(__name__)

# Report output directory
REPORTS_DIR = Path(os.path.expanduser("~/.local/share/crypto_trader/reports"))
REPORTS_DIR.mkdir(parents=True, exist_ok=True)


def async_to_sync(awaitable):
    """Run async awaitable synchronously."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(awaitable)
    finally:
        loop.close()


@shared_task(bind=True, name="src.worker.tasks.train_model_task")
def train_model_task(
    self, 
    force_retrain: bool = False, 
    bootstrap: bool = True, 
    symbols: List[str] = None,
    days: int = None,
    timeframe: str = None,
    min_samples_per_strategy: int = None
) -> Dict[str, Any]:
    """Train ML model in background (Offloaded from API).
    
    Args:
        force_retrain: Force retraining even if recent model exists
        bootstrap: Automatically bootstrap training data if insufficient
        symbols: List of symbols to bootstrap (defaults to config value if not provided)
        days: Number of days of historical data (defaults to config value)
        timeframe: Timeframe for data (defaults to config value)
        min_samples_per_strategy: Minimum samples per strategy (defaults to config value)
        
    Returns:
        Training metrics
    """
    logger.info(f"Starting model training workflow (force={force_retrain}, bootstrap={bootstrap}, symbols={symbols})")
    
    # CRITICAL: Update state IMMEDIATELY before any blocking operations
    # This must happen in the sync context, not in async_to_sync
    try:
        self.update_state(state='PROGRESS', meta={'step': 'init', 'progress': 5, 'message': 'Initializing...'})
        logger.info("Initial progress state updated")
    except Exception as e:
        logger.error(f"Failed to update initial state: {e}")
        # Don't fail the task, just log it
    
    async def run_training_workflow():
        try:
            logger.info("Step 1: Resetting singletons...")
            # Update progress - use try/except in case update_state fails
            try:
                self.update_state(state='PROGRESS', meta={'step': 'reset', 'progress': 8, 'message': 'Resetting modules...'})
            except:
                pass
            
            # Force reset singletons to ensure no connection sharing across loops
            from src.autopilot import strategy_selector as ss_module
            from src.core import database as db_module
            
            ss_module._strategy_selector = None
            
            # Also reset PerformanceTracker as it holds a reference to the database
            from src.autopilot import performance_tracker as pt_module
            pt_module._performance_tracker = None
            
            db_module._db_instance = None  # Reset global database instance to ensure fresh connection/engine
            
            # Re-import to be safe
            import src.autopilot.strategy_selector
            src.autopilot.strategy_selector._strategy_selector = None
            
            logger.info("Step 2: Getting strategy selector...")
            try:
                self.update_state(state='PROGRESS', meta={'step': 'get_selector', 'progress': 10, 'message': 'Getting strategy selector...'})
            except:
                pass
            
            selector = get_strategy_selector()
            
            # Use passed config or fall back to selector defaults
            bootstrap_days = days if days is not None else selector.bootstrap_days
            bootstrap_timeframe = timeframe if timeframe is not None else selector.bootstrap_timeframe
            bootstrap_min_samples = min_samples_per_strategy if min_samples_per_strategy is not None else selector.min_samples_per_strategy
            bootstrap_symbols = symbols if symbols else selector.bootstrap_symbols
            
            logger.info(f"Config: days={bootstrap_days}, timeframe={bootstrap_timeframe}, min_samples={bootstrap_min_samples}, symbols={bootstrap_symbols}")
            
            # 1. Check existing training data
            logger.info("Step 3: Checking existing training data...")
            try:
                self.update_state(state='PROGRESS', meta={'step': 'check_data', 'progress': 15, 'message': 'Checking training data...'})
            except:
                pass
            training_data = await selector.performance_tracker.prepare_training_data(min_samples_per_strategy=bootstrap_min_samples)
            logger.info(f"Training data check complete. Has data: {training_data is not None and len(training_data.get('X', [])) > 0}")
            
            # 2. Bootstrap if needed
            if (training_data is None or len(training_data.get('X', [])) == 0) and bootstrap:
                self.update_state(state='PROGRESS', meta={'step': 'bootstrap', 'progress': 20, 'message': 'Bootstrapping data...'})
                
                total_samples = 0
                
                for i, symbol in enumerate(bootstrap_symbols):
                    pct = 20 + int(40 * (i / len(bootstrap_symbols)))
                    self.update_state(state='PROGRESS', meta={
                        'step': 'bootstrap', 
                        'progress': pct, 
                        'message': f'Bootstrapping {symbol} ({i+1}/{len(bootstrap_symbols)})...',
                        'details': {'symbol': symbol, 'symbol_index': i+1, 'total_symbols': len(bootstrap_symbols)}
                    })
                    
                    logger.info(f"Starting bootstrap for {symbol} (days={bootstrap_days}, timeframe={bootstrap_timeframe})")
                    res = await selector.bootstrap_training_data(
                        symbol=symbol, 
                        timeframe=bootstrap_timeframe, 
                        days=bootstrap_days
                    )
                    if "error" not in res:
                        total_samples += res.get("total_samples", 0)
                        logger.info(f"Bootstrap for {symbol} complete: {res.get('total_samples', 0)} samples")
                    else:
                        logger.warning(f"Bootstrap for {symbol} failed: {res.get('error', 'Unknown error')}")
                
                if total_samples == 0:
                    raise Exception("Bootstrap failed: No data collected")
                    
            # 3. Train
            self.update_state(state='PROGRESS', meta={'step': 'training', 'progress': 70, 'message': 'Training model...'})
            metrics = await selector.train_model(force_retrain=force_retrain, min_samples_per_strategy=bootstrap_min_samples)
            
            self.update_state(state='PROGRESS', meta={'step': 'complete', 'progress': 100, 'message': 'Training complete'})
            logger.info(f"Model training workflow completed: {metrics}")
            return metrics

        except Exception as e:
            logger.error(f"Model training logic failed: {e}")
            raise e

    try:
        return async_to_sync(run_training_workflow())
    except Exception as e:
        logger.error(f"Model training task failed: {e}")
        # self.update_state(state='FAILURE', meta={'error': str(e)}) # Let Celery handle failure state
        raise e


@shared_task(bind=True, name="src.worker.tasks.bootstrap_task")
def bootstrap_task(
    self,
    days: int = 90,
    timeframe: str = "1h",
    min_samples: int = 10,
    symbols: Optional[List[str]] = None
) -> Dict[str, Any]:
    """Bootstrap training data in background.
    
    Args:
        days: Number of days to look back
        timeframe: Candle timeframe
        min_samples: Minimum samples per strategy
        symbols: List of training symbols
        
    Returns:
        Bootstrap results
    """
    logger.info(f"Starting data bootstrap (days={days}, timeframe={timeframe})")
    try:
        selector = get_strategy_selector()
        
        results = async_to_sync(selector.bootstrap_training_data(
            days=days,
            timeframe=timeframe,
        ))
        
        logger.info("Bootstrap task completed")
        return results
    except Exception as e:
        logger.error(f"Bootstrap task failed: {e}")
        raise


@shared_task(bind=True, name="src.worker.tasks.optimize_strategy_task")
def optimize_strategy_task(
    self,
    strategy_type: str,
    symbol: str,
    param_ranges: Dict[str, tuple],
    method: str = "genetic",
    population_size: int = 50,
    generations: int = 100,
    maximize: bool = True
) -> Dict[str, Any]:
    """Optimize strategy parameters in background.
    
    Args:
        strategy_type: Type of strategy to optimize (e.g., 'rsi', 'macd')
        symbol: Trading symbol for backtesting
        param_ranges: Dictionary of parameter -> (min, max) range
        method: Optimization method ('genetic', 'grid', 'bayesian')
        population_size: Population size for genetic algorithm
        generations: Number of generations
        maximize: True to maximize objective, False to minimize
        
    Returns:
        Optimization results with best parameters
    """
    logger.info(f"Starting strategy optimization: {strategy_type} on {symbol}")
    
    try:
        self.update_state(state='PROGRESS', meta={
            'step': 'init',
            'progress': 5,
            'message': f'Initializing {method} optimization...'
        })
        
        if method == "genetic":
            from src.optimization.genetic import GeneticOptimizer
            optimizer = GeneticOptimizer()
            
            # Define objective function (backtesting)
            def objective(params: Dict[str, Any]) -> float:
                """Run backtest with given parameters."""
                try:
                    from src.backtesting.engine import BacktestEngine
                    from src.strategies.base import get_strategy_registry
                    
                    registry = get_strategy_registry()
                    strategy_class = registry.get(strategy_type)
                    if not strategy_class:
                        return float('-inf')
                    
                    # Create strategy with params
                    strategy = strategy_class(name=f"{strategy_type}_opt", **params)
                    
                    # Run quick backtest
                    engine = BacktestEngine()
                    results = async_to_sync(engine.run(
                        strategy=strategy,
                        symbol=symbol,
                        start_date=datetime.now().replace(day=1),  # Last month
                        end_date=datetime.now(),
                        initial_capital=10000.0
                    ))
                    
                    # Return Sharpe ratio as objective
                    return results.get('sharpe_ratio', 0.0)
                except Exception as e:
                    logger.warning(f"Objective evaluation failed: {e}")
                    return float('-inf')
            
            self.update_state(state='PROGRESS', meta={
                'step': 'optimizing',
                'progress': 20,
                'message': f'Running genetic algorithm ({generations} generations)...'
            })
            
            result = optimizer.optimize(
                param_ranges=param_ranges,
                objective_function=objective,
                population_size=population_size,
                generations=generations,
                maximize=maximize
            )
            
        elif method == "grid":
            from src.optimization.grid_search import GridSearchOptimizer
            optimizer = GridSearchOptimizer()
            
            self.update_state(state='PROGRESS', meta={
                'step': 'optimizing',
                'progress': 20,
                'message': 'Running grid search...'
            })
            
            # Grid search implementation would go here
            result = {"best_params": {}, "best_score": 0.0, "message": "Grid search not fully implemented"}
            
        else:
            result = {"error": f"Unknown optimization method: {method}"}
        
        self.update_state(state='PROGRESS', meta={
            'step': 'complete',
            'progress': 100,
            'message': 'Optimization complete'
        })
        
        logger.info(f"Optimization completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Strategy optimization failed: {e}")
        self.update_state(state='FAILURE', meta={'error': str(e)})
        raise


@shared_task(bind=True, name="src.worker.tasks.generate_report_task")
def generate_report_task(self, report_type: str, params: Dict[str, Any]) -> Dict[str, Any]:
    """Generate report in background.
    
    Args:
        report_type: 'performance', 'backtest', 'tax', 'trades'
        params: Report parameters including:
            - start_date: Start date for report
            - end_date: End date for report
            - format: 'pdf' or 'csv'
            
    Returns:
        Result with file path
    """
    logger.info(f"Starting report generation: {report_type}")
    
    try:
        self.update_state(state='PROGRESS', meta={
            'step': 'init',
            'progress': 10,
            'message': f'Preparing {report_type} report...'
        })
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_format = params.get('format', 'pdf')
        
        if report_type == "performance":
            self.update_state(state='PROGRESS', meta={
                'step': 'collecting',
                'progress': 30,
                'message': 'Collecting performance metrics...'
            })
            
            # Collect metrics from database
            from src.portfolio.analytics import get_portfolio_analytics
            analytics = get_portfolio_analytics()
            metrics = async_to_sync(analytics.get_portfolio_metrics())
            
            if report_format == 'pdf':
                from src.reporting.pdf_generator import get_pdf_generator
                generator = get_pdf_generator()
                
                filepath = REPORTS_DIR / f"performance_{timestamp}.pdf"
                
                self.update_state(state='PROGRESS', meta={
                    'step': 'generating',
                    'progress': 70,
                    'message': 'Generating PDF...'
                })
                
                success = generator.generate_performance_report(filepath, metrics)
                
                if success:
                    result = {"status": "success", "file": str(filepath), "type": "pdf"}
                else:
                    result = {"status": "error", "message": "PDF generation failed"}
            else:
                # CSV export
                from src.reporting.csv_exporter import get_csv_exporter
                exporter = get_csv_exporter()
                
                filepath = REPORTS_DIR / f"performance_{timestamp}.csv"
                success = exporter.export_metrics(filepath, metrics)
                result = {"status": "success" if success else "error", "file": str(filepath)}
                
        elif report_type == "trades":
            self.update_state(state='PROGRESS', meta={
                'step': 'collecting',
                'progress': 30,
                'message': 'Collecting trade history...'
            })
            
            from src.reporting.csv_exporter import get_csv_exporter
            from src.core.database import get_database, Trade
            from sqlalchemy import select
            
            db = get_database()
            
            async def get_trades():
                async with db.get_session() as session:
                    stmt = select(Trade).order_by(Trade.timestamp.desc()).limit(1000)
                    result = await session.execute(stmt)
                    return result.scalars().all()
            
            trades = async_to_sync(get_trades())
            
            filepath = REPORTS_DIR / f"trades_{timestamp}.csv"
            exporter = get_csv_exporter()
            
            self.update_state(state='PROGRESS', meta={
                'step': 'generating',
                'progress': 70,
                'message': 'Exporting trades...'
            })
            
            success = exporter.export_trades(filepath, trades)
            result = {"status": "success" if success else "error", "file": str(filepath), "count": len(trades)}
            
        elif report_type == "tax":
            self.update_state(state='PROGRESS', meta={
                'step': 'calculating',
                'progress': 30,
                'message': 'Calculating tax obligations...'
            })
            
            from src.reporting.tax_reporter import get_tax_reporter
            reporter = get_tax_reporter()
            
            year = params.get('year', datetime.now().year)
            method = params.get('method', 'fifo')  # fifo, lifo, specific
            
            filepath = REPORTS_DIR / f"tax_report_{year}_{timestamp}.pdf"
            
            self.update_state(state='PROGRESS', meta={
                'step': 'generating',
                'progress': 70,
                'message': 'Generating tax report...'
            })
            
            success = reporter.generate_annual_report(year, method, filepath)
            result = {"status": "success" if success else "error", "file": str(filepath), "year": year}
            
        else:
            result = {"status": "error", "message": f"Unknown report type: {report_type}"}
        
        self.update_state(state='PROGRESS', meta={
            'step': 'complete',
            'progress': 100,
            'message': 'Report generated'
        })
        
        logger.info(f"Report generation completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Report generation failed: {e}")
        self.update_state(state='FAILURE', meta={'error': str(e)})
        raise


@shared_task(bind=True, name="src.worker.tasks.export_data_task")
def export_data_task(
    self,
    export_type: str,
    params: Dict[str, Any]
) -> Dict[str, Any]:
    """Export data in background.
    
    Args:
        export_type: 'orders', 'positions', 'market_data'
        params: Export parameters
        
    Returns:
        Result with file path
    """
    logger.info(f"Starting data export: {export_type}")
    
    try:
        self.update_state(state='PROGRESS', meta={
            'step': 'init',
            'progress': 10,
            'message': f'Preparing {export_type} export...'
        })
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        from src.reporting.csv_exporter import get_csv_exporter
        exporter = get_csv_exporter()
        
        filepath = REPORTS_DIR / f"{export_type}_{timestamp}.csv"
        
        self.update_state(state='PROGRESS', meta={
            'step': 'exporting',
            'progress': 50,
            'message': 'Exporting data...'
        })
        
        # Export based on type
        if export_type == "orders":
            from src.core.database import get_database, Order
            from sqlalchemy import select
            
            db = get_database()
            
            async def get_orders():
                async with db.get_session() as session:
                    stmt = select(Order).order_by(Order.created_at.desc())
                    result = await session.execute(stmt)
                    return result.scalars().all()
            
            data = async_to_sync(get_orders())
            success = exporter.export_orders(filepath, data)
            
        elif export_type == "positions":
            from src.core.database import get_database, Position
            from sqlalchemy import select
            
            db = get_database()
            
            async def get_positions():
                async with db.get_session() as session:
                    stmt = select(Position)
                    result = await session.execute(stmt)
                    return result.scalars().all()
            
            data = async_to_sync(get_positions())
            success = exporter.export_positions(filepath, data)
            
        else:
            success = False
            
        result = {
            "status": "success" if success else "error",
            "file": str(filepath),
            "export_type": export_type
        }
        
        self.update_state(state='PROGRESS', meta={
            'step': 'complete',
            'progress': 100,
            'message': 'Export complete'
        })
        
        logger.info(f"Data export completed: {result}")
        return result
        
    except Exception as e:
        logger.error(f"Data export failed: {e}")
        self.update_state(state='FAILURE', meta={'error': str(e)})
        raise
