#!/usr/bin/env python3
"""Bound Telegram context and backfill durable API-session metadata."""

from __future__ import annotations

import argparse
import json
import re
import sqlite3
from pathlib import Path
from typing import Any


MAX_RESPONSE_BYTES = 32 << 20
MAX_HISTORY_ITEMS = 80
MAX_HISTORY_BYTES = 192_000
MAX_RECENT_TURNS = 12
MAX_ITEM_CHARS = 8_000
MAX_SUMMARY_CHARS = 12_000
SUMMARY_PREFIX = "Telegram continuity summary (older bounded context):"
ATTACHMENT_NOTE = (
    "[Earlier Telegram image omitted from bounded context; attach it again "
    "in Hermes WebUI if the original pixels are required.]"
)
_CONVERSATION = re.compile(r"^telegram(?:-topic-[a-z0-9-]{1,160})?$")
_MEDIA_PATH = re.compile(
    r'(?:MEDIA:\s*)?(?:/opt/data/(?:cache/images|workspace)|/workspace)/[^\s\]\[)}`"]+',
    re.IGNORECASE,
)
_DATA_IMAGE = re.compile(r"data:image/[^;,\s]+;base64,[A-Za-z0-9+/=]+", re.IGNORECASE)
_MESSAGE_IDENTITY_COLUMNS = (
    "session_id",
    "role",
    "content",
    "tool_call_id",
    "tool_calls",
    "tool_name",
    "timestamp",
    "token_count",
    "finish_reason",
    "reasoning",
    "reasoning_content",
    "reasoning_details",
    "codex_reasoning_items",
    "codex_message_items",
    "platform_message_id",
    "observed",
    "active",
    "compacted",
)


def _text(value: Any) -> str:
    """Return bounded context text without reusable transport-only images."""
    if isinstance(value, str):
        text = value
    elif isinstance(value, list):
        parts: list[str] = []
        for part in value:
            if not isinstance(part, dict):
                continue
            kind = str(part.get("type") or "").lower()
            if "image" in kind or "image_url" in part:
                parts.append(ATTACHMENT_NOTE)
                continue
            candidate = part.get("text", part.get("content", ""))
            if isinstance(candidate, str):
                parts.append(candidate)
        text = "\n".join(parts)
    else:
        try:
            text = json.dumps(value, ensure_ascii=False, sort_keys=True)
        except (TypeError, ValueError):
            text = str(value)
    text = _DATA_IMAGE.sub(ATTACHMENT_NOTE, text)
    text = _MEDIA_PATH.sub("[private attachment]", text)
    return text[:MAX_ITEM_CHARS]


def _normalized_item(raw: Any) -> dict[str, Any] | None:
    if not isinstance(raw, dict):
        return None
    role = str(raw.get("role") or "").strip().lower()
    if role not in {"assistant", "system", "tool", "user"}:
        return None
    content = _text(raw.get("content", ""))
    if not content and role != "assistant":
        return None
    item: dict[str, Any] = {"role": role, "content": content}
    if raw.get("_db_persisted") is True:
        item["_db_persisted"] = True
    for key in ("name", "tool_call_id", "tool_name"):
        value = raw.get(key)
        if isinstance(value, str) and value:
            item[key] = value[:256]
    calls = raw.get("tool_calls")
    if isinstance(calls, list):
        bounded_calls = []
        for call in calls[:16]:
            if not isinstance(call, dict):
                continue
            function = (
                call.get("function") if isinstance(call.get("function"), dict) else {}
            )
            bounded = {
                "id": str(call.get("id") or "")[:256],
                "type": str(call.get("type") or "function")[:32],
                "function": {
                    "name": str(function.get("name") or "")[:256],
                    "arguments": _text(function.get("arguments", ""))[:2_000],
                },
            }
            bounded_calls.append(bounded)
        if bounded_calls:
            item["tool_calls"] = bounded_calls
    return item


def _fingerprint(item: dict[str, Any]) -> str:
    try:
        visible = {key: value for key, value in item.items() if not key.startswith("_")}
        return json.dumps(
            visible, ensure_ascii=False, sort_keys=True, separators=(",", ":")
        )
    except (TypeError, ValueError):
        return repr(item)


def _deduplicate(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Drop replay duplicates while preserving intentional repeated prose."""
    result: list[dict[str, Any]] = []
    stable_ids: set[tuple[str, str]] = set()
    for item in items:
        stable = ""
        if isinstance(item.get("tool_call_id"), str):
            stable = item["tool_call_id"]
        elif isinstance(item.get("id"), str):
            stable = item["id"]
        if stable:
            key = (item["role"], stable)
            if key in stable_ids:
                continue
            stable_ids.add(key)
        if result and _fingerprint(result[-1]) == _fingerprint(item):
            continue
        result.append(item)
    # Response replays can append one exact transcript to itself. Remove only
    # complete adjacent blocks, not ordinary repeated user phrases.
    changed = True
    while changed and len(result) > 1:
        changed = False
        for width in range(len(result) // 2, 0, -1):
            if [_fingerprint(x) for x in result[-2 * width : -width]] == [
                _fingerprint(x) for x in result[-width:]
            ]:
                del result[-width:]
                changed = True
                break
    return result


def _summary(older: list[dict[str, Any]]) -> dict[str, str] | None:
    lines: list[str] = []
    for item in older:
        role = item["role"]
        content = str(item.get("content") or "").strip()
        if role == "system" and content.startswith(SUMMARY_PREFIX):
            candidates = content[len(SUMMARY_PREFIX) :].strip().splitlines()
        elif role in {"user", "assistant"} and content:
            candidates = [f"- {role.title()}: {content[:600]}"]
        else:
            continue
        for line in candidates:
            clean = " ".join(line.split())
            if clean and clean not in lines:
                lines.append(clean)
    content = SUMMARY_PREFIX
    for line in lines[-40:]:
        candidate = content + "\n" + line
        if len(candidate) > MAX_SUMMARY_CHARS:
            break
        content = candidate
    return (
        {
            "role": "system",
            "content": content,
            "_db_persisted": True,
            "_compressed_summary": True,
        }
        if lines
        else None
    )


def compact_telegram_history(history: Any) -> list[dict[str, Any]]:
    """Return a deterministic summary plus bounded recent Telegram turns."""
    if not isinstance(history, list):
        return []
    normalized = _deduplicate(
        [item for raw in history if (item := _normalized_item(raw)) is not None]
    )
    recent: list[dict[str, Any]] = []
    recent_bytes = 2
    user_turns = 0
    split = len(normalized)
    for index in range(len(normalized) - 1, -1, -1):
        item = normalized[index]
        encoded = len(_fingerprint(item).encode("utf-8")) + 1
        next_turns = user_turns + (item["role"] == "user")
        if recent and (
            len(recent) >= MAX_HISTORY_ITEMS
            or recent_bytes + encoded > MAX_HISTORY_BYTES
            or next_turns > MAX_RECENT_TURNS
        ):
            break
        recent.insert(0, item)
        recent_bytes += encoded
        user_turns = next_turns
        split = index
    summary = _summary(normalized[:split])
    result = ([summary] if summary else []) + recent
    while (
        len(json.dumps(result, ensure_ascii=False).encode("utf-8"))
        > (MAX_HISTORY_BYTES + MAX_SUMMARY_CHARS)
        and len(recent) > 1
    ):
        recent.pop(0)
        result = ([summary] if summary else []) + recent
    return result


def telegram_sessions(response_store: Path) -> dict[str, str]:
    """Return safe session ID to stable Telegram conversation 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-%')
                  AND LENGTH(r.data) <= ?
                ORDER BY c.name, r.response_id""",
            (MAX_RESPONSE_BYTES,),
        ).fetchall()
    except sqlite3.DatabaseError:
        return {}
    finally:
        connection.close()
    result: dict[str, str] = {}
    for conversation, raw in rows:
        name = str(conversation or "")
        if not _CONVERSATION.fullmatch(name):
            continue
        try:
            payload = json.loads(raw)
            session_id = str(payload.get("session_id") or "").strip()
        except (AttributeError, TypeError, ValueError, json.JSONDecodeError):
            continue
        if (
            session_id
            and len(session_id) <= 256
            and not re.search(r"[\r\n\x00/\\]", session_id)
        ):
            result.setdefault(session_id, name)
    return result


def deduplicate_stored_messages(
    connection: sqlite3.Connection, session_ids: list[str]
) -> int:
    """Remove only byte-for-byte replay rows for known Telegram API sessions."""
    columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(messages)")}
    required = {"id", "session_id", "role", "content", "timestamp"}
    if not required.issubset(columns):
        return 0
    identity = [column for column in _MESSAGE_IDENTITY_COLUMNS if column in columns]
    partition = ", ".join(identity)
    changed = 0
    for session_id in session_ids:
        query = f"""DELETE FROM messages WHERE id IN (
                    SELECT id FROM (
                        SELECT id, ROW_NUMBER() OVER (
                            PARTITION BY {partition} ORDER BY id
                        ) AS replay_number
                        FROM messages
                        WHERE session_id = ? AND EXISTS (
                            SELECT 1 FROM sessions
                            WHERE sessions.id = messages.session_id
                              AND sessions.source = 'api_server'
                        )
                    ) WHERE replay_number > 1
                )"""  # noqa: S608 -- identifiers come only from the fixed allowlist
        cursor = connection.execute(query, (session_id,))
        changed += max(cursor.rowcount, 0)
    return changed


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)")
        }
        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=(",", ":"),
            )
            default_title = (
                "Telegram · General" if conversation == "telegram" else "Telegram"
            )
            cursor = connection.execute(
                """UPDATE sessions
                      SET session_key = ?, chat_type = 'private',
                          display_name = 'Telegram', origin_json = ?,
                          title = CASE
                            WHEN title IS NULL OR TRIM(title) = '' OR
                                 LOWER(TRIM(title)) = 'unassigned'
                            THEN ? ELSE title END
                    WHERE id = ? AND source = 'api_server'
                      AND (COALESCE(session_key, '') != ? OR
                           COALESCE(chat_type, '') != 'private' OR
                           COALESCE(display_name, '') != 'Telegram' OR
                           COALESCE(origin_json, '') != ? OR title IS NULL OR
                           TRIM(title) = '' OR LOWER(TRIM(title)) = 'unassigned')""",
                (
                    conversation,
                    origin,
                    default_title,
                    session_id,
                    conversation,
                    origin,
                ),
            )
            changed += max(cursor.rowcount, 0)
        deduplicate_stored_messages(connection, list(mappings))
        connection.commit()
    except sqlite3.DatabaseError:
        connection.rollback()
        return 0
    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())