"""Background pipeline: CSV loading, SVD, and pipeline worker."""
import os
import threading
import traceback
import logging

from analysis.appdata import ensure_appdata_dir
from analysis.constants import MAX_CLUSTER_FEATURES, RANDOM_SEED
from analysis.models import AnalysisRun
from analysis.profile_util import profile

logger = logging.getLogger(__name__)


def _run_pipeline_worker(run_id, analysis_fn, **ctx):
    """Unified background worker: setup / ORM / progress / error handling."""
    import django
    os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
    os.environ['DJANGO_ALLOW_ASYNC_UNSAFE'] = 'true'
    django.setup()

    from analysis.db_utils import retry_on_lock
    from django.db import close_old_connections

    close_old_connections()

    @retry_on_lock()
    def _get_run():
        return AnalysisRun.objects.get(id=run_id)

    run = _get_run()
    try:
        analysis_fn(run, ctx)
    except Exception:
        tb = traceback.format_exc()
        logger.error(tb)
        run.status = 'failed'
        run.error_message = tb
        run.run_log += f'\n[FATAL] {tb}'
        retry_on_lock()(lambda: run.save(update_fields=['status', 'error_message', 'run_log']))()


@profile
def _background_process(run_id, upload_dir):
    """Background: stream-load CSVs → save to SQLite → ready for analysis."""
    import django
    os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
    django.setup()
    from analysis.models import AnalysisRun
    from analysis.session_store import SessionStore

    run = AnalysisRun.objects.get(id=run_id)
    store = SessionStore()
    try:
        # Check upload directory still exists (may have been temp-cleaned)
        if not upload_dir.is_dir():
            msg = '上传数据目录已不存在，请重新上传'
            logger.error('[BACKGROUND] %s (run_id=%s, upload_dir=%s)', msg, run_id, upload_dir)
            run.status = 'failed'
            run.error_message = msg
            run.save(update_fields=['status', 'error_message'])
            return

        import glob
        import polars as pl
        csv_paths = sorted(glob.glob(str(upload_dir / '*.csv')))
        if not csv_paths:
            msg = '没有找到 CSV 文件'
            logger.error('[BACKGROUND] %s (run_id=%s)', msg, run_id)
            run.status = 'failed'
            run.error_message = msg
            run.save(update_fields=['status', 'error_message'])
            return

        # ── Detect lat/lon columns for schema override ──
        import csv as _csv
        _LAT_LON_KW = ('latd', 'lond', 'latitude', 'longitude')
        lat_lon_overrides = {}
        try:
            with open(csv_paths[0], 'r', encoding='utf-8') as _f:
                _reader = _csv.reader(_f)
                _headers = next(_reader, [])
            for _col in _headers:
                _norm = _col.lower().replace('-', '_').replace('.', '_').replace(' ', '_')
                if any(kw in _norm for kw in _LAT_LON_KW):
                    lat_lon_overrides[_col] = pl.Utf8
                    logger.info('[BACKGROUND] Forcing lat/lon column %r to Utf8 for safe parsing', _col)
        except Exception:
            logger.warning('[BACKGROUND] Failed to detect lat/lon columns', exc_info=True)

        total_files = len(csv_paths)
        logger.info('[BACKGROUND] Found %d CSV files in %s', total_files, upload_dir)
        # Adaptive batch size: aim for ~20 batches total, cap at 500 files per batch
        batch_size = max(1, min(500, total_files // 20 or 1))
        total_batches = (total_files + batch_size - 1) // batch_size

        run.status = 'loading'; run.progress_pct = 10
        run.progress_msg = f'正在流式加载 {total_files} 个文件（共 {total_batches} 批）...'
        run.save(update_fields=['status', 'progress_pct', 'progress_msg'])

        from analysis.data_loader import save_to_db, drop_sqlite_table, load_from_db, _dtype_to_sqlite

        table_name = f'_data_{run_id}'

        # ── Phase 1: Pre-scan all files for full union column set ──
        all_columns = []
        for _path in csv_paths:
            with open(_path, 'r', encoding='utf-8') as _f:
                _reader = _csv.reader(_f)
                _headers = next(_reader, [])
                for _h in _headers:
                    if _h not in all_columns:
                        all_columns.append(_h)
        reference_columns = sorted(all_columns, key=str.lower)
        schema = {col: 'Utf8' for col in all_columns}

        # ── Phase 2: Create SQLite table once with all columns as TEXT ──
        import sqlite3 as _sqlite3
        from django.conf import settings as _dj_settings
        _db_path = _dj_settings.DATABASES['default']['NAME']
        _conn = _sqlite3.connect(_db_path)
        try:
            _col_defs = [f'"{_col}" TEXT' for _col in reference_columns]
            _conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
            _conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(_col_defs)})')
            _conn.commit()
        finally:
            _conn.close()

        # ── Phase 3: Batched collection → SQLite append ──
        schema_overrides = {col: pl.Utf8 for col in reference_columns}
        completed_batches = 0
        total_rows = 0
        for i in range(0, total_files, batch_size):
            batch = csv_paths[i:i+batch_size]
            # Scan each file INDIVIDUALLY — batch scan fails when files have different schemas
            dfs = []
            for fp in batch:
                lf = pl.scan_csv(fp, infer_schema_length=0,
                                 schema_overrides=schema_overrides)
                df = lf.collect(streaming=True)
                # Fill missing columns with NULL
                for col in reference_columns:
                    if col not in df.columns:
                        df = df.with_columns(pl.lit(None).alias(col))
                # Select all reference_columns in canonical order
                df = df.select(reference_columns)
                dfs.append(df)
            df_batch = pl.concat(dfs, how='diagonal_relaxed')

            # Append to pre-created SQLite table
            save_to_db(run.id, df_batch.lazy(), mode='append')

            # Delete temp files after successful SQLite write
            for f in batch:
                try:
                    os.unlink(f)
                except OSError:
                    pass

            completed_batches += 1
            run.progress_pct = int((completed_batches / total_batches) * 80 + 10)
            run.progress_msg = f'正在加载数据... ({completed_batches}/{total_batches} 批)'
            run.save(update_fields=['progress_pct', 'progress_msg'])

        # ── Phase 3: Reload from SQLite with correct types ──
        lf = load_from_db(table_name, schema_overrides=schema)
        if lf is None:
            raise RuntimeError(f'SQLite 表 {table_name} 加载失败')

        # Apply _coerce_to_float on lat/lon columns (restore Float64 from Utf8 storage)
        from analysis.data_loader import _coerce_to_float as _coerce_to_float_fn
        for lat_lon_col in lat_lon_overrides:
            if lat_lon_col in lf.columns:
                lf = lf.with_columns(
                    pl.col(lat_lon_col).map_batches(_coerce_to_float_fn, return_dtype=pl.Float64).alias(lat_lon_col)
                )

        # ── Restore column types after SQLite storage ──
        # All columns were stored as Utf8 for SQLite.  Recover original types
        # from the classification system — ONLY the 8 allowed types.
        from analysis.type_classifier import classify_schema, DataType
        type_map = classify_schema(lf)
        # Apply Polars-level type casts based on classification
        dtypes = {
            DataType.FLOAT: pl.Float64,
            DataType.INT: pl.Int64,
            DataType.BOOL: pl.Boolean,
            DataType.TIMESTAMP: pl.Float64,
        }
        casts = []
        for col, dt in type_map.items():
            target = dtypes.get(dt)
            if target and col in lf.collect_schema().names():
                casts.append(pl.col(col).cast(target, strict=False))
                schema[col] = str(target)
        if casts:
            lf = lf.with_columns(casts)
            restored = [(col, dt.value) for col, dt in type_map.items() if dt in dtypes][:15]
            logger.info('[BACKGROUND] Type restoration: %s', restored)

        # ── SVD dimensionality reduction ──
        svd_components = 0
        # Only operate on columns with known numeric types in the schema
        numeric_dtypes = {'Int64', 'Float64', 'Int32', 'Float32',
                          'Int8', 'Int16', 'UInt8', 'UInt16', 'UInt32', 'UInt64'}
        numeric_cols = [c for c, dt in schema.items() if dt in numeric_dtypes]
        n_numeric = len(numeric_cols)

        if n_numeric >= 2:
            n_components = min(MAX_CLUSTER_FEATURES, n_numeric - 1)
            try:
                from sklearn.decomposition import TruncatedSVD  # fmt: skip
                import numpy as np  # fmt: skip

                # Collect numeric data for SVD (full materialise → transform → add back)
                df_full = lf.select(numeric_cols).collect(streaming=True)
                X = df_full.to_numpy()
                X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)

                svd = TruncatedSVD(n_components=n_components, random_state=RANDOM_SEED)
                X_svd = svd.fit_transform(X)
                svd_components = n_components

                # Build a new LazyFrame that includes _svd_* columns
                # Strategy: collect full LazyFrame once, add SVD columns, re-wrap as lazy
                df_all = lf.collect(streaming=True)
                for i in range(n_components):
                    col_name = f'_svd_{i}'
                    schema[col_name] = 'Float64'
                    df_all = df_all.with_columns(
                        pl.Series(col_name, X_svd[:, i])
                    )
                lf = df_all.lazy()

                logger.info('[BACKGROUND] SVD: %d components from %d numeric columns '
                            '(explained variance ratio sum=%.4f)',
                            svd_components, n_numeric,
                            float(svd.explained_variance_ratio_.sum()))
            except Exception as e:
                logger.warning('[BACKGROUND] SVD failed (non-fatal): %s', e)
        elif n_numeric == 1:
            logger.info('[BACKGROUND] SVD skipped: only 1 numeric column (%s)',
                        numeric_cols[0])

        # Get row count
        df_count = lf.select(pl.len()).collect(streaming=True)
        total_rows = df_count[0, 0] if df_count.height > 0 else 0

        run.sqlite_table = table_name
        run.total_flows = total_rows
        run.save(update_fields=['sqlite_table', 'total_flows'])

        ds_id = f'upload_{run.display_id}'
        store.store_dataset(ds_id, lf, schema=schema, metadata={
            'row_count': total_rows, 'file_count': total_files, 'upload_dir': str(upload_dir),
            'csv_glob': str(upload_dir / '*.csv'),
            'sqlite_table': table_name,
            'svd_components': svd_components,
        })

        # ── Phase 3: Data ready for filtering/clustering via MCP tools ──
        run.status = 'ready'; run.progress_pct = 60; run.progress_msg = '数据准备就绪'
        run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
    except Exception:
        tb = traceback.format_exc()
        logger.error(tb)
        # Drop partial SQLite table on failure (idempotent)
        try:
            from analysis.data_loader import drop_sqlite_table  # fmt: skip
            drop_sqlite_table(f'_data_{run_id}')
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            pass  # Expected: drop_sqlite_table may fail if table already gone
        run.status = 'failed'
        run.error_message = tb
        run.run_log += f'\n[ERROR] {tb}'
        run.save(update_fields=['status', 'error_message', 'run_log'])