"""Operational storage, account checks, cycle decisions and non-overwriting archive."""
import contextlib
import fcntl
import hashlib
import json
import os
from pathlib import Path
import pwd
import socket
import subprocess
import sys
import rcms as core

MODEL_NAMES = {'wrf': 'CLWRF', 'regcm': 'RegCM'}


def required_products(model):
    if model == 'wrf':
        return [v + '.nc' for v in 'slp u10m v10m u v w rh2 rh tsk t2 tk theta geopt q2 qvapor rainc rainnc olr xtrm'.split()]
    return ([f'daily/{v}.nc' for v in 'pr prmax tas tasmax tasmin'.split()] +
            [f'3hourly/{v}.nc' for v in 'pr prc uas vas ts huss hurs psl evspsbl zmla hfss hfls'.split()] +
            [f'6hourly/{v}.nc' for v in 'ua va wa ta hus rh hgt mslp ts rlut prw'.split()])


def storage(cfg):
    path = Path(cfg.get('runs_root', 'runs'))
    return (path if path.is_absolute() else core.HERE / path).resolve()


def archive(cfg, model, cycle):
    return Path(cfg['current_result_root']) / MODEL_NAMES[model] / cycle


def catalog(path, model):
    missing = [n for n in required_products(model) if not (Path(path)/n).is_file() or (Path(path)/n).stat().st_size == 0]
    return {'path':str(path), 'complete_catalog':not missing, 'missing':missing,
            'exists':os.path.lexists(path)}


def username():
    return pwd.getpwuid(os.geteuid()).pw_name


def require_user(cfg):
    owner = cfg.get('run_user')
    if owner:
        core.need(username() == owner,
                  'Run this write/submission command as %s (current account: %s). No automatic sudo/chmod/chown.' % (owner, username()))


def nearest_existing(path):
    path = Path(path).resolve()
    while not path.exists():
        core.need(path.parent != path, 'Cannot locate existing storage ancestor')
        path = path.parent
    return path


def write_preflight(cfg):
    require_user(cfg)
    paths = [storage(cfg)]
    if cfg.get('operational_layout'):
        paths += [Path(base)/'wdir' for base in cfg['legacy_dirs'].values()]
        paths += [Path(cfg['regcm_dataset'])/'CFS']
    if cfg.get('current_result_root'):
        paths += [Path(cfg['current_result_root'])/m for m in MODEL_NAMES.values()]
        paths += [Path(cfg['archive_versions_root'])]
    for path in paths:
        ancestor = nearest_existing(path)
        core.need(ancestor.is_dir() and os.access(ancestor, os.W_OK | os.X_OK),
                  'Account cannot write storage: ' + str(ancestor))


@contextlib.contextmanager
def cycle_lock(cfg, cycle):
    core.date10(cycle)
    root = storage(cfg) / '.locks'
    root.mkdir(parents=True, exist_ok=True)
    # Persistent lock file, released by close. Never removed or truncated.
    with (root/(cycle + '.lock')).open('a') as f:
        try:
            fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise RuntimeError('Another controller is handling cycle ' + cycle)
        try:
            yield
        finally:
            fcntl.flock(f, fcntl.LOCK_UN)


def queue(cfg):
    owner = cfg.get('run_user', username())
    p = subprocess.run(['squeue','-h','-u',owner,'-o','%i|%j|%T'], stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE, universal_newlines=True)
    core.need(p.returncode == 0, 'Cannot verify active jobs: ' + p.stderr.strip())
    rows = {}
    for line in p.stdout.splitlines():
        fields = line.strip().split('|')
        core.need(len(fields) == 3, 'Unrecognized squeue output: ' + line)
        rows[fields[0]] = {'job_id':fields[0], 'name':fields[1], 'state':fields[2]}
    return rows


def attempts(cfg, cycle):
    return sorted((storage(cfg)/cycle).glob('*/manifest.json'), reverse=True)


def known_jobs(root, model=None):
    result = []
    for p in sorted(Path(root).glob('*.job.json')):
        row = json.loads(p.read_text())
        if model is None or row['model'] == model:
            result.append(row)
    return result


def active_for_cycle(cfg, cycle, model, rows):
    found = []
    prefix = 'rcms-' + model + '-' + cycle + '-'
    legacy_name = {'wrf':'clWRF_run','regcm':'RegCM_run'}[model]
    for row in rows.values():
        if row['name'].startswith(prefix) or row['name'] == legacy_name:
            found.append(row)
    # Also recognize older package jobs, but never trust job ID alone after reuse.
    for m in attempts(cfg, cycle):
        for job in known_jobs(m.parent, model):
            row = rows.get(job['job_id'])
            if row and row['name'] == job.get('slurm_name') and row not in found:
                found.append(row)
    return found


def legacy_state(cfg, cycle, model):
    base = Path(cfg['legacy_dirs'][model])
    marker = 'SUCCESS COMPLETE WRF' if model == 'wrf' else 'RegCM V4 simulation successfully reached end'
    logs = ([base/'wdir'/cycle/'rsl.error.0000',base/'wdir2'/cycle/'rsl.error.0000'] if model == 'wrf'
            else [base/('logs_run_'+cycle)])
    success_log = next((p for p in logs if p.is_file() and marker in p.read_text(errors='replace')), None)
    for branch in ['wdir2','wdir']:
        stage = base/branch/cycle
        choices = [stage/'output',stage] if model == 'wrf' else [stage/'output']
        for raw in choices:
            pattern = 'wrfout_d01_*' if model == 'wrf' else 'd01_ATM.*.nc'
            if any(raw.glob(pattern)) and (success_log or branch == 'wdir2'):
                return {'state':'RAW_CANDIDATE', 'raw':str(raw),
                        'evidence':str(success_log) if success_log else 'legacy wdir2 (must validate complete time coverage)'}
    occupied = [base/b/cycle for b in ['wdir','wdir2'] if (base/b/cycle).exists()]
    if model == 'regcm' and (base/'Dataset/CFS'/cycle).exists():
        occupied.append(base/'Dataset/CFS'/cycle)
    return {'state':'UNCONFIRMED' if occupied else 'EMPTY', 'paths':[str(x) for x in occupied]}


def decide(cfg, cycle, model, rows, rebuild=False):
    running = active_for_cycle(cfg, cycle, model, rows)
    if running:
        return {'action':'WAIT', 'reason':'Active job; refusing duplicate submission', 'jobs':running}
    stored = catalog(archive(cfg,model,cycle),model)
    if stored['complete_catalog'] and not rebuild:
        return {'action':'SKIP', 'reason':'Historical archive has all nonempty expected products; no rerun',
                'archive':stored, 'verification':'catalog only, not a new physical/NetCDF validation'}
    if not rebuild:
        for p in attempts(cfg, cycle):
            root = p.parent
            for stage in sorted((root/model).glob('post-*'), reverse=True):
                if (stage/'SUCCESS.json').is_file():
                    return {'action':'PUBLISH', 'run':str(root), 'post':str(stage), 'archive':stored}
            if (root/model/'integration/SUCCESS.json').is_file():
                return {'action':'POST', 'run':str(root), 'archive':stored}
        old = legacy_state(cfg,cycle,model)
        if old['state'] == 'RAW_CANDIDATE':
            return {'action':'ADOPT_POST', 'legacy':old, 'archive':stored}
        if old['state'] == 'UNCONFIRMED':
            return {'action':'BLOCKED', 'legacy':old,
                    'reason':'Old state has no confirmed integration. Inspect the existing cycle before running; automatic overwrite is refused.'}
        if any(model in json.loads(p.read_text()).get('models', []) for p in attempts(cfg,cycle)):
            return {'action':'BLOCKED', 'reason':'A previous model attempt is incomplete/failed; inspect its logs. No automatic rerun.'}
    return {'action':'NEW', 'archive':stored,
            'reason':'New attempt; previous files remain at their existing paths'}


def sha256(path):
    h=hashlib.sha256()
    with Path(path).open('rb') as f:
        for block in iter(lambda:f.read(1024*1024),b''):
            h.update(block)
    return h.hexdigest()


def publish(root, model, post):
    root, post = Path(root).resolve(), Path(post).resolve()
    man=json.loads((root/'manifest.json').read_text())
    cfg=man['config']
    require_user(cfg)
    core.need(post.parent == root/model and post.name.startswith('post-'), 'Post stage must belong to this run/model')
    success=post/'SUCCESS.json'
    core.need(success.is_file(), 'Cannot archive unverified postprocessing')
    record=json.loads(success.read_text())
    products=post/'products'
    core.need(Path(record['products']).resolve() == products.resolve(), 'Product path mismatch')
    core.need(catalog(products,model)['complete_catalog'], 'Missing expected products; archive refused')
    dest=archive(cfg,model,man['cycle'])
    if (post/'ARCHIVED.json').is_file():
        saved=json.loads((post/'ARCHIVED.json').read_text())
        version=Path(saved['version'])
        core.need(dest.exists() and dest.resolve() == version.resolve(), 'Archive link differs from recorded publication')
        for name, meta in saved['files'].items():
            core.need((version/name).stat().st_size == meta['size'] and sha256(version/name) == meta['sha256'], 'Archived file changed: '+name)
        return saved
    version=Path(cfg['archive_versions_root'])/MODEL_NAMES[model]/man['cycle']/(root.name+'-'+post.name)
    receipt_path=post/('archive-'+version.name+'.json')
    previous_receipt=receipt_path if receipt_path.is_file() else version/'RCMS_RECEIPT.json'
    if previous_receipt.is_file():
        receipt=json.loads(previous_receipt.read_text())
        core.need(receipt['run']==str(root) and receipt['post']==str(post), 'Archive provenance mismatch')
        for name,meta in receipt['files'].items():
            core.need(sha256(version/name)==meta['sha256'] and sha256(products/name)==meta['sha256'], 'Archive/source changed')
    else:
        if version.exists():
            # A partially copied version is retained; retry creates another version.
            version=version.with_name(version.name+'-'+core.stamp())
        version.mkdir(parents=True)
        receipt={'cycle':man['cycle'],'end':(json.loads((post/'RAW_VALIDATED.json').read_text())['end'] if (post/'RAW_VALIDATED.json').exists() else man['end']),'model':model,'run':str(root),'post':str(post),
                 'version':str(version),'destination':str(dest),'account':username(),'host':socket.gethostname(),
                 'time':core.stamp(),'files':{}}
        for name in required_products(model):
            source, target=products/name,version/name
            target.parent.mkdir(parents=True,exist_ok=True)
            core.copy_new(source,target)
            digest=sha256(source)
            core.need(sha256(target)==digest,'Archive copy checksum mismatch: '+name)
            receipt['files'][name]={'size':target.stat().st_size,'sha256':digest}
        core.put(post/('archive-'+version.name+'.json'),receipt)
    dest.parent.mkdir(parents=True,exist_ok=True)
    if os.path.lexists(dest):
        core.need(dest.resolve()==version.resolve(),
                  'Existing historical result preserved: %s. New verified version is at %s; no overwrite.' % (dest,version))
    else:
        # Exclusive symlink creation: complete copied version becomes visible atomically.
        # A concurrent publisher cannot overwrite an existing destination.
        dest.symlink_to(version)
    core.put(post/'ARCHIVED.json',receipt)
    return receipt


def managed_run(value):
    root=Path(value).resolve()
    man=json.loads((root/'manifest.json').read_text())
    core.need(root.parent.parent == storage(man['config']), 'Run is outside its configured storage root')
    core.need(root.parent.name==man['cycle'], 'Run cycle does not match manifest')
    return root,man


def cancel(value, execute=False):
    root,man=managed_run(value)
    rows=queue(man['config'])
    jobs=[]
    for j in known_jobs(root):
        if j['job_id'] in rows:
            core.need(rows[j['job_id']]['name']==j.get('slurm_name'), 'Job ID/name mismatch; refusing cancellation')
            jobs.append(j['job_id'])
    print('Active jobs belonging to this run:', ', '.join(jobs) or '(none)')
    if execute and jobs:
        require_user(man['config'])
        with cycle_lock(man['config'],man['cycle']):
            p=subprocess.run(['scancel']+jobs,stdout=subprocess.PIPE,stderr=subprocess.PIPE,universal_newlines=True)
            core.put(root/('cancel-'+core.stamp()+'.json'),{'jobs':jobs,'returncode':p.returncode,'stdout':p.stdout,'stderr':p.stderr,'account':username()})
            core.need(p.returncode==0,'scancel failed: '+p.stderr)
    elif not execute:
        print('PLAN ONLY. Add --execute to cancel these jobs. Files are retained.')


def doctor(cfg):
    print('Current account:',username(),'; required execution account:',cfg.get('run_user'))
    for label,path in [('Run storage',storage(cfg)),('Historical results',Path(cfg['current_result_root'])),
                       ('Archived versions',Path(cfg['archive_versions_root']))]:
        p=nearest_existing(path)
        print(label,':',path,'; nearest existing owner:',pwd.getpwuid(p.stat().st_uid).pw_name,
              '; writable by current account:',os.access(p,os.W_OK|os.X_OK))
    for model in MODEL_NAMES:
        done=Path(cfg['legacy_dirs'][model])/'done'
        print(MODEL_NAMES[model],'done ->',done.resolve(),'; expected:',Path(cfg['current_result_root'])/MODEL_NAMES[model])
        core.need(done.resolve()==(Path(cfg['current_result_root'])/MODEL_NAMES[model]).resolve(), 'Legacy done target differs from configured archive')


def plan(cfg, cycle, models, rows):
    import recovery
    core.date10(cycle)
    decisions = {m: decide(cfg, cycle, m, rows) for m in models}
    for model, decision in decisions.items():
        if decision['action'] in ('BLOCKED','ADOPT_POST','NEW'):
            candidate = recovery.discover(cfg,cycle,model)
            if candidate:
                decisions[model] = candidate
    return decisions


def new_run(cfg, cycle, models, decisions=None):
    import shutil
    core.validate_config(cfg, models)
    decisions = decisions or {}
    if models and all(decisions.get(m, {}).get('action') == 'ADOPT_POST' for m in models):
        end, entries = None, []
    else:
        end, entries = core.inventory(cfg, cycle)
    cfg = dict(cfg, runs_root=str(storage(cfg)))  # Freeze before copying runtime code.
    root = storage(cfg) / cycle / core.stamp()
    root.mkdir(parents=True)
    runtime = root / 'runtime'
    runtime.mkdir()
    for name in ('rcms.py', 'operations.py', 'recovery.py'):
        core.copy_new(core.HERE / name, runtime / name)
    shutil.copytree(core.HERE / 'templates', runtime / 'templates')
    man = {'cycle':cycle, 'end':end, 'config':cfg, 'models':models, 'inputs':entries,
           'python':sys.executable, 'account':username(),
           'recovery':{m:decisions[m] for m in models if decisions.get(m,{}).get('action') in ('RESUME','ADOPT_POST')},
           'tasks':{m:cfg['slurm'][m]['nodes'] * cfg['slurm'][m]['tasks_per_node'] for m in MODEL_NAMES},
           'snapshot_sha256':{str(f.relative_to(runtime)):sha256(f) for f in runtime.rglob('*') if f.is_file()}}
    core.put(root / 'manifest.json', man)
    for model in models:
        (root / model).mkdir()
    return root, man


def resume_post(root, man, model):
    core.need(model in man['models'], 'Model is not part of this run')
    core.need((root/model/'integration/SUCCESS.json').is_file() or man.get('recovery',{}).get(model,{}).get('action') == 'ADOPT_POST', 'Verified integration or legacy raw source required')
    core.need(not active_for_cycle(man['config'], man['cycle'], model, queue(man['config'])), 'Active model job; retry refused')
    return core.submit_job(root, man, model, 'post', core.stamp())


@contextlib.contextmanager
def no_lock():
    yield


def main():
    import argparse
    parser = argparse.ArgumentParser(description='RCMs under nampq, preserving operational input/output paths')
    subs = parser.add_subparsers(dest='command')
    for name in ('plan', 'submit'):
        p = subs.add_parser(name)
        p.add_argument('cycle', help='YYYYMMDD00')
        p.add_argument('--model', choices=['both','wrf','regcm'], default='both')
        p.add_argument('--config', default=str(core.HERE/'config.json'))
    p = subs.add_parser('doctor'); p.add_argument('--config', default=str(core.HERE/'config.json'))
    p = subs.add_parser('status'); p.add_argument('run')
    p = subs.add_parser('retry-post'); p.add_argument('run'); p.add_argument('--model', choices=list(MODEL_NAMES), required=True)
    p = subs.add_parser('publish'); p.add_argument('run'); p.add_argument('--model', choices=list(MODEL_NAMES), required=True); p.add_argument('--post', required=True)
    args = parser.parse_args()
    if args.command is None:
        parser.error('a command is required')
    if args.command == 'status':
        return core.status(args.run)
    if args.command in ('retry-post', 'publish'):
        root, man = managed_run(args.run)
        write_preflight(man['config'])
        with cycle_lock(man['config'], man['cycle']):
            if args.command == 'retry-post':
                return resume_post(root, man, args.model)
            return print(json.dumps(publish(root, args.model, Path(args.post)), indent=2))
    cfg = json.loads(Path(args.config).read_text())
    if args.command == 'doctor':
        core.validate_config(cfg, list(MODEL_NAMES))
        return doctor(cfg)
    models = list(MODEL_NAMES) if args.model == 'both' else [args.model]
    if args.command == 'submit':
        write_preflight(cfg)  # Reject the wrong account before creating any file.
    core.validate_config(cfg, models)
    if args.command == 'submit':
        doctor(cfg)  # Verify legacy done destinations before submission.
    # A read-only plan never creates a lock or output file.
    lock = cycle_lock(cfg, args.cycle) if args.command == 'submit' else no_lock()
    with lock:
        decisions = plan(cfg, args.cycle, models, queue(cfg))
        print(json.dumps({'cycle':args.cycle, 'account':username(), 'required_account':cfg['run_user'],
                          'decisions':decisions}, indent=2), flush=True)
        blocked = any(d['action'] in ('WAIT', 'BLOCKED') for d in decisions.values())
        if args.command == 'plan':
            if any(d['action'] in ('NEW','RESUME') for d in decisions.values()):
                end, entries = core.inventory(cfg, args.cycle)
                print('CFS interval:', args.cycle, 'to', end, '; pairs:', len(entries)//2)
            print('PLAN ONLY: no files created or jobs submitted.')
            return
        if blocked:
            print('Models marked WAIT/BLOCKED are retained; independent ready models may proceed.', flush=True)
        fresh = [m for m,d in decisions.items() if d['action'] in ('NEW','RESUME','ADOPT_POST')]
        if fresh:
            root, man = new_run(cfg, args.cycle, fresh, decisions)
            print('Run directory:', root, flush=True)
            for model in fresh:
                job = None
                if decisions[model]['action'] != 'ADOPT_POST':
                    job = core.submit_job(root, man, model, 'integrate', 'initial')
                core.submit_job(root, man, model, 'post', 'initial', job)
        for model, decision in decisions.items():
            if decision['action'] == 'POST':
                root, man = managed_run(decision['run'])
                resume_post(root, man, model)
            elif decision['action'] == 'PUBLISH':
                publish(decision['run'], model, decision['post'])
        core.need(not any(d['action'] == 'BLOCKED' for d in decisions.values()),
                  'Some models remain BLOCKED; inspect decisions and recorded job IDs before retrying.')