#!/usr/bin/env python3
"""Offline, non-destructive migration from bot-chat data to unified-UID chat data.

The source chat directory and identity accounts are opened read-only in spirit:
no source file is replaced or removed. All destination text writes are atomic,
and each SQLite migration runs in one IMMEDIATE transaction with FTS validation.
"""
import argparse
import hashlib
import json
import os
import sqlite3
import tempfile
from pathlib import Path

TG_USER_OFFSET = 10 ** 10
SCHEMA = """
CREATE TABLE IF NOT EXISTS facts(
    id INTEGER PRIMARY KEY,
    uid INTEGER NOT NULL,
    content TEXT NOT NULL,
    importance INTEGER NOT NULL DEFAULT 3,
    keywords TEXT NOT NULL DEFAULT '',
    created_at REAL NOT NULL,
    archived INTEGER NOT NULL DEFAULT 0,
    source TEXT NOT NULL DEFAULT ''
);
CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts USING fts5(
    content, keywords, content='facts', content_rowid='id', tokenize='trigram'
);
CREATE TRIGGER IF NOT EXISTS facts_ai AFTER INSERT ON facts BEGIN
    INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
END;
CREATE TRIGGER IF NOT EXISTS facts_ad AFTER DELETE ON facts BEGIN
    INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
    VALUES ('delete', old.id, old.content, old.keywords);
END;
CREATE TRIGGER IF NOT EXISTS facts_au AFTER UPDATE ON facts BEGIN
    INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
    VALUES ('delete', old.id, old.content, old.keywords);
    INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
END;
CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT);
CREATE TABLE IF NOT EXISTS triples(
    id INTEGER PRIMARY KEY,
    uid INTEGER NOT NULL,
    subject TEXT NOT NULL,
    predicate TEXT NOT NULL,
    object TEXT NOT NULL,
    importance INTEGER NOT NULL DEFAULT 3,
    valid_from REAL NOT NULL,
    valid_until REAL,
    created_at REAL NOT NULL,
    source TEXT NOT NULL DEFAULT 'model'
);
CREATE INDEX IF NOT EXISTS triples_uid ON triples(uid);
CREATE INDEX IF NOT EXISTS triples_subj ON triples(uid, subject);
CREATE TABLE IF NOT EXISTS relations(
    uid INTEGER PRIMARY KEY,
    first_met REAL,
    msg_count INTEGER NOT NULL DEFAULT 0,
    affinity INTEGER NOT NULL DEFAULT 50,
    last_seen REAL,
    notes TEXT NOT NULL DEFAULT '[]'
);
"""


def _atomic_bytes(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(prefix=".migrate-", dir=str(path.parent))
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, path)
    except Exception:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def _atomic_text(path, text):
    _atomic_bytes(path, text.encode("utf-8"))


def _atomic_json(path, value):
    _atomic_text(path, json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n")


def _read_json(path):
    with open(path, "r", encoding="utf-8") as fh:
        return json.load(fh)


def _int(value):
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _source_name(source, path):
    try:
        return str(Path(path).relative_to(source))
    except ValueError:
        return str(path)


class ChatMigrator:
    def __init__(self, source, identity_accounts, dest, raw_tg=False, provision_tg=False):
        self.source = Path(source).resolve()
        self.accounts = Path(identity_accounts).resolve()
        self.dest = Path(dest).resolve()
        self.raw_tg = bool(raw_tg)
        self.provision_tg = bool(provision_tg)
        if not self.source.is_dir():
            raise ValueError("source chat directory does not exist: %s" % self.source)
        if not self.accounts.is_dir():
            raise ValueError("identity accounts directory does not exist: %s" % self.accounts)
        if self.source == self.dest:
            raise ValueError("source and destination must differ")
        self.dest.mkdir(parents=True, exist_ok=True)
        self.report = {
            "version": 1,
            "source": str(self.source),
            "identity_accounts": str(self.accounts),
            "dest": str(self.dest),
            "mapping": {"tg_ns_to_uid": {}, "old_key_to_uid": {}},
            "files": {
                "history_lines": 0, "memory_lines": 0, "group_lines": 0,
                "status_files": 0, "setting_merged": False, "stickers_copied": 0,
            },
            "sqlite": {
                "source_found": False, "facts": 0, "triples": 0,
                "relations": 0, "meta": 0, "fact_conflicts": 0,
                "triple_conflicts": 0, "meta_conflicts": 0,
                "unmatched": {"facts": 0, "triples": 0, "relations": 0},
                "fts_validated": False,
            },
            "quarantine": [],
            "conflicts": [],
            "provisioned_tg": [],
        }
        self.old_to_uid = {}
        self.uid_set = set()
        self._unmatched_db = []

    def _conflict(self, kind, source, detail):
        self.report["conflicts"].append({
            "kind": kind, "source": _source_name(self.source, source), "detail": detail,
        })

    def _quarantine(self, path, category, reason):
        base = self.dest / "quarantine" / category / Path(path).name
        target = base
        data = Path(path).read_bytes()
        if target.exists() and target.read_bytes() != data:
            # Same basenames may come from different source directories, and a
            # later migration must not overwrite evidence from an earlier run.
            digest = hashlib.sha256(data).hexdigest()[:12]
            target = base.with_name("%s-%s%s" % (base.stem, digest, base.suffix))
            serial = 2
            while target.exists() and target.read_bytes() != data:
                target = base.with_name("%s-%s-%d%s" % (
                    base.stem, digest, serial, base.suffix))
                serial += 1
        if not target.exists():
            _atomic_bytes(target, data)
        self.report["quarantine"].append({
            "source": _source_name(self.source, path),
            "target": str(target.relative_to(self.dest)),
            "reason": reason,
        })

    def _legacy_private_keys(self):
        keys = set()
        for category, pattern in (("history", "*.txt"), ("memory", "*.txt"), ("status", "*.json")):
            for path in (self.source / category).glob(pattern):
                value = _int(path.stem)
                if value is not None and value > 0:
                    keys.add(value)
        source_db = self.source / "memory.db"
        if source_db.is_file():
            con = sqlite3.connect("file:%s?mode=ro" % source_db.as_posix(), uri=True)
            try:
                for table in ("facts", "triples", "relations"):
                    try:
                        for (value,) in con.execute("SELECT DISTINCT uid FROM %s" % table):
                            value = _int(value)
                            if value is not None and value > 0:
                                keys.add(value)
                    except sqlite3.Error:
                        pass
            finally:
                con.close()
        return sorted(keys)

    def provision_legacy_tg(self):
        """Explicitly create missing identities for the old TG-only raw-ID layout."""
        if not self.provision_tg:
            return
        import sys
        sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
        import identity
        identity.configure(str(self.accounts))
        for old_key in self._legacy_private_keys():
            ns = old_key if old_key >= TG_USER_OFFSET else TG_USER_OFFSET + old_key
            before = identity.resolve("tg", ns)
            uid = identity.ensure_tg(ns)
            if before is None and uid is not None:
                self.report["provisioned_tg"].append({"old_key": old_key, "tg_ns": ns, "uid": uid})

    def load_identity(self):
        candidates = {}
        tg_candidates = {}
        users = self.accounts / "users"
        for path in sorted(users.glob("*.json")):
            try:
                rec = _read_json(path)
            except (OSError, UnicodeError, json.JSONDecodeError) as exc:
                self._conflict("identity_invalid", path, str(exc))
                continue
            if not isinstance(rec, dict):
                self._conflict("identity_invalid", path, "record is not an object")
                continue
            uid = _int(rec.get("uid"))
            if uid is None:
                self._conflict("identity_invalid", path, "missing numeric uid")
                continue
            self.uid_set.add(uid)
            candidates.setdefault(str(uid), set()).add(uid)
            tg = rec.get("tg") or {}
            if isinstance(tg, dict):
                ns = _int(tg.get("ns"))
                raw = _int(tg.get("raw"))
                if ns is None and raw is not None:
                    ns = TG_USER_OFFSET + raw
                if ns is not None:
                    candidates.setdefault(str(ns), set()).add(uid)
                    tg_candidates.setdefault(str(ns), set()).add(uid)
            qq_items = rec.get("qq") or []
            if isinstance(qq_items, list):
                for item in qq_items:
                    qq = _int(item.get("id")) if isinstance(item, dict) else None
                    if qq is not None:
                        candidates.setdefault(str(qq), set()).add(uid)
        if self.raw_tg:
            # The legacy chat process was TG-only. A small integer is a raw TG
            # user ID, never a QQ guess; namespace IDs are used as-is.
            for key in self._legacy_private_keys():
                ns = key if key >= TG_USER_OFFSET else key + TG_USER_OFFSET
                owners = tg_candidates.get(str(ns), set())
                if len(owners) == 1:
                    self.old_to_uid[str(key)] = next(iter(owners))
                elif len(owners) > 1:
                    self._conflict("identity_key_ambiguous", self.accounts,
                                   "%s -> %s" % (key, sorted(owners)))
        else:
            for key, owners in sorted(candidates.items()):
                if len(owners) == 1:
                    self.old_to_uid[key] = next(iter(owners))
                else:
                    self._conflict("identity_key_ambiguous", self.accounts, "%s -> %s" % (key, sorted(owners)))
        self.report["mapping"]["old_key_to_uid"] = dict(sorted(self.old_to_uid.items()))
        self.report["mapping"]["tg_ns_to_uid"] = {
            key: self.old_to_uid[key] for key in sorted(tg_candidates) if key in self.old_to_uid
        }

    def _merge_text(self, source, target):
        with open(source, "r", encoding="utf-8") as fh:
            incoming = [line.rstrip("\r\n") for line in fh if line.strip()]
        existing = []
        if Path(target).exists():
            with open(target, "r", encoding="utf-8") as fh:
                existing = [line.rstrip("\r\n") for line in fh if line.strip()]
        seen = set(existing)
        additions = []
        for line in incoming:
            if line not in seen:
                seen.add(line)
                additions.append(line)
        if additions or not Path(target).exists():
            _atomic_text(target, "\n".join(existing + additions) + ("\n" if existing or additions else ""))
        return len(additions)

    def _merge_json(self, current, incoming, source, prefix=""):
        for key, value in sorted(incoming.items()):
            path = "%s.%s" % (prefix, key) if prefix else str(key)
            if key not in current:
                current[key] = value
            elif isinstance(current[key], dict) and isinstance(value, dict):
                self._merge_json(current[key], value, source, path)
            elif current[key] != value:
                self._conflict("json_value", source, "%s: destination value retained" % path)

    def _merge_json_file(self, source, target):
        incoming = _read_json(source)
        if not isinstance(incoming, dict):
            raise ValueError("JSON root must be an object")
        current = _read_json(target) if Path(target).exists() else {}
        if not isinstance(current, dict):
            raise ValueError("destination JSON root must be an object: %s" % target)
        self._merge_json(current, incoming, source)
        _atomic_json(target, current)

    def _migrate_private_text(self, category, suffix):
        src_dir = self.source / category
        dst_dir = self.dest / category
        if not src_dir.is_dir():
            return
        for path in sorted(src_dir.glob("*" + suffix)):
            if not path.is_file():
                continue
            key = path.name[:-len(suffix)]
            uid = self.old_to_uid.get(key)
            if uid is None:
                self._quarantine(path, category, "no unambiguous identity mapping for old key " + key)
                continue
            target = dst_dir / (str(uid) + suffix)
            try:
                if suffix == ".txt":
                    added = self._merge_text(path, target)
                    self.report["files"]["history_lines" if category == "history" else "memory_lines"] += added
                else:
                    self._merge_json_file(path, target)
                    self.report["files"]["status_files"] += 1
            except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc:
                self._quarantine(path, category, "invalid source data: %s" % exc)

    def migrate_files(self):
        self._migrate_private_text("history", ".txt")
        self._migrate_private_text("memory", ".txt")
        self._migrate_private_text("status", ".json")

        group_dir = self.source / "history" / "group"
        if group_dir.is_dir():
            for path in sorted(group_dir.glob("*.txt")):
                if path.is_file():
                    # Group namespaces deliberately remain unchanged.
                    self.report["files"]["group_lines"] += self._merge_text(
                        path, self.dest / "history" / "group" / path.name)

        setting = self.source / "setting.json"
        if setting.is_file():
            try:
                self._merge_json_file(setting, self.dest / "setting.json")
                self.report["files"]["setting_merged"] = True
            except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc:
                self._quarantine(setting, "setting", "invalid setting: %s" % exc)

        stickers = self.source / "stickers"
        if stickers.is_dir():
            for path in sorted(p for p in stickers.rglob("*") if p.is_file()):
                rel = path.relative_to(stickers)
                target = self.dest / "stickers" / rel
                data = path.read_bytes()
                if target.exists():
                    if target.read_bytes() == data:
                        continue
                    self._quarantine(path, "stickers", "destination sticker differs; destination retained")
                    continue
                _atomic_bytes(target, data)
                self.report["files"]["stickers_copied"] += 1

    def _rows(self, con, table, columns):
        available = {row[1] for row in con.execute("PRAGMA table_info(%s)" % table)}
        if not available:
            return []
        missing = set(columns) - available
        if missing:
            raise ValueError("source %s lacks columns: %s" % (table, sorted(missing)))
        con.row_factory = sqlite3.Row
        order = "id" if "id" in available else ("uid" if "uid" in available else "key")
        return [dict(row) for row in con.execute("SELECT * FROM %s ORDER BY %s" % (table, order))]

    def _new_uid(self, table, row):
        old_uid = _int(row.get("uid"))
        uid = self.old_to_uid.get(str(old_uid))
        if uid is None:
            self.report["sqlite"]["unmatched"][table] += 1
            self._unmatched_db.append({"table": table, "row": row})
        return uid

    @staticmethod
    def _notes(value):
        try:
            parsed = json.loads(value or "[]")
            return parsed if isinstance(parsed, list) else []
        except (TypeError, json.JSONDecodeError):
            return []

    def _merge_relation(self, con, source_uid, uid, row):
        marker = "chat_migration:relation:%s" % source_uid
        previous = con.execute("SELECT value FROM meta WHERE key=?", (marker,)).fetchone()
        old_count = 0
        if previous:
            try:
                state = json.loads(previous[0])
                if int(state["uid"]) != uid:
                    raise ValueError("identity target changed")
                old_count = int(state["msg_count"])
            except (ValueError, TypeError, KeyError, json.JSONDecodeError) as exc:
                raise ValueError("invalid migration relation marker %s: %s" % (marker, exc))
        count = int(row.get("msg_count") or 0)
        delta = count - old_count
        existing = con.execute(
            "SELECT first_met,msg_count,affinity,last_seen,notes FROM relations WHERE uid=?", (uid,)
        ).fetchone()
        if existing is None:
            con.execute(
                "INSERT INTO relations(uid,first_met,msg_count,affinity,last_seen,notes) VALUES(?,?,?,?,?,?)",
                (uid, row.get("first_met"), count, int(row.get("affinity") or 50),
                 row.get("last_seen"), row.get("notes") or "[]"),
            )
        else:
            first_values = [v for v in (existing[0], row.get("first_met")) if v is not None]
            first_met = min(first_values) if first_values else None
            last_values = [v for v in (existing[3], row.get("last_seen")) if v is not None]
            last_seen = max(last_values) if last_values else None
            notes = self._notes(existing[4])
            serialized = {json.dumps(item, ensure_ascii=False, sort_keys=True) for item in notes}
            for item in self._notes(row.get("notes")):
                key = json.dumps(item, ensure_ascii=False, sort_keys=True)
                if key not in serialized:
                    notes.append(item)
                    serialized.add(key)
            notes = sorted(notes, key=lambda item: item.get("ts", 0) if isinstance(item, dict) else 0)[-5:]
            con.execute(
                "UPDATE relations SET first_met=?,msg_count=?,last_seen=?,notes=? WHERE uid=?",
                (first_met, max(0, int(existing[1]) + delta), last_seen,
                 json.dumps(notes, ensure_ascii=False), uid),
            )
        con.execute("INSERT OR REPLACE INTO meta(key,value) VALUES(?,?)", (
            marker, json.dumps({"uid": uid, "msg_count": count}, sort_keys=True),
        ))
        self.report["sqlite"]["relations"] += 1

    def migrate_sqlite(self):
        source_db = self.source / "memory.db"
        dest_db = self.dest / "memory.db"
        if source_db == dest_db:
            raise ValueError("source and destination memory.db must differ")
        src = None
        if source_db.is_file():
            self.report["sqlite"]["source_found"] = True
            src = sqlite3.connect("file:%s?mode=ro" % source_db.as_posix(), uri=True)
        out = sqlite3.connect(str(dest_db), timeout=30)
        try:
            out.execute("PRAGMA journal_mode=WAL")
            out.executescript(SCHEMA)
            if src is None:
                out.close()
                out = None
                return
            fact_rows = self._rows(src, "facts", {
                "id", "uid", "content", "importance", "keywords", "created_at", "archived", "source",
            })
            triple_rows = self._rows(src, "triples", {
                "id", "uid", "subject", "predicate", "object", "importance",
                "valid_from", "valid_until", "created_at", "source",
            })
            relation_rows = self._rows(src, "relations", {
                "uid", "first_met", "msg_count", "affinity", "last_seen", "notes",
            })
            meta_rows = self._rows(src, "meta", {"key", "value"})
            fact_rows.sort(key=lambda row: (int(row["uid"]), int(row["id"])))
            triple_rows.sort(key=lambda row: (int(row["uid"]), int(row["id"])))
            relation_rows.sort(key=lambda row: int(row["uid"]))
            out.execute("BEGIN IMMEDIATE")
            for row in fact_rows:
                uid = self._new_uid("facts", row)
                if uid is None:
                    continue
                duplicate = out.execute(
                    "SELECT id FROM facts WHERE uid=? AND content=?", (uid, row["content"])
                ).fetchone()
                if duplicate:
                    self.report["sqlite"]["fact_conflicts"] += 1
                    continue
                out.execute(
                    "INSERT INTO facts(uid,content,importance,keywords,created_at,archived,source)"
                    " VALUES(?,?,?,?,?,?,?)",
                    (uid, row["content"], int(row["importance"]), row["keywords"] or "",
                     float(row["created_at"]), int(row["archived"]), row["source"] or ""),
                )
                self.report["sqlite"]["facts"] += 1
            for row in triple_rows:
                uid = self._new_uid("triples", row)
                if uid is None:
                    continue
                duplicate = out.execute(
                    "SELECT id FROM triples WHERE uid=? AND subject=? AND predicate=? AND object=?"
                    " AND valid_from=? AND valid_until IS ?",
                    (uid, row["subject"], row["predicate"], row["object"],
                     float(row["valid_from"]), row["valid_until"]),
                ).fetchone()
                if duplicate:
                    self.report["sqlite"]["triple_conflicts"] += 1
                    continue
                out.execute(
                    "INSERT INTO triples(uid,subject,predicate,object,importance,valid_from,valid_until,created_at,source)"
                    " VALUES(?,?,?,?,?,?,?,?,?)",
                    (uid, row["subject"], row["predicate"], row["object"], int(row["importance"]),
                     float(row["valid_from"]), row["valid_until"], float(row["created_at"]), row["source"] or "model"),
                )
                self.report["sqlite"]["triples"] += 1
            for row in relation_rows:
                source_uid = _int(row.get("uid"))
                uid = self._new_uid("relations", row)
                if uid is not None:
                    self._merge_relation(out, source_uid, uid, row)
            for row in sorted(meta_rows, key=lambda item: item["key"]):
                key, value = row["key"], row["value"]
                if key.startswith("chat_migration:"):
                    continue
                if key.startswith("imported:"):
                    old_key = key[len("imported:"):]
                    mapped = self.old_to_uid.get(old_key)
                    if mapped is None:
                        self._unmatched_db.append({"table": "meta", "row": row})
                        continue
                    key = "imported:%s" % mapped
                existing = out.execute("SELECT value FROM meta WHERE key=?", (key,)).fetchone()
                if existing is not None:
                    if existing[0] != value:
                        self.report["sqlite"]["meta_conflicts"] += 1
                    continue
                out.execute("INSERT INTO meta(key,value) VALUES(?,?)", (key, value))
                self.report["sqlite"]["meta"] += 1
            facts_count = out.execute("SELECT COUNT(*) FROM facts").fetchone()[0]
            fts_count = out.execute("SELECT COUNT(*) FROM facts_fts").fetchone()[0]
            if facts_count != fts_count:
                raise ValueError("FTS row count mismatch: facts=%d fts=%d" % (facts_count, fts_count))
            out.execute("INSERT INTO facts_fts(facts_fts) VALUES('integrity-check')")
            out.commit()
            self.report["sqlite"]["fts_validated"] = True
        except Exception:
            if out is not None:
                out.rollback()
            raise
        finally:
            if src is not None:
                src.close()
            if out is not None:
                out.close()
        if self._unmatched_db:
            target = self.dest / "quarantine" / "sqlite_unmatched.json"
            _atomic_json(target, self._unmatched_db)
            self.report["quarantine"].append({
                "source": "memory.db", "target": str(target.relative_to(self.dest)),
                "reason": "SQLite rows without an unambiguous identity mapping",
            })

    def run(self):
        if self.raw_tg:
            self.provision_legacy_tg()
        self.load_identity()
        self.migrate_files()
        self.migrate_sqlite()
        return self.report


def migrate(source, identity_accounts, dest, report=None, raw_tg=False, provision_tg=False):
    migrator = ChatMigrator(source, identity_accounts, dest, raw_tg=raw_tg, provision_tg=provision_tg)
    result = migrator.run()
    report_path = Path(report) if report else migrator.dest / "migration-report.json"
    _atomic_json(report_path, result)
    return result


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, help="old bot-chat directory")
    parser.add_argument("--identity-accounts", required=True, help="identity accounts directory")
    parser.add_argument("--dest", required=True, help="new bot-utility/chat directory")
    parser.add_argument("--report", help="JSON report path (default: DEST/migration-report.json)")
    parser.add_argument("--raw-tg", action="store_true",
                        help="interpret legacy private keys as raw/namespaced TG IDs only")
    parser.add_argument("--provision-tg", action="store_true",
                        help="with --raw-tg, explicitly create missing identity records")
    args = parser.parse_args(argv)
    result = migrate(args.source, args.identity_accounts, args.dest, args.report,
                     raw_tg=args.raw_tg, provision_tg=args.provision_tg)
    print(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()