#!/usr/bin/env python3
#
# tg_backup.py - incremental Telegram media backup
#
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <https://www.gnu.org/licenses/>.

import argparse
import asyncio
import json
import os
import shutil
import sys
from pathlib import Path

from telethon import TelegramClient
from telethon.errors import FloodWaitError
from telethon.tl.types import MessageMediaPhoto, MessageMediaDocument

CONFIG_DIR = Path.home() / ".config" / "telegram_backup"
CONFIG_FILE = CONFIG_DIR / "config.json"
# Shared across every archive dir: one login covers all chats. Telethon appends
# ".session", so the file on disk is session.session.
SESSION_FILE = CONFIG_DIR / "session"
MAX_RETRIES = 5

def bootstrap_config():
    # Called before argparse: --target is required, but a user with no config
    # has nothing useful to pass yet. Without this, a bare run exits on the
    # missing argument and never creates the template.
    if not CONFIG_FILE.exists():
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        CONFIG_DIR.chmod(0o700)  # holds API keys and the shared session
        template = {"api_id": "YOUR_API_ID", "api_hash": "YOUR_API_HASH"}
        CONFIG_FILE.write_text(json.dumps(template, indent=4))
        CONFIG_FILE.chmod(0o600)
        sys.exit(f"Config created at {CONFIG_FILE}. Fill in API keys and run again.")

def load_config():
    bootstrap_config()

    with open(CONFIG_FILE, 'r') as f:
        config = json.load(f)

    api_id, api_hash = config.get('api_id'), config.get('api_hash')
    if not api_id or not api_hash or 'YOUR_' in (str(api_id) + str(api_hash)):
        sys.exit("Error: Invalid or missing API keys in config.")
    return config

def load_state(archive_dir):
    state_file = archive_dir / "state.json"
    if state_file.exists():
        with open(state_file, 'r') as f:
            return json.load(f)
    return {'last_id': 0}

def save_state(archive_dir, state):
    # Written after every message, so an interrupted write is a matter of time.
    # Write to a temp file and rename: readers see the old state or the new one,
    # never a truncated file that would make the next run unstartable.
    state_file = archive_dir / "state.json"
    tmp_file = state_file.with_name(state_file.name + ".tmp")
    with open(tmp_file, 'w') as f:
        json.dump(state, f, indent=4)
    os.replace(tmp_file, state_file)

def record_failure(archive_dir, message_id, filename, error):
    # State advances past a permanently failed download on purpose, so without
    # this the message ID is lost and the backup silently has a hole in it.
    # Append-only log, one JSON object per line: a partial last line costs one
    # record, not the whole file.
    entry = {'id': message_id, 'file': filename, 'error': str(error)}
    with open(archive_dir / "failures.json", 'a') as f:
        f.write(json.dumps(entry) + "\n")

def safe_filename(message_id, original_name, ext):
    # original_name comes from the sender. Strip path components so a name like
    # "../../.ssh/authorized_keys" resolves to a plain file inside archive_dir.
    name = Path(original_name or f"media{ext}").name
    if name in ('', '.', '..'):
        name = f"media{ext}"
    return f"{message_id}_{name}"

async def download_with_retry(client, message, filepath):
    # Download to a .part file so an interrupted run never leaves a truncated
    # file that the exists() check would mistake for a complete download.
    partpath = filepath.with_name(filepath.name + ".part")
    last_error = f"no attempt made (MAX_RETRIES={MAX_RETRIES})"
    for attempt in range(MAX_RETRIES):
        try:
            await client.download_media(message, file=str(partpath))
            os.replace(partpath, filepath)
            return None
        except FloodWaitError as e:
            last_error = e
            print(f"Rate limited on download. Waiting {e.seconds}s... (Attempt {attempt+1}/{MAX_RETRIES})")
            await asyncio.sleep(e.seconds)
        except (KeyboardInterrupt, asyncio.CancelledError):
            partpath.unlink(missing_ok=True)
            raise
        except Exception as e:
            # Network drops, expired file references, disk errors. Backing off
            # and retrying beats killing an overnight run on one bad socket.
            last_error = e
            delay = 2 ** attempt
            print(f"Download error on {filepath.name}: {e}. "
                  f"Retrying in {delay}s... (Attempt {attempt+1}/{MAX_RETRIES})")
            partpath.unlink(missing_ok=True)
            await asyncio.sleep(delay)
    partpath.unlink(missing_ok=True)
    print(f"Failed to download {filepath.name} after {MAX_RETRIES} retries.")
    # Returns the last error, or None on success: the caller logs it to failures.json.
    return last_error

def migrate_session(archive_dir):
    # Sessions used to live per-archive-dir, which forced a fresh login for every
    # chat. Move an old one to the shared location instead of making the user
    # re-authenticate. Only runs when there is no shared session yet.
    shared = SESSION_FILE.with_suffix(".session")
    old = archive_dir / "session.session"
    if not shared.exists() and old.exists():
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        shutil.move(str(old), str(shared))  # may cross filesystems
        shared.chmod(0o600)
        print(f"Moved existing login from {old} to {shared}.")

def parse_target(target):
    # Telethon resolves a numeric *string* as a username and fails; only a real
    # int hits the session's cached-entity lookup. --list-chats prints ints, so
    # convert them back before handing off to get_entity().
    t = target.strip()
    if t.lstrip('-').isdigit():
        return int(t)
    return t

def same_target(a, b):
    # Only catches targets that are textually the same chat. "@foo" and "foo"
    # match; "@foo" and its numeric ID do not, since resolving them needs a
    # logged-in client and the check runs before that.
    if a is None or b is None:
        return True  # nothing saved to contradict
    norm = lambda t: str(t).strip().lstrip('@').lower()
    return norm(a) == norm(b)

def make_client(archive_dir):
    # One session in CONFIG_DIR, shared by every archive dir: log in once, back
    # up any number of chats.
    config = load_config()
    archive_dir.mkdir(parents=True, exist_ok=True)
    migrate_session(archive_dir)
    return TelegramClient(str(SESSION_FILE), config['api_id'], config['api_hash'])

async def start_client(client):
    # Telethon prompts for phone/code on first login. Without a TTY that
    # surfaces as a bare EOFError traceback, which says nothing useful.
    try:
        await client.start()
    except EOFError:
        sys.exit("First login needs an interactive terminal (phone + code). "
                 "Run this once by hand, then the saved session works unattended.")

def entity_username(entity):
    # Telethon exposes both a legacy .username and a newer .usernames list
    # (multiple public handles). .username can be None while .usernames is set.
    username = getattr(entity, 'username', None)
    if username:
        return username
    for u in getattr(entity, 'usernames', None) or []:
        name = getattr(u, 'username', None)
        if name:
            return name
    return None

async def list_chats(archive_dir):
    client = make_client(archive_dir)
    await start_client(client)
    try:
        print(f"{'ID':>15}  {'TYPE':<8}  {'USERNAME':<20}  NAME")
        async for dialog in client.iter_dialogs():
            kind = 'user' if dialog.is_user else 'group' if dialog.is_group else 'channel'
            # Not every chat has one: private groups and users who never set a
            # public handle resolve by numeric ID only.
            username = entity_username(dialog.entity)
            handle = f"@{username}" if username else "-"
            print(f"{dialog.id:>15}  {kind:<8}  {handle:<20}  {dialog.name}")
    finally:
        await client.disconnect()
    print("\nBack up with: --target=<ID>   (the = matters for negative IDs)")

async def run_backup(target, archive_dir, force_target=False):
    # last_id is a high-water mark for one specific chat. Applied as min_id to a
    # different one it silently skips that chat's history, and the run reports
    # success. Refuse before touching the network.
    saved = load_state(archive_dir).get('target')
    if not force_target and not same_target(saved, target):
        sys.exit(f"Archive dir {archive_dir} was last backed up with target "
                 f"'{saved}', not '{target}'.\n"
                 "Its saved resume position belongs to that chat and would skip "
                 "this one's history. Use a separate archive dir per chat, or "
                 "pass --force-target to overwrite (only correct if the same "
                 "chat changed handle or ID).")

    client = make_client(archive_dir)
    await start_client(client)

    try:
        entity = await client.get_entity(parse_target(target))
    except Exception as e:
        sys.exit(f"Failed to resolve target '{target}': {e}\n"
                 "Run --list-chats first: numeric IDs only resolve once the "
                 "chat is in the session's entity cache.")

    state = load_state(archive_dir)
    # Remember the target so later runs can resume with just --archive-dir.
    state['target'] = target
    last_id = state['last_id']
    save_state(archive_dir, state)
    print(f"Resuming from message ID: {last_id}")
    
    # reverse=True ensures oldest-to-newest processing. 
    # Critical for safe incremental state tracking.
    while True:
        try:
            async for message in client.iter_messages(entity, min_id=last_id, reverse=True):
                if isinstance(message.media, (MessageMediaPhoto, MessageMediaDocument)):
                    ext = message.file.ext or '.jpg'
                    filename = safe_filename(message.id, message.file.name, ext)
                    filepath = archive_dir / filename
                    
                    if not filepath.exists():
                        print(f"Downloading {filename}...")
                        error = await download_with_retry(client, message, filepath)
                        if error is not None:
                            # If download permanently fails, we still advance state
                            # so we don't get stuck on a broken file forever.
                            # Log it so the gap is auditable instead of silent.
                            record_failure(archive_dir, message.id, filename, error)
                            print(f"Skipping file due to max retries "
                                  f"(logged to {archive_dir / 'failures.json'}).")
                    else:
                        print(f"Skipping {filename} (exists).")
                        
                # Save state after every message to ensure progress is kept
                state['last_id'] = message.id
                save_state(archive_dir, state)
            break # Loop completes successfully
            
        except FloodWaitError as e:
            print(f"Rate limited on history fetch. Waiting {e.seconds}s...")
            save_state(archive_dir, state) # Save before sleeping
            await asyncio.sleep(e.seconds)
            # Loop restarts, fetching from the last saved ID

    print("Backup complete.")

def main():
    bootstrap_config()
    parser = argparse.ArgumentParser(description="Incremental Telegram media backup.")
    parser.add_argument("--target", help="Group username (@name), ID, or invite link.")
    parser.add_argument("--archive-dir", type=Path, default=Path("./tg_archive"),
                        help="Directory to store media and state.")
    parser.add_argument("--list-chats", action="store_true",
                        help="List your dialogs with their IDs and exit.")
    parser.add_argument("--force-target", action="store_true",
                        help="Back up --target into an archive dir saved with a "
                             "different one, keeping its resume position.")

    args = parser.parse_args()
    if args.list_chats:
        asyncio.run(list_chats(args.archive_dir))
        return
    # An archive dir already backed up once knows its own target.
    target = args.target or load_state(args.archive_dir).get('target')
    if target:
        asyncio.run(run_backup(target, args.archive_dir, args.force_target))
    else:
        parser.error(f"--target is required the first time ({args.archive_dir}/state.json "
                     "has no saved target), or use --list-chats")

def _self_check():
    """Run with --self-check. Guards the path-traversal and state-write logic."""
    import tempfile
    archive = Path(tempfile.mkdtemp()).resolve()

    for hostile in ['../../.ssh/authorized_keys', '..', '.', '', None,
                    '/etc/passwd', 'a/../../b', '....//x', 'ok.jpg']:
        path = (archive / safe_filename(7, hostile, '.jpg')).resolve()
        assert path.parent == archive, f"escaped archive dir: {hostile!r} -> {path}"
    assert safe_filename(7, 'ok.jpg', '.jpg') == '7_ok.jpg'
    assert safe_filename(7, None, '.jpg') == '7_media.jpg'

    # Username extraction must survive both telethon shapes and neither.
    class _E: pass
    class _U:
        def __init__(s, n): s.username = n
    legacy = _E(); legacy.username = 'examplecontact'
    assert entity_username(legacy) == 'examplecontact'
    multi = _E(); multi.username = None; multi.usernames = [_U('newstyle')]
    assert entity_username(multi) == 'newstyle'
    assert entity_username(_E()) is None
    none_set = _E(); none_set.username = None; none_set.usernames = []
    assert entity_username(none_set) is None

    # Numeric targets must reach get_entity() as ints or the cache is missed.
    assert parse_target('123456789') == 123456789
    assert parse_target('-1001234567890') == -1001234567890
    assert parse_target(' 42 ') == 42
    assert parse_target('@examplecontact') == '@examplecontact'
    assert parse_target('me') == 'me'
    assert parse_target('https://t.me/foo') == 'https://t.me/foo'

    # State survives a rewrite and never leaves a stray temp file behind.
    save_state(archive, {'last_id': 1})
    save_state(archive, {'last_id': 42})
    assert load_state(archive) == {'last_id': 42}
    assert not list(archive.glob('*.tmp')), "temp state file left behind"
    assert load_state(Path(tempfile.mkdtemp())) == {'last_id': 0}

    # A saved target survives, and a fresh dir has none to fall back on.
    save_state(archive, {'last_id': 9, 'target': '-1001234567890'})
    assert load_state(archive).get('target') == '-1001234567890'
    assert load_state(Path(tempfile.mkdtemp())).get('target') is None

    # A dir keeps its chat. Handle case and the @ prefix are noise, a different
    # chat is not. Nothing saved yet must never block a first run.
    assert same_target('@examplegroup', 'examplegroup')
    assert same_target('@ExampleGroup', ' @examplegroup ')
    assert same_target('-1001234567890', '-1001234567890')
    assert same_target(-1001234567890, '-1001234567890')  # state may hold an int
    assert same_target(None, '@examplegroup')  # first run, nothing saved
    assert same_target('@examplegroup', None)  # resumed run, no --target given
    assert not same_target('@examplegroup', '@othergroup')
    assert not same_target('-1001234567890', '-1009876543210')
    assert not same_target('@examplegroup', '-1001234567890')  # unresolvable, refuse

    # Failures append, one parseable record per line, and survive an exception
    # object being passed straight in.
    record_failure(archive, 11, '11_a.jpg', OSError('disk full'))
    record_failure(archive, 12, '12_b.mp4', 'timeout')
    lines = (archive / 'failures.json').read_text().splitlines()
    assert len(lines) == 2, f"expected 2 failure records, got {len(lines)}"
    first, second = json.loads(lines[0]), json.loads(lines[1])
    assert first == {'id': 11, 'file': '11_a.jpg', 'error': 'disk full'}
    assert second['id'] == 12 and second['error'] == 'timeout'

    print("self-check OK")

if __name__ == "__main__":
    if '--self-check' in sys.argv:
        _self_check()
    else:
        main()