#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import argparse
from collections import Counter
import json
from pathlib import Path


def parse_expected(value: str) -> list[int]:
    values: list[int] = []
    for token in value.split(','):
        token = token.strip()
        if not token:
            continue
        if '-' in token:
            a, b = token.split('-', 1)
            start, end = int(a), int(b)
            if end < start:
                raise ValueError(f'invalid expected zoom range: {token}')
            values.extend(range(start, end + 1))
        else:
            values.append(int(token))
    return sorted(set(values))


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--plan', default='build/production/cm95_production_plan_v1.json')
    p.add_argument('--output', default='build/production/cm95_production_plan_audit_v1.json')
    p.add_argument('--expected-zooms', default='0-12')
    return p.parse_args()


def main():
    a = parse_args()
    plan = json.loads(Path(a.plan).read_text(encoding='utf-8'))
    jobs = plan.get('jobs') or []
    keys = [(int(j['z']), int(j['x']), int(j['y'])) for j in jobs]
    unique = set(keys)
    by_zoom = Counter(z for z, _, _ in keys)
    invalid = []
    for z, x, y in keys:
        n = 1 << z
        if z < 0 or x < 0 or y < 0 or x >= n or y >= n:
            invalid.append({'z': z, 'x': x, 'y': y})

    source_stats = plan.get('source_stats') or {}
    stats_mismatch = []
    for z, count in sorted(by_zoom.items()):
        declared = int((source_stats.get(str(z)) or {}).get('tile_count') or 0)
        if declared != count:
            stats_mismatch.append({'z': z, 'declared': declared, 'actual': count})

    expected = parse_expected(a.expected_zooms)
    planned_zooms = sorted(int(z) for z in (plan.get('zooms') or []))
    missing_plan_zooms = sorted(set(expected) - set(planned_zooms))
    unexpected_plan_zooms = sorted(set(planned_zooms) - set(expected))
    missing_job_zooms = sorted(z for z in expected if by_zoom.get(z, 0) == 0)

    failures = []
    if len(unique) != len(keys):
        failures.append('duplicate_jobs')
    if invalid:
        failures.append('invalid_xyz')
    if int(plan.get('job_count') or 0) != len(keys):
        failures.append('job_count_mismatch')
    if stats_mismatch:
        failures.append('source_stats_mismatch')
    if missing_plan_zooms:
        failures.append('missing_plan_zooms')
    if unexpected_plan_zooms:
        failures.append('unexpected_plan_zooms')
    if missing_job_zooms:
        failures.append('missing_job_zooms')

    result = {
        'schema': 'navsea.cm95.production-plan-audit/v1.1',
        'plan_id': plan.get('plan_id'),
        'release_id': plan.get('release_id'),
        'dataset_id': plan.get('dataset_id'),
        'builder_contract_version': plan.get('builder_contract_version'),
        'expected_zooms': expected,
        'planned_zooms': planned_zooms,
        'job_count': len(keys),
        'unique_job_count': len(unique),
        'tile_count_by_zoom': {str(k): v for k, v in sorted(by_zoom.items())},
        'missing_plan_zooms': missing_plan_zooms,
        'unexpected_plan_zooms': unexpected_plan_zooms,
        'missing_job_zooms': missing_job_zooms,
        'invalid_xyz': invalid,
        'source_stats_mismatch': stats_mismatch,
        'failures': failures,
        'overall': 'PASS' if not failures else 'FAIL',
    }
    out = Path(a.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')

    print(f"Plan ID        : {result['plan_id']}")
    print(f"Release ID     : {result['release_id']}")
    if result.get('builder_contract_version'):
        print(f"Builder        : {result['builder_contract_version']}")
    print(f"Expected zooms : {a.expected_zooms}")
    print(f"Jobs           : {result['job_count']}")
    for z in expected:
        stat = source_stats.get(str(z)) or {}
        enum_scales = stat.get('enumeration_scales') or []
        ladder = stat.get('scale_ladder') or []
        print(f"z{z} tiles      : {by_zoom.get(z, 0)}  enum={enum_scales}  fallback={ladder}")
    print(f"Duplicate jobs : {len(keys)-len(unique)}")
    print(f"Invalid XYZ    : {len(invalid)}")
    print(f"Missing zooms  : {len(missing_plan_zooms) + len(missing_job_zooms)}")
    print(f"OVERALL        : {result['overall']}")


if __name__ == '__main__':
    main()