"""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 numeric types for Utf8 columns coerced during CSV load ──
        # All columns were forced to Utf8 for safe SQLite storage.  Now try to
        # recover the original numeric types (Int64 / Float64) from the data.
        numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                         pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                         pl.Float32, pl.Float64)
        # Get current schema to know which columns are still Utf8
        live_schema = lf.collect_schema()
        live_names = live_schema.names()
        live_dtypes = list(live_schema.dtypes())
        utf8_cols = [
            live_names[i] for i, dt in enumerate(live_dtypes)
            if dt == pl.Utf8 and not live_names[i].startswith('_')
        ]
        if utf8_cols:
            # Infer types from a small sample to decide Float64 vs Int64
            sample_df = lf.select(utf8_cols).head(1000).collect(streaming=True)
            type_map: dict[str, pl.DataType] = {}
            for col in utf8_cols:
                try:
                    # Try Float64 first (most general numeric type)
                    casted = sample_df[col].cast(pl.Float64, strict=False)
                    non_null = casted.is_not_null().sum()
                    if non_null > len(casted) * 0.5:
                        # Check if all non-null values are integers → use Int64
                        all_int = True
                        for v in casted.head(100).to_list():
                            if v is not None and v != int(v):
                                all_int = False
                                break
                        type_map[col] = pl.Int64 if all_int else pl.Float64
                except Exception:
                    logger.error("unknown failed: {}".format(traceback.format_exc()))
                    pass

            if type_map:
                # Update schema dict to reflect restored types
                for col, dtype in type_map.items():
                    schema[col] = str(dtype)
                # Apply casts to the LazyFrame
                casts = [pl.col(col).cast(dtype, strict=False) for col, dtype in type_map.items()]
                lf = lf.with_columns(casts)
                logger.info('[BACKGROUND] Restored numeric types for %d columns: %s',
                            len(type_map), list(type_map.keys())[:10])

        # ── 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'])
