"""Import data from SQL backup with content-based deduplication."""

from __future__ import annotations

import logging
import re
import sqlite3
from pathlib import Path
from typing import TYPE_CHECKING, Any

from open_agent_kit.features.team.activity.store.backup.hashes import (
    compute_activity_hash,
    compute_observation_hash,
    compute_prompt_batch_hash,
    compute_resolution_event_hash,
)
from open_agent_kit.features.team.activity.store.backup.models import ImportResult
from open_agent_kit.features.team.activity.store.backup.paths import (
    extract_machine_id_from_filename,
)
from open_agent_kit.features.team.activity.store.backup.sql_parser import (
    _extract_sql_statements,
    _parse_insert_statement,
    _parse_sql_values_as_strings,
    _remove_column_from_insert,
    _replace_column_value,
)
from open_agent_kit.features.team.activity.store.schema import SCHEMA_VERSION

if TYPE_CHECKING:
    from open_agent_kit.features.team.activity.store.core import ActivityStore

logger = logging.getLogger(__name__)


# ---- Table-to-attribute mappings for ImportResult counters ----

_TABLE_TO_IMPORTED_ATTR: dict[str, str] = {
    "sessions": "sessions_imported",
    "prompt_batches": "batches_imported",
    "memory_observations": "observations_imported",
    "activities": "activities_imported",
    "agent_schedules": "schedules_imported",
    "resolution_events": "resolution_events_imported",
    "governance_audit_events": "gov_audit_imported",
}

_TABLE_TO_SKIPPED_ATTR: dict[str, str] = {
    "sessions": "sessions_skipped",
    "prompt_batches": "batches_skipped",
    "memory_observations": "observations_skipped",
    "activities": "activities_skipped",
    "agent_schedules": "schedules_skipped",
    "resolution_events": "resolution_events_skipped",
    "governance_audit_events": "gov_audit_skipped",
}


def import_from_sql_with_dedup(
    store: ActivityStore,
    backup_path: Path,
    dry_run: bool = False,
    replace_machine: bool = False,
    vector_store: Any | None = None,
) -> ImportResult:
    """Import data from SQL backup with content-based deduplication.

    The database should already have the current schema (via _ensure_schema).
    This method imports data only, not schema. All imported observations
    are marked as unembedded to trigger ChromaDB rebuild.

    Uses content hashes to detect duplicates across machines:
    - Sessions: deduplicated by primary key (session ID)
    - Prompt batches: deduplicated by content_hash (session_id + prompt_number)
    - Observations: deduplicated by content_hash (observation + type + context)
    - Activities: deduplicated by content_hash (session_id + timestamp + tool_name)

    Foreign key handling:
    - prompt_batches: id column removed (auto-generated by SQLite)
    - activities: id column removed, prompt_batch_id remapped to new IDs

    When ``replace_machine=True``, all existing records from the backup's
    source machine are deleted before importing.  This prevents memory
    amplification when observations are regenerated with different text.

    Args:
        store: The ActivityStore instance.
        backup_path: Path to SQL backup file.
        dry_run: If True, preview what would be imported without making changes.
        replace_machine: If True, delete existing records from the backup's
            source machine before importing (drop-and-replace semantics).
        vector_store: Optional vector store for ChromaDB cleanup during replace.

    Returns:
        ImportResult with detailed statistics.
    """
    logger.info(f"Importing from backup: {backup_path} (dry_run={dry_run})")

    result = ImportResult()
    content = backup_path.read_text(encoding="utf-8")
    lines = content.split("\n")

    # Check backup schema version
    from open_agent_kit.features.team.activity.store.backup.exporter import (
        _parse_backup_schema_version,
    )

    backup_schema_version = _parse_backup_schema_version(lines)
    if backup_schema_version is not None and backup_schema_version != SCHEMA_VERSION:
        if backup_schema_version < SCHEMA_VERSION:
            logger.warning(
                f"Backup file is from older schema version {backup_schema_version} "
                f"(current: {SCHEMA_VERSION}). Import will proceed - missing columns "
                "will use default values and hashes will be computed from content."
            )
        else:
            logger.warning(
                f"Backup file is from newer schema version {backup_schema_version} "
                f"(current: {SCHEMA_VERSION}). Import may have issues if backup "
                "contains columns not in current schema."
            )

    # Drop-and-replace: delete existing records from this backup's source machine
    source_machine_id = extract_machine_id_from_filename(backup_path.name)
    if replace_machine and source_machine_id != store.machine_id and not dry_run:
        from open_agent_kit.features.team.activity.store.delete import (
            delete_records_by_machine,
        )

        delete_counts = delete_records_by_machine(store, source_machine_id, vector_store)
        result.sessions_deleted = delete_counts.get("sessions", 0)
        result.batches_deleted = delete_counts.get("prompt_batches", 0)
        result.observations_deleted = delete_counts.get("memory_observations", 0)
        result.activities_deleted = delete_counts.get("activities", 0)
        result.runs_deleted = delete_counts.get("agent_runs", 0)
        result.gov_audit_deleted = delete_counts.get("governance_audit_events", 0)
        logger.info(
            f"Replace mode: deleted {result.total_deleted} existing records "
            f"for machine {source_machine_id}"
        )

    # Load existing IDs/hashes for deduplication
    existing_session_ids = store.get_all_session_ids()
    existing_batch_hashes = store.get_all_prompt_batch_hashes()
    existing_obs_hashes = store.get_all_observation_hashes()
    existing_activity_hashes = store.get_all_activity_hashes()
    existing_schedule_names = store.get_all_schedule_task_names()
    existing_resolution_hashes = store.get_all_resolution_event_hashes()

    # Get current machine ID for schedule filtering
    current_machine_id = store.machine_id

    logger.debug(
        f"Existing records: {len(existing_session_ids)} sessions, "
        f"{len(existing_batch_hashes)} batches, {len(existing_obs_hashes)} observations, "
        f"{len(existing_activity_hashes)} activities, {len(existing_schedule_names)} schedules, "
        f"{len(existing_resolution_hashes)} resolution_events"
    )

    # Parse INSERT statements - use proper SQL statement extraction for multi-line values
    statements_by_table: dict[str, list[tuple[str, dict]]] = {
        "sessions": [],
        "prompt_batches": [],
        "memory_observations": [],
        "activities": [],
        "agent_schedules": [],
        "resolution_events": [],
        "governance_audit_events": [],
    }

    # Extract complete SQL statements (handles multi-line INSERT with newlines in values)
    sql_statements = _extract_sql_statements(content)
    logger.debug(f"Extracted {len(sql_statements)} SQL statements from backup")

    for stmt in sql_statements:
        # Extract table name
        table_match = re.match(r"INSERT INTO (\w+)", stmt)
        if not table_match:
            continue
        table = table_match.group(1)

        if table not in statements_by_table:
            continue

        # Parse column names and values
        parsed = _parse_insert_statement(stmt)
        if parsed:
            statements_by_table[table].append((stmt, parsed))

    # Process each table in order (sessions first due to foreign keys)
    conn = store._get_connection()

    # Get valid columns for each table (for schema compatibility filtering)
    table_columns: dict[str, set[str]] = {}
    for table in statements_by_table:
        table_columns[table] = _get_table_columns(conn, table)
    logger.debug(f"Schema columns loaded for tables: {list(table_columns.keys())}")

    # Track old prompt_batch_id -> (session_id, prompt_number) from backup file
    # This is needed to remap activities' prompt_batch_id foreign keys
    old_batch_id_to_key: dict[int, tuple[str, int]] = {}
    for _stmt, row_dict in statements_by_table["prompt_batches"]:
        old_id = row_dict.get("id")
        session_id = row_dict.get("session_id")
        prompt_number = row_dict.get("prompt_number")
        if old_id is not None and session_id and prompt_number is not None:
            old_batch_id_to_key[int(old_id)] = (str(session_id), int(prompt_number))

    logger.debug(f"Tracked {len(old_batch_id_to_key)} old prompt_batch_id -> key mappings")

    # Initialize the batch ID mapping (built after prompt_batches import)
    old_to_new_batch_id: dict[int, int] = {}

    # Track imported session IDs for parent validation
    imported_session_ids: set[str] = set()

    # Track columns removed during filtering (for summary logging)
    all_removed_columns: dict[str, set[str]] = {t: set() for t in statements_by_table}

    for table in [
        "sessions",
        "prompt_batches",
        "memory_observations",
        "activities",
        "agent_schedules",
        "resolution_events",
        "governance_audit_events",
    ]:
        # After importing sessions, validate parent_session_id references
        if table == "prompt_batches" and not dry_run and imported_session_ids:
            _validate_parent_session_ids(conn, imported_session_ids)

        # After importing prompt_batches, build the ID mapping and remap self-references
        if table == "memory_observations" and not dry_run:
            # Build mapping from (session_id, prompt_number) -> new_prompt_batch_id
            new_key_to_batch_id = _build_prompt_batch_id_map(conn)
            # Combine: old_batch_id -> key -> new_batch_id
            for old_id, key in old_batch_id_to_key.items():
                if key in new_key_to_batch_id:
                    old_to_new_batch_id[old_id] = new_key_to_batch_id[key]
            logger.debug(
                f"Built prompt_batch_id remap: {len(old_to_new_batch_id)} mappings "
                f"(from {len(old_batch_id_to_key)} old IDs)"
            )

            # Remap source_plan_batch_id self-references in prompt_batches
            _remap_source_plan_batch_id(conn, old_to_new_batch_id)

        # Governance audit events use delete-then-import (no hashing needed).
        # Delete existing events from the backup's source machine before importing.
        if table == "governance_audit_events" and source_machine_id and not dry_run:
            deleted = conn.execute(
                "DELETE FROM governance_audit_events WHERE source_machine_id = ?",
                (source_machine_id,),
            ).rowcount
            if deleted > 0:
                result.gov_audit_deleted = deleted
                logger.debug(
                    f"Cleared {deleted} governance audit events from machine {source_machine_id}"
                )

        for stmt, row_dict in statements_by_table[table]:
            try:
                # Filter columns for schema compatibility (handles newer backups)
                filtered_stmt, filtered_row_dict, removed_cols = _filter_columns_for_schema(
                    stmt, row_dict, table_columns[table]
                )
                all_removed_columns[table].update(removed_cols)

                should_skip, reason = _should_skip_record(
                    table,
                    filtered_row_dict,
                    existing_session_ids,
                    existing_batch_hashes,
                    existing_obs_hashes,
                    existing_activity_hashes,
                    existing_schedule_names,
                    current_machine_id,
                    existing_resolution_hashes,
                )

                if should_skip:
                    _increment_skipped(result, table)
                    logger.debug(f"Skipping {table} record: {reason}")
                    continue

                rows_affected = 1  # Default for dry_run
                if not dry_run:
                    # Modify statement for proper import
                    modified_stmt = _prepare_statement_for_import(filtered_stmt, table)

                    # For memory_observations and activities, remap prompt_batch_id to new ID
                    if table in ("memory_observations", "activities"):
                        modified_stmt = _remap_prompt_batch_id(
                            modified_stmt, filtered_row_dict, old_to_new_batch_id
                        )

                    cursor = conn.execute(modified_stmt)
                    rows_affected = cursor.rowcount

                    # Track imported session IDs for parent validation
                    if table == "sessions":
                        session_id = filtered_row_dict.get("id")
                        if session_id:
                            imported_session_ids.add(str(session_id))

                    # Update existing sets to avoid duplicates within same import
                    _update_existing_sets(
                        table,
                        filtered_row_dict,
                        existing_session_ids,
                        existing_batch_hashes,
                        existing_obs_hashes,
                        existing_activity_hashes,
                        existing_schedule_names,
                        existing_resolution_hashes,
                    )

                # Only count as imported if a row was actually inserted
                # (INSERT OR IGNORE returns rowcount=0 when ignored due to ID collision)
                if rows_affected > 0:
                    _increment_imported(result, table)
                else:
                    _increment_skipped(result, table)
                    logger.debug(f"Skipped {table} record due to ID collision")

            except sqlite3.Error as e:
                result.errors += 1
                error_msg = f"Error importing {table} record: {e}"
                result.error_messages.append(error_msg)
                logger.warning(error_msg)

    if not dry_run:
        conn.commit()

        # Backfill sessions.summary from imported session_summary observations.
        # Old backups store summaries as memory_observations with memory_type='session_summary'.
        # New schema stores them directly in sessions.summary. Migrate on import.
        _backfill_session_summaries_from_observations(conn)

    # Log summary of filtered columns (schema compatibility)
    for table, removed in all_removed_columns.items():
        if removed:
            logger.info(
                f"Schema compatibility: filtered {len(removed)} unknown columns "
                f"from {table}: {sorted(removed)}"
            )

    logger.info(
        f"Import complete: {result.total_imported} imported, "
        f"{result.total_skipped} skipped (duplicates), {result.errors} errors"
    )
    return result


# ---- Helper functions ----


def _should_skip_record(
    table: str,
    row_dict: dict,
    existing_session_ids: set[str],
    existing_batch_hashes: set[str],
    existing_obs_hashes: set[str],
    existing_activity_hashes: set[str],
    existing_schedule_names: set[str] | None = None,
    current_machine_id: str | None = None,
    existing_resolution_hashes: set[str] | None = None,
) -> tuple[bool, str]:
    """Determine if a record should be skipped due to duplication.

    Args:
        table: Table name.
        row_dict: Parsed row data.
        existing_*: Sets of existing IDs/hashes.
        existing_schedule_names: Set of existing schedule instance names.
        current_machine_id: Current machine identifier (for schedule filtering).

    Returns:
        Tuple of (should_skip, reason).
    """
    if table == "sessions":
        session_id = row_dict.get("id", "")
        if session_id in existing_session_ids:
            return True, f"session {session_id} already exists"

    elif table == "prompt_batches":
        # Check content hash first, then compute if missing
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            session_id = str(row_dict.get("session_id", ""))
            prompt_number = int(row_dict.get("prompt_number", 0))
            content_hash = compute_prompt_batch_hash(session_id, prompt_number)
        if content_hash in existing_batch_hashes:
            return True, f"prompt_batch with hash {content_hash} already exists"

    elif table == "memory_observations":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            observation = str(row_dict.get("observation", ""))
            memory_type = str(row_dict.get("memory_type", ""))
            context = row_dict.get("context")
            content_hash = compute_observation_hash(observation, memory_type, context)
        if content_hash in existing_obs_hashes:
            return True, f"observation with hash {content_hash} already exists"

    elif table == "activities":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            session_id = str(row_dict.get("session_id", ""))
            timestamp_epoch = int(row_dict.get("timestamp_epoch", 0))
            tool_name = str(row_dict.get("tool_name", ""))
            content_hash = compute_activity_hash(session_id, timestamp_epoch, tool_name)
        if content_hash in existing_activity_hashes:
            return True, f"activity with hash {content_hash} already exists"

    elif table == "agent_schedules":
        # Schedules are user preferences - only import from same machine
        # This prevents team backups from overwriting personal schedule settings
        backup_machine_id = row_dict.get("source_machine_id")
        if current_machine_id and backup_machine_id and backup_machine_id != current_machine_id:
            return True, f"schedule from different machine ({backup_machine_id})"

        # Check if schedule already exists by task_name (primary key)
        task_name = row_dict.get("task_name", "")
        if existing_schedule_names and task_name in existing_schedule_names:
            return True, f"schedule {task_name} already exists"

    elif table == "resolution_events":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            observation_id = str(row_dict.get("observation_id", ""))
            action = str(row_dict.get("action", ""))
            source_machine_id = str(row_dict.get("source_machine_id", ""))
            superseded_by = str(row_dict.get("superseded_by", ""))
            content_hash = compute_resolution_event_hash(
                observation_id, action, source_machine_id, superseded_by
            )
        if existing_resolution_hashes and content_hash in existing_resolution_hashes:
            return True, f"resolution_event with hash {content_hash} already exists"

    return False, ""


def _update_existing_sets(
    table: str,
    row_dict: dict,
    existing_session_ids: set[str],
    existing_batch_hashes: set[str],
    existing_obs_hashes: set[str],
    existing_activity_hashes: set[str],
    existing_schedule_names: set[str] | None = None,
    existing_resolution_hashes: set[str] | None = None,
) -> None:
    """Update existing sets after successful import to prevent duplicates within batch.

    Args:
        table: Table name.
        row_dict: Imported row data.
        existing_*: Sets to update.
    """
    if table == "sessions":
        existing_session_ids.add(str(row_dict.get("id", "")))

    elif table == "prompt_batches":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            session_id = str(row_dict.get("session_id", ""))
            prompt_number = int(row_dict.get("prompt_number", 0))
            content_hash = compute_prompt_batch_hash(session_id, prompt_number)
        existing_batch_hashes.add(content_hash)

    elif table == "memory_observations":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            observation = str(row_dict.get("observation", ""))
            memory_type = str(row_dict.get("memory_type", ""))
            context = row_dict.get("context")
            content_hash = compute_observation_hash(observation, memory_type, context)
        existing_obs_hashes.add(content_hash)

    elif table == "activities":
        content_hash = row_dict.get("content_hash")
        if not content_hash:
            session_id = str(row_dict.get("session_id", ""))
            timestamp_epoch = int(row_dict.get("timestamp_epoch", 0))
            tool_name = str(row_dict.get("tool_name", ""))
            content_hash = compute_activity_hash(session_id, timestamp_epoch, tool_name)
        existing_activity_hashes.add(content_hash)

    elif table == "agent_schedules":
        if existing_schedule_names is not None:
            task_name = str(row_dict.get("task_name", ""))
            existing_schedule_names.add(task_name)

    elif table == "resolution_events":
        if existing_resolution_hashes is not None:
            content_hash = row_dict.get("content_hash")
            if not content_hash:
                observation_id = str(row_dict.get("observation_id", ""))
                action = str(row_dict.get("action", ""))
                source_machine_id = str(row_dict.get("source_machine_id", ""))
                superseded_by = str(row_dict.get("superseded_by", ""))
                content_hash = compute_resolution_event_hash(
                    observation_id, action, source_machine_id, superseded_by
                )
            existing_resolution_hashes.add(content_hash)


def _prepare_statement_for_import(stmt: str, table: str) -> str:
    """Modify INSERT statement for proper import.

    For memory_observations, marks embedded=0 to trigger ChromaDB rebuild.
    For prompt_batches, marks plan_embedded=0 for re-indexing and removes id column.
    For activities, removes id column to avoid PRIMARY KEY conflicts.
    For agent_schedules, uses INSERT OR IGNORE for TEXT primary key handling.

    The id column is removed for prompt_batches and activities because these use
    auto-increment INTEGER PRIMARY KEY, which would conflict when importing from
    different machines that have overlapping id sequences.

    Args:
        stmt: Original INSERT statement.
        table: Table name.

    Returns:
        Modified statement ready for execution.
    """
    if table == "memory_observations":
        # Use INSERT OR IGNORE to handle rare UUID collisions gracefully
        # (if ID exists but content_hash is different, skip the import)
        stmt = stmt.replace("INSERT INTO", "INSERT OR IGNORE INTO", 1)
        return _replace_column_value(stmt, "embedded", "0")
    elif table == "prompt_batches":
        # Remove id column to let SQLite auto-generate, and mark unembedded
        stmt = _remove_column_from_insert(stmt, "id")
        stmt = _replace_column_value(stmt, "plan_embedded", "0")
        return stmt
    elif table == "activities":
        # Remove id column to let SQLite auto-generate
        return _remove_column_from_insert(stmt, "id")
    elif table == "sessions":
        # Use INSERT OR IGNORE to handle potential session ID collisions
        return stmt.replace("INSERT INTO", "INSERT OR IGNORE INTO", 1)
    elif table == "agent_schedules":
        # Use INSERT OR IGNORE for TEXT primary key (task_name)
        # Schedules are already filtered by machine in _should_skip_record
        return stmt.replace("INSERT INTO", "INSERT OR IGNORE INTO", 1)
    elif table == "resolution_events":
        # Use INSERT OR IGNORE and mark as unapplied (needs replay)
        stmt = stmt.replace("INSERT INTO", "INSERT OR IGNORE INTO", 1)
        return _replace_column_value(stmt, "applied", "0")
    elif table == "governance_audit_events":
        # Remove auto-increment id to let SQLite generate new IDs
        return _remove_column_from_insert(stmt, "id")
    return stmt


def _build_prompt_batch_id_map(conn: sqlite3.Connection) -> dict[tuple[str, int], int]:
    """Build mapping from (session_id, prompt_number) to prompt_batch id.

    Used to remap activities' prompt_batch_id foreign keys after importing
    prompt_batches with auto-generated IDs.

    Args:
        conn: SQLite connection.

    Returns:
        Dictionary mapping (session_id, prompt_number) tuples to prompt_batch id.
    """
    cursor = conn.execute("SELECT id, session_id, prompt_number FROM prompt_batches")
    return {(row[1], row[2]): row[0] for row in cursor.fetchall()}


def _remap_prompt_batch_id(
    stmt: str,
    row_dict: dict,
    old_to_new_batch_id: dict[int, int],
) -> str:
    """Remap prompt_batch_id to the correct new ID.

    After prompt_batches are imported with auto-generated IDs, activities
    and memory_observations still reference the OLD prompt_batch_id from
    the source machine. This function looks up the correct new ID using
    the pre-computed mapping.

    Args:
        stmt: The INSERT statement for the record.
        row_dict: Parsed row data from the record.
        old_to_new_batch_id: Mapping from old prompt_batch_id to new id.

    Returns:
        Modified statement with corrected prompt_batch_id, or original if lookup fails.
    """
    old_batch_id = row_dict.get("prompt_batch_id")
    if old_batch_id is None:
        return stmt

    old_batch_id_int = int(old_batch_id)
    new_batch_id = old_to_new_batch_id.get(old_batch_id_int)

    if new_batch_id is None:
        logger.warning(
            f"No mapping found for old prompt_batch_id {old_batch_id}, "
            f"activity may reference wrong batch"
        )
        return stmt

    logger.debug(f"Remapped activity prompt_batch_id: {old_batch_id} -> {new_batch_id}")
    return _replace_column_value(stmt, "prompt_batch_id", str(new_batch_id))


def _increment_imported(result: ImportResult, table: str) -> None:
    """Increment the appropriate imported counter."""
    attr = _TABLE_TO_IMPORTED_ATTR.get(table)
    if attr:
        setattr(result, attr, getattr(result, attr) + 1)


def _increment_skipped(result: ImportResult, table: str) -> None:
    """Increment the appropriate skipped counter."""
    attr = _TABLE_TO_SKIPPED_ATTR.get(table)
    if attr:
        setattr(result, attr, getattr(result, attr) + 1)


def _get_table_columns(conn: sqlite3.Connection, table: str) -> set[str]:
    """Get column names for a table from the current schema.

    Args:
        conn: SQLite connection.
        table: Table name.

    Returns:
        Set of column names in the table.
    """
    cursor = conn.execute(f"PRAGMA table_info({table})")  # noqa: S608 - trusted table name
    return {row[1] for row in cursor.fetchall()}


def _filter_columns_for_schema(
    stmt: str,
    row_dict: dict,
    valid_columns: set[str],
) -> tuple[str, dict, list[str]]:
    """Filter INSERT statement to only include columns in current schema.

    This enables importing backups from newer schema versions by stripping
    columns that don't exist in the current schema.

    Args:
        stmt: Original INSERT statement.
        row_dict: Parsed row data.
        valid_columns: Set of valid column names from current schema.

    Returns:
        Tuple of (filtered_stmt, filtered_row_dict, removed_columns).
    """
    # Parse the statement to get columns and values
    cols_match = re.search(r"INSERT INTO (\w+) \(([^)]+)\) VALUES \(", stmt, re.IGNORECASE)
    if not cols_match:
        return stmt, row_dict, []

    table_name = cols_match.group(1)
    cols_str = cols_match.group(2)
    columns = [c.strip() for c in cols_str.split(",")]

    # Find columns to remove
    removed_columns = [c for c in columns if c not in valid_columns]
    if not removed_columns:
        return stmt, row_dict, []

    # Parse values
    values_start = stmt.upper().find("VALUES")
    paren_start = stmt.find("(", values_start)
    values_section = stmt[paren_start + 1 :].rstrip(");")
    values = _parse_sql_values_as_strings(values_section)

    if len(values) != len(columns):
        logger.warning(f"Column/value count mismatch, skipping column filter for {table_name}")
        return stmt, row_dict, []

    # Build filtered columns and values
    filtered_columns = []
    filtered_values = []
    filtered_row_dict = {}

    for col, val in zip(columns, values, strict=False):
        if col in valid_columns:
            filtered_columns.append(col)
            filtered_values.append(val)
            if col in row_dict:
                filtered_row_dict[col] = row_dict[col]

    # Rebuild statement
    filtered_stmt = (
        f"INSERT INTO {table_name} ({', '.join(filtered_columns)}) "
        f"VALUES ({', '.join(filtered_values)});"
    )

    logger.debug(
        f"Filtered {len(removed_columns)} unknown columns from {table_name}: {removed_columns}"
    )
    return filtered_stmt, filtered_row_dict, removed_columns


def _remap_source_plan_batch_id(
    conn: sqlite3.Connection,
    old_to_new_batch_id: dict[int, int],
) -> int:
    """Remap source_plan_batch_id FKs in prompt_batches after import.

    Since prompt_batches.id is auto-generated during import, any
    source_plan_batch_id values reference old IDs that no longer exist.
    This function updates them to point to the correct new IDs.

    Args:
        conn: SQLite connection.
        old_to_new_batch_id: Mapping from old prompt_batch_id to new id.

    Returns:
        Number of records updated.
    """
    if not old_to_new_batch_id:
        return 0

    # Find all prompt_batches with source_plan_batch_id set
    cursor = conn.execute(
        "SELECT id, source_plan_batch_id FROM prompt_batches WHERE source_plan_batch_id IS NOT NULL"
    )
    rows = cursor.fetchall()

    updated = 0
    for batch_id, old_source_id in rows:
        new_source_id = old_to_new_batch_id.get(old_source_id)
        if new_source_id is not None and new_source_id != old_source_id:
            conn.execute(
                "UPDATE prompt_batches SET source_plan_batch_id = ? WHERE id = ?",
                (new_source_id, batch_id),
            )
            updated += 1
            logger.debug(
                f"Remapped prompt_batch {batch_id} source_plan_batch_id: "
                f"{old_source_id} -> {new_source_id}"
            )
        elif new_source_id is None and old_source_id not in old_to_new_batch_id:
            # Old source batch wasn't imported - set to NULL to avoid FK violation
            conn.execute(
                "UPDATE prompt_batches SET source_plan_batch_id = NULL WHERE id = ?",
                (batch_id,),
            )
            logger.warning(
                f"prompt_batch {batch_id} references unknown source_plan_batch_id "
                f"{old_source_id}, setting to NULL"
            )
            updated += 1

    if updated > 0:
        logger.info(f"Remapped {updated} source_plan_batch_id references")
    return updated


def _validate_parent_session_ids(
    conn: sqlite3.Connection,
    imported_session_ids: set[str],
) -> int:
    """Validate parent_session_id references after session import.

    Sets parent_session_id to NULL for sessions that reference
    parent sessions not included in the import or existing data.

    Args:
        conn: SQLite connection.
        imported_session_ids: Set of session IDs that were imported.

    Returns:
        Number of orphaned references fixed.
    """
    # Get all existing session IDs (includes both imported and pre-existing)
    cursor = conn.execute("SELECT id FROM sessions")
    all_session_ids = {row[0] for row in cursor.fetchall()}

    # Find sessions with orphaned parent references
    cursor = conn.execute(
        "SELECT id, parent_session_id FROM sessions WHERE parent_session_id IS NOT NULL"
    )
    rows = cursor.fetchall()

    fixed = 0
    for session_id, parent_id in rows:
        if parent_id not in all_session_ids:
            conn.execute(
                "UPDATE sessions SET parent_session_id = NULL, parent_session_reason = NULL "
                "WHERE id = ?",
                (session_id,),
            )
            logger.warning(
                f"Session {session_id} references non-existent parent {parent_id}, "
                f"unlinking (parent may not have been included in backup)"
            )
            fixed += 1

    if fixed > 0:
        logger.info(f"Fixed {fixed} orphaned parent_session_id references")
    return fixed


def _backfill_session_summaries_from_observations(conn: sqlite3.Connection) -> int:
    """Backfill sessions.summary from imported session_summary observations.

    Old backups store summaries as memory_observations with memory_type='session_summary'.
    New schema stores them directly in sessions.summary column. This migrates
    the data on import: copies the most recent session_summary observation into
    sessions.summary/summary_updated_at, then removes the migrated observations.

    Args:
        conn: SQLite connection (already committed after import).

    Returns:
        Number of sessions backfilled.
    """
    cursor = conn.execute("""
        UPDATE sessions SET
          summary = (SELECT observation FROM memory_observations
                     WHERE memory_observations.session_id = sessions.id
                     AND memory_observations.memory_type = 'session_summary'
                     ORDER BY created_at_epoch DESC LIMIT 1),
          summary_updated_at = (SELECT created_at_epoch FROM memory_observations
                     WHERE memory_observations.session_id = sessions.id
                     AND memory_observations.memory_type = 'session_summary'
                     ORDER BY created_at_epoch DESC LIMIT 1)
        WHERE summary IS NULL AND EXISTS (
          SELECT 1 FROM memory_observations
          WHERE memory_observations.session_id = sessions.id
          AND memory_observations.memory_type = 'session_summary'
        )
        """)
    backfilled = cursor.rowcount

    if backfilled > 0:
        conn.execute("DELETE FROM memory_observations WHERE memory_type = 'session_summary'")
        conn.commit()
        logger.info(
            f"Backfilled {backfilled} session summaries from legacy session_summary observations"
        )

    return backfilled
