#!/usr/bin/env python3
import argparse
import os
import sys
from typing import List

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

from analytics import AnalyticsRepository, AnalyticsService, _is_model_traffic_line, _parse_traffic_size
from config import ANALYTICS_TRAFFIC_ROOT, MODEL_TRAFFIC_THRESHOLD, MONITORING_DB_PATH


def _split_packages(raw: str) -> List[str]:
    return [item.strip() for item in str(raw or "").split(",") if item.strip()]


def _load_all_packages(service: AnalyticsService, tag: str = "", limit: int = 0) -> List[str]:
    normalized_tag = str(tag or "").strip()
    with service.repo._connect() as connection:
        summaries = service.repo._list_catalog_summaries(connection)
    packages = []
    for item in sorted(summaries, key=lambda value: str(value.get("package_name") or "")):
        if normalized_tag and normalized_tag not in (item.get("incremental_batch_tags") or []):
            continue
        package_name = str(item.get("package_name") or "").strip()
        if package_name:
            packages.append(package_name)
    return packages[:limit] if limit > 0 else packages


def _update_model_stats(repo: AnalyticsRepository, package_name: str, flow_count: int, traffic_bytes: int) -> bool:
    with repo._write_lock, repo._connect() as connection:
        latest = connection.execute(
            """
            SELECT package_name, batch_tag, run_kind, attempt
              FROM collection_task
             WHERE package_name = ?
             ORDER BY
                  datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC,
                  attempt DESC,
                  batch_tag DESC,
                  run_kind DESC
             LIMIT 1
            """,
            (package_name,),
        ).fetchone()
        if not latest:
            return False
        connection.execute(
            """
            UPDATE collection_task
               SET model_flow_count = ?,
                   model_traffic_bytes = ?
             WHERE package_name = ?
               AND batch_tag = ?
               AND run_kind = ?
               AND attempt = ?
            """,
            (
                flow_count,
                traffic_bytes,
                latest["package_name"],
                latest["batch_tag"],
                latest["run_kind"],
                latest["attempt"],
            ),
        )
        return True


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Backfill model traffic stats for all or selected packages."
    )
    parser.add_argument("--db-path", default=MONITORING_DB_PATH)
    parser.add_argument("--traffic-root", default=ANALYTICS_TRAFFIC_ROOT)
    parser.add_argument("--packages", default="", help="Comma-separated package names. Defaults to all packages.")
    parser.add_argument("--tag", default="", help="Only process packages whose incremental_batch_tag contains this tag.")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--auto-promote", action="store_true",
                        help="Auto-set model_eligible=1 and collection_status=pending if model_flow_count > 0.")
    args = parser.parse_args()

    scanned = 0
    matched = 0
    backfilled = 0
    promoted = 0
    missing_task = 0
    service = AnalyticsService(
        db_path=args.db_path,
        traffic_root=args.traffic_root,
        start_worker=False,
    )
    try:
        packages = _split_packages(args.packages) or _load_all_packages(service, args.tag, args.limit)

        for package_name in packages:
            scanned += 1
            traffic_files = service._find_traffic_files(package_name)
            if not traffic_files:
                continue
            model_flow_count = 0
            model_traffic_bytes = 0
            for file_path in traffic_files:
                try:
                    with open(file_path, "r", encoding="utf-8", errors="ignore") as handle:
                        for raw_line in handle:
                            line = raw_line.strip()
                            if not line:
                                continue
                            fields = [item.strip() for item in line.split(",")]
                            if not fields or fields[0].strip() != package_name:
                                continue
                            if _is_model_traffic_line(fields):
                                model_flow_count += 1
                                model_traffic_bytes += _parse_traffic_size(fields[6])
                except OSError:
                    continue
            if model_flow_count == 0:
                continue
            matched += 1
            eligible = 1 if 0 < model_flow_count <= MODEL_TRAFFIC_THRESHOLD else 0
            if args.dry_run:
                print(f"DRY-RUN {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}")
                continue
            if not _update_model_stats(service.repo, package_name, model_flow_count, model_traffic_bytes):
                missing_task += 1
                continue
            backfilled += 1
            if args.auto_promote and eligible:
                promoted += service.set_packages_pending([package_name], reason="model_traffic_backfill")
            print(f"backfilled {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}")
    finally:
        service.close()

    print(f"done scanned={scanned} matched={matched} backfilled={backfilled} promoted={promoted} missing_task={missing_task}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())