#!/usr/bin/env python3
"""Backfill Telegram origin metadata for durable API conversations."""

from __future__ import annotations

import argparse
import json
import sqlite3
from pathlib import Path


def telegram_sessions(response_store: Path) -> dict[str, str]:
    """Return session ID to stable Telegram conversation-key mappings."""
    if not response_store.exists():
        return {}
    connection = sqlite3.connect(f"file:{response_store}?mode=ro", uri=True)
    try:
        rows = connection.execute(
            """SELECT c.name, r.data
                 FROM conversations AS c
                 JOIN responses AS r ON r.response_id = c.response_id
                WHERE c.name = 'telegram' OR c.name LIKE 'telegram-topic-%'"""
        ).fetchall()
    finally:
        connection.close()
    result: dict[str, str] = {}
    for conversation, raw in rows:
        try:
            session_id = str(json.loads(raw).get("session_id") or "").strip()
        except (TypeError, ValueError, json.JSONDecodeError):
            continue
        if session_id:
            result[session_id] = str(conversation)
    return result


def migrate(state_database: Path, response_store: Path) -> int:
    """Annotate known sessions without creating or rewriting conversations."""
    mappings = telegram_sessions(response_store)
    if not mappings or not state_database.exists():
        return 0
    connection = sqlite3.connect(state_database)
    changed = 0
    try:
        columns = {
            str(row[1])
            for row in connection.execute("PRAGMA table_info(sessions)").fetchall()
        }
        required = {
            "id", "source", "session_key", "chat_type", "display_name",
            "origin_json", "title",
        }
        if not required.issubset(columns):
            return 0
        for session_id, conversation in mappings.items():
            origin = json.dumps(
                {"platform": "telegram", "session_key": conversation},
                separators=(",", ":"),
            )
            title = "Telegram · General" if conversation == "telegram" else "Telegram"
            cursor = connection.execute(
                """UPDATE sessions
                      SET session_key = COALESCE(session_key, ?),
                          chat_type = COALESCE(chat_type, 'private'),
                          display_name = COALESCE(display_name, 'Telegram'),
                          origin_json = COALESCE(origin_json, ?),
                          title = COALESCE(title, ?)
                    WHERE id = ? AND source = 'api_server'
                      AND (
                          session_key IS NULL OR chat_type IS NULL OR
                          display_name IS NULL OR origin_json IS NULL OR title IS NULL
                      )""",
                (conversation, origin, title, session_id),
            )
            changed += max(cursor.rowcount, 0)
        connection.commit()
    finally:
        connection.close()
    return changed


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("state_database", type=Path)
    parser.add_argument("response_store", type=Path)
    args = parser.parse_args()
    migrate(args.state_database, args.response_store)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())