# Выгрузка файлов транзакций ДЭГ с портала наблюдения stat.vybory.gov.ru
#
# Установка зависимостей:  pip install -r requirements.txt   (Windows: py -m pip ...)
#
# Авторизация: портал работает по HttpOnly-cookie сессии ЕСИА (AUTH_MODE = cookie).
# Bearer-токен больше не используется. Cookie нужно скопировать из браузера:
#   Chrome -> DevTools -> Application -> Cookies -> https://stat.vybory.gov.ru
# и сохранить в файл cookie.txt в виде строки  имя=значение  (несколько cookie — через "; ").
# Реально нужна только stat_acc_t (JWT ЕСИА, живёт 3 часа); после истечения — обновить cookie.txt.
#
# Запуск:
#   python3 file-downloader.py             — все доступные регионы
#   python3 file-downloader.py 21 33       — только Чувашия и Владимирская область
#   python3 file-downloader.py --out ПАПКА — докачать в существующую папку (уже скачанные файлы пропускаются)
#   python3 file-downloader.py --test      — тестовый стенд teststat.deg.rt.ru
#   python3 file-downloader.py --help      — все параметры
#
# Файлы качаются в несколько потоков (--workers) с общим для всех потоков
# ограничением частоты запросов к порталу (--rate) — защита от блокировки WAF.
#
# Лог пишется в консоль и в файл download.log внутри папки выгрузки.

import argparse
import base64
import json
import logging
import os
import re
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

import requests
import urllib3

try:
    from rich.console import Console
    from rich.logging import RichHandler
    from rich.progress import Progress, BarColumn, MofNCompleteColumn, TaskProgressColumn, TextColumn
except ImportError:  # rich не установлен — работаем без прогресс-баров
    Progress = None

from utils import getStrTimeNow

PROD_URL = 'https://stat.vybory.gov.ru/api/'
TEST_URL = 'https://teststat.deg.rt.ru/api/'

HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/151.0.0.0 Safari/537.36',
    'Accept': 'application/json, text/plain, */*',
}

# символы, недопустимые в именах папок (Windows + POSIX), и управляющие
_BAD_CHARS = re.compile(r'[\\/:*?"<>|\x00-\x1f]+')

log = logging.getLogger('downloader')


class AuthError(Exception):
    """Ошибка, при которой продолжать бессмысленно (авторизация, WAF, TLS)."""


def setup_logging(out_dir, verbose, console=None):
    """Консоль: INFO (или DEBUG при --verbose); файл download.log: всегда DEBUG.

    При включённом прогресс-баре в консоль идут только WARNING и выше
    (ход работы виден по барам), подробности — в download.log.
    """
    log.setLevel(logging.DEBUG)
    if console is not None:
        con = RichHandler(console=console, show_time=False, show_path=False, show_level=True, markup=False)
        con.setLevel(logging.DEBUG if verbose else logging.WARNING)
    else:
        con = logging.StreamHandler(sys.stdout)
        con.setLevel(logging.DEBUG if verbose else logging.INFO)
    con.setFormatter(logging.Formatter('%(message)s'))
    log.addHandler(con)

    fh = logging.FileHandler(os.path.join(out_dir, 'download.log'), encoding='utf-8')
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(logging.Formatter('%(asctime)s %(levelname)-7s %(message)s', '%Y-%m-%d %H:%M:%S'))
    log.addHandler(fh)


def safe_name(name, limit=120):
    """Имя папки из названия с портала: убираем запрещённые символы и переносы строк."""
    name = _BAD_CHARS.sub(' ', str(name))
    name = re.sub(r'\s+', ' ', name).strip(' .')
    return name[:limit] or 'unnamed'


def write_info(path, data):
    with open(os.path.join(path, 'info.json'), 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def load_cookies(path):
    """cookie.txt: строка вида `stat_acc_t=abc` или `a=1; b=2` (как в заголовке Cookie)."""
    try:
        with open(path, encoding='utf-8') as f:
            raw = f.read().strip()
    except FileNotFoundError:
        sys.exit(f'Не найден файл {path}. Скопируйте cookie сессии из браузера (см. шапку скрипта).')
    if raw.lower().startswith('cookie:'):
        raw = raw[len('cookie:'):].strip()
    cookies = {}
    for part in raw.split(';'):
        if '=' in part:
            k, v = part.split('=', 1)
            cookies[k.strip()] = v.strip()
    if not cookies:
        sys.exit(f'В файле {path} не найдено ни одной cookie вида имя=значение.')
    return cookies


def token_expiry(cookies):
    """Срок действия stat_acc_t (JWT ЕСИА) либо None, если разобрать не удалось."""
    token = cookies.get('stat_acc_t')
    if not token or token.count('.') != 2:
        return None
    try:
        payload = token.split('.')[1]
        payload += '=' * (-len(payload) % 4)
        exp = json.loads(base64.urlsafe_b64decode(payload))['exp']
        return datetime.fromtimestamp(exp)
    except (ValueError, KeyError, TypeError):
        return None


class RateLimiter:
    """Общий на все потоки лимит частоты запросов к порталу (защита от WAF)."""

    def __init__(self, rate):
        self.interval = 1.0 / rate if rate > 0 else 0.0
        self.lock = threading.Lock()
        self.next_at = 0.0

    def wait(self):
        if not self.interval:
            return
        with self.lock:
            now = time.monotonic()
            slot = max(now, self.next_at)
            self.next_at = slot + self.interval
        delay = slot - time.monotonic()
        if delay > 0:
            time.sleep(delay)


class ProgressUI:
    """Четыре строки прогресса: регионы / кампании / округа / файлы, с текущим элементом справа.

    Без rich (или при --no-progress) все методы — пустышки.
    """
    LEVELS = [('region', 'Регионы'), ('election', 'Кампании'), ('district', 'Округа'), ('file', 'Файлы')]

    def __init__(self, enabled):
        self.progress = None
        self.console = None
        self.tasks = {}
        if enabled and Progress is not None:
            self.console = Console(highlight=False)
            self.progress = Progress(
                TextColumn('{task.description:<9}'),
                BarColumn(bar_width=20),
                MofNCompleteColumn(),
                TaskProgressColumn(),
                TextColumn('{task.fields[cur]}', style='cyan'),
                console=self.console, transient=False, expand=False,
            )

    def __enter__(self):
        if self.progress:
            self.progress.__enter__()
            for key, title in self.LEVELS:
                self.tasks[key] = self.progress.add_task(title, total=None, cur='')
        return self

    def __exit__(self, *exc):
        if self.progress:
            self.progress.__exit__(*exc)

    def start(self, level, total, cur=''):
        """Начать новый список на уровне: сбросить счётчик и задать общее число."""
        if self.progress:
            self.progress.reset(self.tasks[level], total=total or None, completed=0, cur=cur)

    def step(self, level, cur, done=None):
        """Перейти к следующему элементу уровня и показать его название."""
        if self.progress:
            kw = {'completed': done} if done is not None else {'advance': 1}
            self.progress.update(self.tasks[level], cur=self._short(cur), **kw)

    def clear(self, level):
        if self.progress:
            self.progress.update(self.tasks[level], cur='')

    def _short(self, text, limit=None):
        limit = limit or max(30, self.console.width - 60)
        text = ' '.join(str(text).split())
        return text if len(text) <= limit else text[:limit - 1] + '…'


class Portal:
    def __init__(self, base_url, cookies, retries=3, timeout=60, verify=True, rate=10.0, workers=4):
        self.url = base_url
        self.retries = retries
        self.timeout = timeout
        self.workers = workers
        self.limiter = RateLimiter(rate)
        self.sess = requests.Session()
        self.sess.headers.update(HEADERS)
        self.sess.cookies.update(cookies)
        self.sess.verify = verify
        adapter = requests.adapters.HTTPAdapter(pool_connections=workers, pool_maxsize=workers)
        self.sess.mount('https://', adapter)
        self.sess.mount('http://', adapter)
        self._lock = threading.Lock()
        self._pool = None
        if verify is False:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
            log.warning('ВНИМАНИЕ: проверка TLS-сертификата отключена (--insecure)')
        self.errors = 0
        self.requests = 0

    def _count(self, errors=0, reqs=0):
        with self._lock:
            self.errors += errors
            self.requests += reqs

    def _request(self, method, path, **kw):
        last = None
        for attempt in range(1, self.retries + 1):
            try:
                self.limiter.wait()
                self._count(reqs=1)
                t0 = time.time()
                r = self.sess.request(method, self.url + path, timeout=self.timeout, **kw)
                log.debug('%s %s -> %s (%d байт, %.2f с)', method, path, r.status_code,
                          len(r.content), time.time() - t0)
                if r.status_code == 401:
                    raise AuthError('Портал ответил 401 — cookie сессии протухла. Обновите cookie.txt.')
                if r.status_code == 403 and 'json' not in r.headers.get('content-type', ''):
                    raise AuthError('Портал ответил 403 (WAF): доступ запрещён по политике безопасности. '
                                    'Подождите несколько минут и уменьшите частоту запросов '
                                    '(--rate, --workers).')
                if r.status_code < 500:
                    return r
                last = f'HTTP {r.status_code}'
            except AuthError:
                raise
            except requests.exceptions.SSLError as e:
                if 'CERTIFICATE_VERIFY_FAILED' in str(e):
                    import ssl
                    raise AuthError(
                        'Не удалось проверить TLS-сертификат портала (%s). '
                        'Обычно это старый системный Python на macOS (LibreSSL). Варианты: '
                        'установить Python 3.10+ (brew install python) и создать .venv, '
                        'либо запустить с флагом --insecure.' % ssl.OPENSSL_VERSION)
                last = repr(e)
            except requests.RequestException as e:
                last = repr(e)
            log.warning('%s %s: %s, попытка %d/%d', method, path, last, attempt, self.retries)
            time.sleep(2 * attempt)
        log.error('%s %s: не удалось после %d попыток', method, path, self.retries)
        self._count(errors=1)
        return None

    def get_data(self, path, **params):
        """GET JSON-эндпоинта; возвращает содержимое поля data либо None."""
        r = self._request('GET', path, params=params or None)
        if r is None:
            return None
        try:
            body = r.json()
        except ValueError:
            log.error('%s: не JSON (HTTP %s): %s', path, r.status_code, r.text[:200])
            self._count(errors=1)
            return None
        if 'data' not in body:
            log.error('%s: %s', path, body.get('error') or body)
            self._count(errors=1)
            return None
        return body['data']

    def download(self, filename, contract_id, dest):
        """Скачивание файла выгрузки. Портал принимает только POST."""
        r = self._request('POST', f'files/{filename}/contract/{contract_id}/download')
        if r is None:
            return False
        ct = r.headers.get('content-type', '')
        if r.status_code != 200 or 'json' in ct:
            log.error('%s: HTTP %s %s', filename, r.status_code, r.text[:200])
            self._count(errors=1)
            return False
        tmp = dest + '.part'
        with open(tmp, 'wb') as f:
            f.write(r.content)
        os.replace(tmp, dest)
        return len(r.content)

    def download_all(self, filenames, contract_id, voting_path, stats, ui):
        """Скачивание файлов одного голосования в несколько потоков."""
        todo = []
        for filename in filenames:
            dest = os.path.join(voting_path, safe_name(filename, limit=200))
            if os.path.isfile(dest) and os.path.getsize(dest) > 0:
                stats['skipped'] += 1
                log.debug('        = %s уже есть', filename)
            else:
                todo.append((filename, dest))
        done = len(filenames) - len(todo)
        ui.step('file', '', done=done)
        if not todo:
            return

        def job(item):
            filename, dest = item
            size = self.download(filename, contract_id, dest)
            if size:
                log.info('        + %s (%d байт)', filename, size)
            return filename, size

        pool = self._pool or ThreadPoolExecutor(max_workers=self.workers)
        for filename, size in pool.map(job, todo):
            done += 1
            if size:
                stats['downloaded'] += 1
                stats['bytes'] += size
            ui.step('file', filename, done=done)

    def whoami(self):
        me = self.get_data('idm/user/me')
        if not me:
            raise AuthError('Не удалось получить данные пользователя — проверьте cookie.txt.')
        fio = ' '.join(filter(None, [me.get('lastName'), me.get('firstName'), me.get('middleName')]))
        mode = (self.get_data('mode') or {}).get('mode')
        lines = ['Пользователь: %s, роли: %s' % (fio, ', '.join(me.get('roles', [])))]
        for od in me.get('observerData', []):
            lines.append('  наблюдатель от: %s, доступ до %s' % (od.get('stakeholderName'),
                                                                  od.get('observerAdditionLastDate')))
        lines.append('Режим портала: %s' % mode)
        for line in lines:
            log.info(line)
        return lines


def run(portal, out_dir, wanted, stats, ui):
    regions = (portal.get_data('voting/regions') or {}).get('regions', [])
    log.info('Регионов на портале: %d', len(regions))
    unknown = wanted - {str(r['code']) for r in regions}
    if unknown:
        log.warning('Регионы не найдены на портале: %s', ', '.join(sorted(unknown)))
    if wanted:
        regions = [r for r in regions if str(r['code']) in wanted]

    ui.start('region', len(regions))
    for i, region in enumerate(regions):
        log.info('')
        log.info('%s - %s', region['code'], region['description'])
        ui.step('region', f'{region["code"]} - {region["description"]}', done=i)
        region_path = os.path.join(out_dir, f'{region["code"]} - {safe_name(region["description"])}')

        levels = (portal.get_data(f'elections/region/{region["code"]}/elections') or {}).get('elections', [])
        elections = [e for level in levels for e in level.get('elections', [])]
        ui.start('election', len(elections))
        for j, election in enumerate(elections):
            log.info('  %s', election['electionName'])
            ui.step('election', election['electionName'], done=j)
            election_path = os.path.join(region_path, safe_name(election['electionName']))

            districts = portal.get_data('elections/districts',
                                        regionCode=region['code'], electionId=election['electionId']) or []
            ui.start('district', len(districts))
            for k, district in enumerate(districts):
                log.info('    %s', district['name'])
                ui.step('district', district['name'], done=k)
                district_path = os.path.join(election_path, safe_name(district['name']))

                votings = (portal.get_data('statistics/voting',
                                           electionId=election['electionId'], districtId=district['id'])
                           or {}).get('votings', [])
                for voting in votings:
                    contract_id = voting.get('counters', {}).get('contractId')
                    if not contract_id:
                        log.debug('голосование без contractId: %s', voting)
                        continue
                    voting_path = os.path.join(district_path, contract_id)
                    filenames = portal.get_data(f'files/{contract_id}') or []
                    log.info('      %s %s (%s): файлов %d', voting.get('type'), contract_id,
                             voting.get('status'), len(filenames))
                    if not filenames:
                        continue

                    os.makedirs(voting_path, exist_ok=True)
                    write_info(election_path, election)
                    write_info(district_path, district)
                    write_info(voting_path, voting)

                    ui.start('file', len(filenames), cur=f'{voting.get("type")} {contract_id}')
                    portal.download_all(filenames, contract_id, voting_path, stats, ui)
                    ui.step('file', '', done=len(filenames))
                ui.step('district', '', done=k + 1)
            ui.step('election', '', done=j + 1)
        ui.step('region', '', done=i + 1)


def main():
    ap = argparse.ArgumentParser(description='Выгрузка файлов транзакций ДЭГ с портала наблюдения')
    ap.add_argument('regions', nargs='*', help='коды субъектов РФ (пусто — все доступные)')
    ap.add_argument('--out', help='папка выгрузки (по умолчанию — новая с текущим временем)')
    ap.add_argument('--cookie-file', default='cookie.txt', help='файл с cookie сессии (по умолчанию cookie.txt)')
    ap.add_argument('--test', action='store_true', help='использовать тестовый стенд ' + TEST_URL)
    ap.add_argument('--url', help='произвольный адрес API (переопределяет --test)')
    ap.add_argument('--retries', type=int, default=3, help='число повторов запроса при сбое')
    ap.add_argument('--rate', type=float, default=10.0,
                    help='максимум запросов к порталу в секунду (по умолчанию 10; 0 — без ограничения)')
    ap.add_argument('--workers', type=int, default=4,
                    help='число потоков скачивания файлов (по умолчанию 4)')
    ap.add_argument('--insecure', action='store_true', help='не проверять TLS-сертификат портала')
    ap.add_argument('--ca-bundle', help='путь к PEM-файлу с корневыми сертификатами')
    ap.add_argument('-v', '--verbose', action='store_true', help='подробный вывод в консоль (все запросы)')
    ap.add_argument('--no-progress', action='store_true', help='без прогресс-баров, построчный вывод')
    args = ap.parse_args()

    out_dir = args.out or getStrTimeNow('YMDHMS')
    os.makedirs(out_dir, exist_ok=True)
    ui = ProgressUI(enabled=not args.no_progress and sys.stdout.isatty())
    setup_logging(out_dir, args.verbose, console=ui.console)
    if ui.console:
        ui.console.print(f'[dim]Папка выгрузки: {out_dir}  (подробный лог: {os.path.join(out_dir, "download.log")})[/dim]')
    log.info('Запуск: %s', ' '.join(sys.argv))
    log.info('Папка выгрузки: %s', out_dir)

    base_url = args.url or (TEST_URL if args.test else PROD_URL)
    log.info('API: %s', base_url)
    cookies = load_cookies(args.cookie_file)
    log.debug('cookie: %s', ', '.join(cookies))
    exp = token_expiry(cookies)
    if exp:
        left = (exp - datetime.now()).total_seconds() / 60
        if left <= 0:
            log.error('stat_acc_t истёк %s — обновите cookie.txt', exp.strftime('%d.%m.%Y %H:%M'))
        else:
            (log.warning if left < 30 else log.info)(
                'stat_acc_t действует до %s (осталось %d мин)', exp.strftime('%H:%M'), left)

    verify = False if args.insecure else (args.ca_bundle or True)
    portal = Portal(base_url, cookies, retries=args.retries, verify=verify,
                    rate=args.rate, workers=args.workers)
    wanted = set(args.regions)
    stats = {'downloaded': 0, 'skipped': 0, 'bytes': 0}
    start = time.time()
    exit_code = 0

    try:
        for line in portal.whoami():
            if ui.console:
                ui.console.print(line, style='dim')
        with ui, ThreadPoolExecutor(max_workers=args.workers) as pool:
            portal._pool = pool
            run(portal, out_dir, wanted, stats, ui)
    except AuthError as e:
        log.error('%s', e)
        log.error('После исправления продолжить: %s --out "%s" %s',
                  os.path.basename(sys.argv[0]), out_dir, ' '.join(args.regions))
        exit_code = 2
    except KeyboardInterrupt:
        log.warning('Прервано пользователем. Докачать: --out "%s"', out_dir)
        exit_code = 130

    summary = ('Итого за %.0f сек: запросов %d, скачано %d файлов (%.1f МБ), пропущено (уже были) %d, ошибок %d'
               % (time.time() - start, portal.requests, stats['downloaded'], stats['bytes'] / 1024 / 1024,
                  stats['skipped'], portal.errors))
    log.info('')
    log.info(summary)
    if ui.console:
        ui.console.print(summary, style='bold' if not portal.errors else 'bold yellow', highlight=False)
    if portal.errors and not exit_code:
        exit_code = 1
    sys.exit(exit_code)


if __name__ == '__main__':
    main()