#!/usr/bin/env python3
"""
chatlog_embed_sweeper.py — lazy embed chat log rows missing embeddings.

Runs on a schedule (default every 30 min via install_schedules.py). Picks up
rows written with embed=False, embeds in batches using memory_core._embed_many,
and drains any spill-to-disk files from the async write queue.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sqlite3
import sys
import uuid
from datetime import datetime, timezone
from glob import glob

# Setup path so we can import bin/ modules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import chatlog_config

logger = logging.getLogger("chatlog_embed_sweeper")


async def drain_spill(conn: sqlite3.Connection) -> int:
    """
    Drain spill JSONL files back into the chat log DB(s).

    Spilled items written after the DB-parameter refactor carry the target
    path they were enqueued against (``_db_path``). We honor that per-item
    so rows from a session that pointed at a dedicated test/benchmark DB
    don't silently land in the live store. Legacy spill items without
    ``_db_path`` fall back to ``conn`` (the sweeper's connection).

    Returns count of rows inserted across all target DBs.
    """
    spill_dir = chatlog_config.SPILL_DIR
    if not os.path.exists(spill_dir):
        return 0

    spill_files = glob(os.path.join(spill_dir, "*.jsonl"))
    if not spill_files:
        return 0

    insert_sql = (
        "INSERT OR IGNORE INTO memory_items "
        "(id, content, title, metadata_json, type, agent_id, conversation_id, is_deleted, created_at, variant) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
    )

    total_drained = 0
    for spill_path in spill_files:
        try:
            # Group rows by their captured target DB. Legacy spills (no
            # _db_path) land on the sweeper's connection under the None key.
            per_db: dict = {}
            with open(spill_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        doc = json.loads(line)
                    except json.JSONDecodeError as e:
                        logger.warning("Skipping malformed JSON in %s: %s", spill_path, e)
                        continue

                    memory_id = doc.get("_id") or str(uuid.uuid4())
                    content = doc.get("_content", doc.get("content", ""))
                    title = doc.get("_title", "")
                    conversation_id = doc.get("conversation_id", "")
                    host_agent = doc.get("host_agent", "unknown")
                    variant = doc.get("variant")
                    timestamp = (
                        doc.get("_created_at")
                        or doc.get("timestamp")
                        or datetime.now(timezone.utc).isoformat()
                    )

                    if doc.get("_metadata_json"):
                        metadata_json = doc["_metadata_json"]
                    else:
                        metadata = {
                            "role": doc.get("role", "unknown"),
                            "provider": doc.get("provider", "unknown"),
                            "model_id": doc.get("model_id", "unknown"),
                            "host_agent": host_agent,
                            "spill_source": True,
                        }
                        metadata_json = json.dumps(metadata)

                    agent_id = doc.get("agent_id") or f"{host_agent}:spill"

                    row = (
                        memory_id, content, title, metadata_json,
                        "chat_log", agent_id, conversation_id, 0,
                        timestamp, variant,
                    )
                    per_db.setdefault(doc.get("_db_path"), []).append(row)

            file_drained = 0
            had_error = False
            for db_path, rows in per_db.items():
                if not rows:
                    continue
                # Open a fresh short-lived connection for any path that
                # differs from the sweeper's conn. This is the drain path —
                # not hot — so the overhead is acceptable.
                target_conn = conn
                own_conn = False
                if db_path and os.path.abspath(db_path) != os.path.abspath(getattr(conn, "_m3_db_path", "")):
                    try:
                        target_conn = sqlite3.connect(db_path, timeout=10)
                        own_conn = True
                    except sqlite3.Error as e:
                        logger.error("Cannot open spill target DB %s: %s", db_path, e)
                        had_error = True
                        continue
                try:
                    target_conn.executemany(insert_sql, rows)
                    target_conn.commit()
                    file_drained += len(rows)
                    logger.info("Drained %d rows from %s -> %s", len(rows), spill_path, db_path or "<sweeper-conn>")
                except sqlite3.Error as e:
                    logger.error("Failed to insert drained rows from %s into %s: %s", spill_path, db_path, e)
                    had_error = True
                finally:
                    if own_conn:
                        target_conn.close()

            total_drained += file_drained

            # Delete the spill file only if every group landed successfully.
            if not had_error:
                try:
                    os.remove(spill_path)
                    logger.info("Deleted spill file: %s", spill_path)
                except OSError as e:
                    logger.warning("Failed to delete spill file %s: %s", spill_path, e)

        except Exception as e:
            logger.error("Error processing spill file %s: %s", spill_path, e)

    return total_drained


async def embed_batch(
    conn: sqlite3.Connection,
    batch: list[tuple[str, str, str, str]],
    dry_run: bool = False,
) -> int:
    """
    Embed a batch of rows. Returns count embedded.
    batch format: [(id, content, title, metadata_json), ...]

    Kept for backwards compatibility. The main() loop now drives
    embed_sweep_lib.run_embed_loop directly; this function remains for
    any external caller that imported it. Internally still uses the
    same memory_core._embed_many primitive — bit-for-bit equivalent
    to the pre-extraction path.
    """
    if not batch:
        return 0

    texts = [content for _, content, _, _ in batch]

    # Import embedding function lazily
    from embedding_utils import pack as _pack
    from memory_core import _embed_many as embed_many

    try:
        embeddings = await embed_many(texts)
    except Exception as e:
        logger.error("Failed to embed batch: %s", e)
        return 0

    if not dry_run:
        try:
            rows_to_insert = []
            for (mem_id, _, _, _), (vec, model_str) in zip(batch, embeddings):
                if vec is not None:
                    packed = _pack(vec)
                    embed_id = str(uuid.uuid4())
                    rows_to_insert.append((
                        embed_id,
                        mem_id,
                        packed,
                        model_str,
                        len(vec),  # dim
                        datetime.now(timezone.utc).isoformat(),
                    ))

            if rows_to_insert:
                conn.executemany(
                    """
                    INSERT OR IGNORE INTO memory_embeddings
                    (id, memory_id, embedding, embed_model, dim, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                    """,
                    rows_to_insert,
                )
                conn.commit()
        except sqlite3.Error as e:
            logger.error("Failed to insert embeddings: %s", e)
            return 0

    return len(batch)


def table_exists(conn: sqlite3.Connection, table_name: str) -> bool:
    """Check if a table exists in the DB."""
    try:
        row = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
            (table_name,),
        ).fetchone()
        return row is not None
    except sqlite3.Error:
        return False


async def get_unembed_count(conn: sqlite3.Connection) -> int:
    """Count rows needing embeddings."""
    if not table_exists(conn, "memory_items") or not table_exists(conn, "memory_embeddings"):
        return 0

    try:
        row = conn.execute(
            """
            SELECT COUNT(*) as cnt
            FROM memory_items
            WHERE type='chat_log'
              AND is_deleted=0
              AND id NOT IN (SELECT memory_id FROM memory_embeddings)
            """
        ).fetchone()
        return row["cnt"] if row else 0
    except sqlite3.Error as e:
        logger.error("Failed to count unembedded rows: %s", e)
        return 0


def load_state() -> dict:
    """Load state file, or return defaults."""
    state_path = chatlog_config.STATE_FILE
    if not os.path.exists(state_path):
        return {
            "embed_backlog": 0,
            "last_sweeper_run_at": None,
            "last_sweeper_rows_embedded": 0,
            "last_sweeper_spill_drained": 0,
            "queue_depth": 0,
        }
    try:
        with open(state_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        logger.warning("Failed to load state file: %s", e)
        return {
            "embed_backlog": 0,
            "last_sweeper_run_at": None,
            "last_sweeper_rows_embedded": 0,
            "last_sweeper_spill_drained": 0,
            "queue_depth": 0,
        }


def save_state(state: dict) -> None:
    """Save state atomically with rename."""
    state_dir = os.path.dirname(chatlog_config.STATE_FILE)
    os.makedirs(state_dir, exist_ok=True)
    tmp_path = chatlog_config.STATE_FILE + ".tmp"
    try:
        with open(tmp_path, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)
        os.replace(tmp_path, chatlog_config.STATE_FILE)
    except OSError as e:
        logger.error("Failed to save state: %s", e)


async def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Lazy embed chat log rows missing embeddings."
    )
    parser.add_argument(
        "--batch",
        type=int,
        default=None,
        help="Batch size (default from config.embed_sweeper.batch_size)",
    )
    parser.add_argument(
        "--max-per-run",
        type=int,
        default=None,
        help="Max rows per run (default from CHATLOG_EMBED_MAX_PER_RUN env or 10000)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Query and log but don't embed",
    )
    parser.add_argument(
        "--drain-spill",
        action="store_true",
        help="Process spill files before embedding",
    )
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from m3_sdk import add_database_arg
    from _task_runtime import add_log_file_arg, setup_task_runtime
    add_log_file_arg(parser)
    add_database_arg(parser)
    args = parser.parse_args()

    setup_task_runtime(args.log_file, lock_name="chatlog_embed_sweeper")

    if args.database:
        # M3_DATABASE unifies main+chatlog unless CHATLOG_DB_PATH overrides.
        os.environ["M3_DATABASE"] = args.database
        chatlog_config.invalidate_cache()

    # Resolve config
    cfg = chatlog_config.resolve_config()
    batch_size = args.batch or cfg.embed_sweeper.batch_size
    max_per_run = (
        args.max_per_run
        or int(os.environ.get("CHATLOG_EMBED_MAX_PER_RUN", "10000"))
    )

    # Open DB connection
    db_path = chatlog_config.chatlog_db_path()
    if not os.path.exists(db_path):
        logger.info("Chat log DB does not exist yet: %s", db_path)
        return 0

    try:
        conn = sqlite3.connect(db_path, timeout=10)
        conn.row_factory = sqlite3.Row
    except sqlite3.Error as e:
        logger.error("Failed to open chat log DB: %s", e)
        return 1

    start_time = asyncio.get_event_loop().time()

    try:
        spill_drained = 0
        rows_embedded = 0
        batches_processed = 0

        # Check if schema exists
        if not table_exists(conn, "memory_items"):
            logger.info("Chat log schema not initialized yet; nothing to do")
            return 0

        # Drain spill if requested or if files exist
        if args.drain_spill or os.path.exists(chatlog_config.SPILL_DIR):
            spill_drained = await drain_spill(conn)
            if spill_drained > 0:
                logger.info("Drained %d rows from spill", spill_drained)

        # Query and embed batches via the shared embed_sweep_lib loop.
        # This delegates concurrency, batching, cursor advance, and per-
        # batch hardening (timeout, oversize/empty/bad-dim skip,
        # consecutive-fail abort) to one place that bin/embed_backfill.py
        # also uses. Behavior change vs. pre-extraction: ORDER BY id ASC
        # instead of ORDER BY created_at ASC. UUIDs aren't time-ordered,
        # so this changes within-batch ordering (negligible at sweeper
        # cadence) and gains infinite-loop protection on skipped rows
        # via the after_id cursor.
        from embed_sweep_lib import Counters as _Counters
        from embed_sweep_lib import run_embed_loop
        from memory_core import (
            _content_hash as _ch,
        )
        from memory_core import (
            _embed_many as _em,
        )
        from memory_core import (
            _pack as _mc_pack,
        )

        counters = _Counters()

        # Fetch callback: pulls (after_id, limit) -> rows. Only chat_log
        # rows still missing an embedding row.
        def _fetch(after_id, limit):
            where = [
                "type='chat_log'",
                "is_deleted=0",
                "id NOT IN (SELECT memory_id FROM memory_embeddings)",
            ]
            params: list = []
            if after_id is not None:
                where.append("id > ?")
                params.append(after_id)
            sql = (
                f"SELECT id, content, title, metadata_json "
                f"FROM memory_items WHERE {' AND '.join(where)} "
                f"ORDER BY id LIMIT ?"
            )
            params.append(limit)
            try:
                return conn.execute(sql, params).fetchall()
            except sqlite3.Error as e:
                logger.error("Failed to query unembedded rows: %s", e)
                return []

        # Write callback: persists one embedding row. Chatlog DB has the
        # content_hash column (added in 003_chroma_sync_queue_align /
        # main migration 021's column add), so we write it.
        def _write(mid: str, vec: list[float], model_str: str, content_hash: str) -> bool:
            if args.dry_run:
                return True  # count it but don't write
            try:
                cur = conn.execute(
                    """
                    INSERT OR IGNORE INTO memory_embeddings
                    (id, memory_id, embedding, embed_model, dim, created_at, content_hash)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        str(uuid.uuid4()),
                        mid,
                        _mc_pack(vec),
                        model_str,
                        len(vec),
                        datetime.now(timezone.utc).isoformat(),
                        content_hash,
                    ),
                )
                conn.commit()
                return cur.rowcount > 0
            except sqlite3.Error as e:
                logger.error("Failed to insert embedding for %s: %s", mid[:8], e)
                return False

        # Drive the loop. limit=max_per_run preserves the existing
        # "max rows per scheduled run" budget.
        await run_embed_loop(
            fetch_candidates=_fetch,
            write_embedding=_write,
            counters=counters,
            embed_many=_em,
            content_hash_fn=_ch,
            transform_text=lambda t, _m: t,  # chatlog: no anchor augmentation
            batch_size=batch_size,
            concurrency=1,                   # sweeper has historically run sequentially
            timeout_s=300.0,                 # generous for a scheduled background job
            deadline_s=None,                 # max_per_run handles the budget
            max_consecutive_fails=5,
            max_row_bytes=32_768,
            expected_dim=None,               # don't reject by dim — chatlog has heterogenous models
            limit=max_per_run,
            log=lambda msg: logger.info("%s", msg),
        )
        rows_embedded = counters.embedded
        batches_processed = counters.batches_completed

        # Get remaining backlog count
        backlog = await get_unembed_count(conn)

        # Update state
        state = load_state()
        state["embed_backlog"] = backlog
        state["last_sweeper_run_at"] = datetime.now(timezone.utc).isoformat()
        state["last_sweeper_rows_embedded"] = rows_embedded
        state["last_sweeper_spill_drained"] = spill_drained
        save_state(state)

        elapsed = asyncio.get_event_loop().time() - start_time

        if rows_embedded > 0 or spill_drained > 0:
            logger.info(
                "Embedded %d rows in %d batches (%.1fs), spill drained: %d, backlog remaining: %d",
                rows_embedded,
                batches_processed,
                elapsed,
                spill_drained,
                backlog,
            )
        else:
            logger.info("Nothing to do (backlog: %d)", backlog)

        return 0

    except Exception as e:
        logger.exception("Sweeper failed: %s", e)
        return 1
    finally:
        conn.close()


if __name__ == "__main__":
    # Logging is configured inside main() via setup_task_runtime once args
    # are parsed. A minimal fallback keeps logger output visible if main()
    # raises before that point.
    logging.basicConfig(
        level=logging.INFO,
        format="%(name)s: [%(levelname)s] %(message)s",
    )
    sys.exit(asyncio.run(main()))
