# Выгрузка файлов транзакций ДЭГ с портала наблюдения 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      — все параметры
#
# Лог пишется в консоль и в файл download.log внутри папки выгрузки.

import argparse
import base64
import json
import logging
import os
import re
import sys
import time
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 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, delay=0.3):
        self.url = base_url
        self.retries = retries
        self.timeout = timeout
        self.delay = delay  # пауза между запросами, чтобы не провоцировать WAF портала
        self.sess = requests.Session()
        self.sess.headers.update(HEADERS)
        self.sess.cookies.update(cookies)
        self.sess.verify = verify
        if verify is False:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
            log.warning('ВНИМАНИЕ: проверка TLS-сертификата отключена (--insecure)')
        self.errors = 0
        self.requests = 0

    def _request(self, method, path, **kw):
        last = None
        for attempt in range(1, self.retries + 1):
            try:
                time.sleep(self.delay)
                self.requests += 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): доступ запрещён по политике безопасности. '
                                    'Подождите несколько минут и уменьшите частоту запросов (--delay).')
                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.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.errors += 1
            return None
        if 'data' not in body:
            log.error('%s: %s', path, body.get('error') or body)
            self.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.errors += 1
            return False
        tmp = dest + '.part'
        with open(tmp, 'wb') as f:
            f.write(r.content)
        os.replace(tmp, dest)
        return True

    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}')
                    for n, filename in enumerate(filenames):
                        dest = os.path.join(voting_path, safe_name(filename, limit=200))
                        ui.step('file', filename, done=n)
                        if os.path.isfile(dest) and os.path.getsize(dest) > 0:
                            stats['skipped'] += 1
                            log.debug('        = %s уже есть', filename)
                            continue
                        if portal.download(filename, contract_id, dest):
                            size = os.path.getsize(dest)
                            stats['downloaded'] += 1
                            stats['bytes'] += size
                            log.info('        + %s (%d байт)', filename, size)
                    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('--delay', type=float, default=0.3, help='пауза между запросами, сек (по умолчанию 0.3)')
    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, delay=args.delay)
    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:
            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()