import requests
import datetime
import ipaddress
import os
import argparse
import logging
import socket
import sys

import db
from storage import StorageError, file_lock, load_json, save_json_atomic

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# Каталог данных: по умолчанию рядом с кодом, в контейнере - том (RIPE_DATA_DIR=/data)
DATA_DIR = os.environ.get("RIPE_DATA_DIR", BASE_DIR)
CONFIG_FILE = os.path.join(DATA_DIR, "config.json")
DB_FILE = os.path.join(DATA_DIR, "ripe.db")  # собранные адреса (SQLite)
# Старые JSON-хранилища: используются только для однократного импорта в DB_FILE (см. db.py)
DATA_FILE = os.path.join(DATA_DIR, "data.json")
FQDN_DATA_FILE = os.path.join(DATA_DIR, "fqdn_data.json")
STATUS_FILE = os.path.join(DATA_DIR, "status.json")  # пишет collector_daemon, читает API (/health)
COLLECT_REQUEST_FILE = os.path.join(DATA_DIR, "collect_request.json")  # API -> демон: POST /collect
# Резервные копии базы (задание backup в демоне); по умолчанию на том же томе, каталог можно вынести
BACKUP_DIR = os.environ.get("RIPE_BACKUP_DIR", os.path.join(DATA_DIR, "backups"))
RESTORE_FILE = os.path.join(DATA_DIR, "last_restore.json")  # след автовосстановления базы (читает /health)
RECREATED_FILE = os.path.join(DATA_DIR, "db_recreated.json")  # база пересоздана после порчи без копий (см. db.recreated_pending)
BASE_URL = "https://stat.ripe.net/data/announced-prefixes/data.json"

DEFAULT_TTL_DAYS = 90
DEFAULT_BACKUP_KEEP = 7  # сколько копий базы хранить
DEFAULT_CHANGES_RETENTION_DAYS = 30  # срок хранения журнала изменений для /addresses/diff
# Демон сверяет расписание с config.json и пишет heartbeat в status.json с этим периодом;
# API считает демон живым, пока heartbeat не старше STATUS_STALE_AFTER секунд
SYNC_INTERVAL = 30
STATUS_STALE_AFTER = 120
COLLECT_TYPES = ("asn", "fqdn")
TRIGGER_POLL_INTERVAL = 5  # как часто демон проверяет запросы на немедленный сбор, секунды

logger = logging.getLogger(__name__)


def load_full_config():
    return load_json(CONFIG_FILE, {"asns": [], "fqdns": []})


def update_config(mutator):
    """Read-modify-write config.json под блокировкой с атомарной записью.

    mutator(config) изменяет словарь на месте и возвращает результат, который отдаётся вызывающему.
    Единственный путь записи конфигурации для CLI и API.
    """
    with file_lock(CONFIG_FILE):
        config = load_full_config()
        result = mutator(config)
        save_json_atomic(CONFIG_FILE, config)
        return result


def add_to_config_list(key, value):
    """Добавляет значение в список конфигурации ("asns"/"fqdns"). Возвращает (добавлено, актуальный список)."""
    def mutate(config):
        items = config.setdefault(key, [])
        if value in items:
            return False, items
        items.append(value)
        return True, items
    return update_config(mutate)


def remove_from_config_list(key, value):
    """Убирает значение из списка конфигурации. Возвращает (удалено, актуальный список)."""
    def mutate(config):
        items = config.setdefault(key, [])
        if value not in items:
            return False, items
        items.remove(value)
        return True, items
    return update_config(mutate)


def request_collection(types):
    """API: ставит запрос на немедленный сбор. Повторные запросы объединяются. Возвращает список типов."""
    with file_lock(COLLECT_REQUEST_FILE):
        pending = set(load_json(COLLECT_REQUEST_FILE, {}).get("types", []))
        merged = sorted(pending | set(types))
        save_json_atomic(COLLECT_REQUEST_FILE, {"types": merged, "requested_at": datetime.datetime.now().isoformat()})
        return merged


def pop_collection_requests():
    """Демон: забирает и удаляет запрос на немедленный сбор. Возвращает список типов (возможно пустой)."""
    if not os.path.exists(COLLECT_REQUEST_FILE):
        return []
    with file_lock(COLLECT_REQUEST_FILE):
        types = load_json(COLLECT_REQUEST_FILE, {}).get("types", [])
        if os.path.exists(COLLECT_REQUEST_FILE):
            os.remove(COLLECT_REQUEST_FILE)
    return [t for t in types if t in COLLECT_TYPES]


def _is_global(address):
    """Публичный адрес (не loopback, не частный, не 0.0.0.0, не link-local и т. п.)."""
    try:
        return ipaddress.ip_address(address).is_global
    except ValueError:
        return False


class CIDRCollector:
    def __init__(self):
        self.config = load_full_config()
        self.asns = self.config.get("asns", [])
        self.ttl_days = self.config.get("ttl_days", DEFAULT_TTL_DAYS)

    def add_asn(self, asn):
        added, self.asns = add_to_config_list("asns", asn)
        logger.info("ASN %s added." if added else "ASN %s already in list.", asn)

    def remove_asn(self, asn):
        removed, self.asns = remove_from_config_list("asns", asn)
        logger.info("ASN %s removed." if removed else "ASN %s not found in list.", asn)

    def list_asns(self):
        print("Current ASNs:", self.asns)

    def fetch_prefixes(self, asn):
        params = {'resource': f'AS{asn}'}
        try:
            response = requests.get(BASE_URL, params=params, timeout=10)
            response.raise_for_status()
            data = response.json()

            prefixes = []
            if 'data' in data and 'prefixes' in data['data']:
                for item in data['data']['prefixes']:
                    if 'prefix' in item:
                        prefixes.append(item['prefix'])
            return prefixes
        except Exception as e:
            logger.error("Error fetching data for AS%s: %s", asn, e)
            return None

    def run_collection(self):
        logger.info("Starting ASN CIDR collection...")
        # Сетевые запросы выполняем до транзакции, чтобы держать блокировку записи минимально
        fetched = {}
        for asn in self.asns:
            logger.info("Processing AS%s...", asn)
            prefixes = self.fetch_prefixes(asn)
            if prefixes is not None:
                fetched[str(asn)] = set(prefixes)

        now = datetime.datetime.now()
        with db.session() as conn, db.transaction(conn):
            # Конфиг перечитываем внутри транзакции: источник, удалённый во время сбора, не воскресает
            configured = {str(a) for a in load_full_config().get("asns", [])}
            for str_asn, prefixes in fetched.items():
                if str_asn not in configured:
                    continue
                added, removed = db.merge_source(conn, "asn", str_asn, prefixes, now, self.ttl_days)
                logger.info("AS%s: +%d / -%d prefixes", str_asn, len(added), len(removed))
            swept = db.sweep_unconfigured(conn, "asn", configured, now, self.ttl_days)
            if swept:
                logger.info("Expired %d prefixes of unconfigured ASNs", swept)
            db.prune_changes(conn, now, self.config.get("changes_retention_days", DEFAULT_CHANGES_RETENTION_DAYS))
        logger.info("CIDR data saved to %s", DB_FILE)
        with db.session() as conn:
            db.settle_recreated(conn)


class FQDNCollector:
    def __init__(self):
        self.config = load_full_config()
        self.fqdns = self.config.get("fqdns", [])
        self.ttl_days = self.config.get("ttl_days", DEFAULT_TTL_DAYS)
        # По умолчанию в список попадают только глобальные адреса; true - для внутренних имён
        self.allow_non_global_ips = bool(self.config.get("allow_non_global_ips", False))

    def add_fqdn(self, fqdn):
        added, self.fqdns = add_to_config_list("fqdns", fqdn)
        logger.info("FQDN %s added." if added else "FQDN %s already in list.", fqdn)

    def remove_fqdn(self, fqdn):
        removed, self.fqdns = remove_from_config_list("fqdns", fqdn)
        logger.info("FQDN %s removed." if removed else "FQDN %s not found in list.", fqdn)

    def list_fqdns(self):
        print("Current FQDNs:", self.fqdns)

    def resolve_fqdn(self, fqdn):
        try:
            # Family 0 - получаем и IPv4 (A), и IPv6 (AAAA)
            results = socket.getaddrinfo(fqdn, None)
            # result[4] - sockaddr, для IP-протоколов индекс 0 - строка с адресом; зона IPv6 (%eth0) не нужна
            addresses = {result[4][0].split("%")[0] for result in results}
        except socket.gaierror as e:
            logger.error("Error resolving %s: %s", fqdn, e)
            return []
        if self.allow_non_global_ips:
            return list(addresses)
        kept = {a for a in addresses if _is_global(a)}
        if kept != addresses:
            logger.warning("%s: ignoring non-global addresses: %s", fqdn, ", ".join(sorted(addresses - kept)))
        return list(kept)

    def run_collection(self):
        logger.info("Starting FQDN IP collection...")
        fetched = {}
        for fqdn in self.fqdns:
            logger.info("Processing %s...", fqdn)
            resolved_ips = self.resolve_fqdn(fqdn)
            if not resolved_ips:
                # Ошибка DNS: TTL не применяем, чтобы сбой не удалил данные
                logger.warning("No IPs resolved for %s", fqdn)
                continue
            fetched[fqdn] = set(resolved_ips)

        now = datetime.datetime.now()
        with db.session() as conn, db.transaction(conn):
            configured = set(load_full_config().get("fqdns", []))
            for fqdn, ips in fetched.items():
                if fqdn not in configured:
                    continue
                added, removed = db.merge_source(conn, "fqdn", fqdn, ips, now, self.ttl_days)
                logger.info("%s: +%d / -%d IPs", fqdn, len(added), len(removed))
            swept = db.sweep_unconfigured(conn, "fqdn", configured, now, self.ttl_days)
            if swept:
                logger.info("Expired %d IPs of unconfigured FQDNs", swept)
            db.prune_changes(conn, now, self.config.get("changes_retention_days", DEFAULT_CHANGES_RETENTION_DAYS))
        logger.info("FQDN data saved to %s", DB_FILE)
        with db.session() as conn:
            db.settle_recreated(conn)


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    parser = argparse.ArgumentParser(description="Collector for RIPE AS CIDRs and FQDN IPs")
    subparsers = parser.add_subparsers(dest="command")

    # Command: run (default)
    parser_run = subparsers.add_parser("run", help="Run the collection process")
    parser_run.add_argument("--mode", choices=["asn", "fqdn", "all"], default="all", help="Collection mode: asn, fqdn, or all (default)")

    # ASN Commands
    parser_add = subparsers.add_parser("add", help="Add an ASN")
    parser_add.add_argument("asn", type=int, help="ASN to add")

    parser_remove = subparsers.add_parser("remove", help="Remove an ASN")
    parser_remove.add_argument("asn", type=int, help="ASN to remove")

    # FQDN Commands
    parser_add_fqdn = subparsers.add_parser("add-fqdn", help="Add an FQDN")
    parser_add_fqdn.add_argument("fqdn", type=str, help="FQDN to add")

    parser_remove_fqdn = subparsers.add_parser("remove-fqdn", help="Remove an FQDN")
    parser_remove_fqdn.add_argument("fqdn", type=str, help="FQDN to remove")

    # Command: list
    subparsers.add_parser("list", help="List ASNs and FQDNs")

    args = parser.parse_args()

    try:
        asn_collector = CIDRCollector()
        fqdn_collector = FQDNCollector()

        if args.command == "add":
            asn_collector.add_asn(args.asn)
        elif args.command == "remove":
            asn_collector.remove_asn(args.asn)
        elif args.command == "add-fqdn":
            fqdn_collector.add_fqdn(args.fqdn)
        elif args.command == "remove-fqdn":
            fqdn_collector.remove_fqdn(args.fqdn)
        elif args.command == "list":
            asn_collector.list_asns()
            fqdn_collector.list_fqdns()
        elif args.command == "run":
            mode = args.mode
            if mode in ["asn", "all"]:
                asn_collector.run_collection()

            if mode == "all":
                print("-" * 20)

            if mode in ["fqdn", "all"]:
                fqdn_collector.run_collection()
        else:
            parser.print_help()
    except StorageError as e:
        logger.error("Aborted: %s", e)
        sys.exit(1)


if __name__ == "__main__":
    main()