#!/usr/bin/env python3
"""Streaming pipeline: extract (veer90) -> organize (our architecture) -> rclone to gdrive -> delete local.
Processes any downloaded archive (has sibling .done) one at a time, dedup by SHA across archives.
Survives the session (run under nohup). Exits after ~24h with no new work."""
import os, sys, subprocess, shutil, time, glob

BASE = os.path.expanduser("~/Разработки/Прошивки")
DOWNLOADS = os.path.join(BASE, "_downloads")
EXTRACT = os.path.join(BASE, "_pipe_extract")
ORG = os.path.join(BASE, "_pipe_organized")
LOGDIR = os.path.join(BASE, "_logs")
LOG = os.path.join(LOGDIR, "pipeline.log")
SEENF = os.path.join(BASE, "_seen_hashes.txt")
GDRIVE = "gdrive:Прошивки avto1000"
PASSWORD = "veer90"

os.makedirs(LOGDIR, exist_ok=True)

# Reuse classification functions from the organizer (run only its header, no side effects)
_src = open('/tmp/firmware_organize.py', encoding='utf-8').read().split('\n')
_ns = {}
exec('\n'.join(_src[:219]), _ns)
classify_brand = _ns['classify_brand']; classify_model = _ns['classify_model']
detect_calibrator = _ns['detect_calibrator']; detect_calibration_id = _ns['detect_calibration_id']
classify_category = _ns['classify_category']; file_hash = _ns['file_hash']; sanitize = _ns['sanitize']

def log(m):
    line = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {m}"
    print(line, flush=True)
    with open(LOG, 'a', encoding='utf-8') as f:
        f.write(line + '\n')

seen = set()
if os.path.exists(SEENF):
    seen = set(open(SEENF, encoding='utf-8').read().split())

def organize_extracted():
    """Copy unique firmware into ORG. Dedup against global `seen` (already-uploaded) plus
    within-archive. Does NOT mutate global `seen` — returns the new hashes so the caller
    commits them ONLY after a successful upload (avoids poisoning dedup on failed uploads)."""
    new_hashes = []
    local = set(seen)
    for root, _, files in os.walk(EXTRACT):
        for fn in files:
            if not fn.lower().endswith(('.bin', '.bbf', '.hex', '.s19', '.ori')):
                continue
            full = os.path.join(root, fn)
            try:
                size = os.path.getsize(full)
            except OSError:
                continue
            if size < 16384 or size > 8 * 1024 * 1024:
                continue
            try:
                data = open(full, 'rb').read()
            except OSError:
                continue
            h = file_hash(data)
            if h in local:
                continue
            rel = os.path.relpath(full, EXTRACT)
            brand = classify_brand(rel)
            model = classify_model(rel, brand)
            calib = detect_calibrator(rel)
            cat = classify_category(rel)
            cal_id = detect_calibration_id(data, rel) or "_без-калибровки"
            tdir = os.path.join(ORG, sanitize(brand), sanitize(model), sanitize(cal_id), sanitize(cat))
            os.makedirs(tdir, exist_ok=True)
            prefix = f"[{calib}] " if calib != "Unknown" else ""
            base, ext = os.path.splitext(fn)
            tname = sanitize(f"{prefix}{base}__{h}{ext}")
            try:
                shutil.copy2(full, os.path.join(tdir, tname))
            except OSError:
                continue
            local.add(h)
            new_hashes.append(h)
    return new_hashes

def extract_one(archive, dest):
    """unar first (robust RAR), 7z fallback. Returns True on success."""
    low = archive.lower()
    order = (['unar', '7z'] if low.endswith('.rar') else ['7z', 'unar'])
    for tool in order:
        if tool == 'unar':
            r = subprocess.run(['unar', '-q', '-p', PASSWORD, '-f', '-D', '-o', dest, archive],
                               capture_output=True, text=True)
        else:
            r = subprocess.run(['7z', 'x', '-y', f'-p{PASSWORD}', f'-o{dest}', archive],
                               capture_output=True, text=True)
        if r.returncode == 0:
            return True
    return False

def extract_recursive(dest, depth=6):
    """Extract nested archives in-place, removing each after extraction. Bounded depth."""
    for _ in range(depth):
        nested = []
        for root, _, files in os.walk(dest):
            for fn in files:
                if fn.lower().endswith(('.rar', '.zip', '.7z')):
                    nested.append(os.path.join(root, fn))
        if not nested:
            return
        for arc in nested:
            sub = arc + '__x'
            os.makedirs(sub, exist_ok=True)
            if extract_one(arc, sub):
                try:
                    os.remove(arc)
                except OSError:
                    pass

def process_archive(arc):
    name = os.path.basename(arc)
    log(f"PROCESS {name} ({os.path.getsize(arc)/1e6:.0f} MB)")
    shutil.rmtree(EXTRACT, ignore_errors=True); os.makedirs(EXTRACT)
    if not extract_one(arc, EXTRACT):
        log(f"  EXTRACT FAIL (unar+7z) {name}")
        open(arc + '.failed', 'w').close()
        shutil.rmtree(EXTRACT, ignore_errors=True)
        return False
    extract_recursive(EXTRACT)
    shutil.rmtree(ORG, ignore_errors=True); os.makedirs(ORG)
    new_hashes = organize_extracted()
    n = len(new_hashes)
    log(f"  organized {n} unique firmware files")
    if n == 0:
        # Don't delete: likely odd structure / all duplicates — keep for manual review.
        log(f"  0 files organized — KEEP archive ({name}) for review")
        open(arc + '.review', 'w').close()
        shutil.rmtree(EXTRACT, ignore_errors=True); shutil.rmtree(ORG, ignore_errors=True)
        return False
    # Upload with retries — gdrive API occasionally returns transient EOF/connection errors.
    uploaded = False
    for attempt in range(1, 5):
        ru = subprocess.run(['rclone', 'move', ORG, GDRIVE,
                             '--transfers=16', '--checkers=16', '--no-traverse',
                             '--drive-pacer-min-sleep=10ms',
                             '--retries=5', '--low-level-retries=10', '--stats=0'],
                            capture_output=True, text=True)
        if ru.returncode == 0:
            uploaded = True
            break
        log(f"  UPLOAD FAIL (try {attempt}/4): {(ru.stderr or '')[-200:]}")
        time.sleep(30 * attempt)
    if not uploaded:
        # Keep archive AND don't commit hashes — leave EXTRACT/ORG cleaned, retry whole archive later.
        shutil.rmtree(EXTRACT, ignore_errors=True); shutil.rmtree(ORG, ignore_errors=True)
        return False
    log(f"  uploaded {n} files -> {GDRIVE}")
    # Commit hashes to global seen ONLY after a confirmed upload.
    seen.update(new_hashes)
    open(SEENF, 'w', encoding='utf-8').write('\n'.join(sorted(seen)))
    shutil.rmtree(EXTRACT, ignore_errors=True); shutil.rmtree(ORG, ignore_errors=True)
    try:
        os.remove(arc)
    except OSError:
        pass
    open(arc + '.processed', 'w').close()
    log(f"  DONE {name} — archive deleted, {len(seen)} total unique hashes")
    return True

def find_ready():
    out = []
    for f in glob.glob(os.path.join(DOWNLOADS, '*')):
        low = f.lower()
        if not low.endswith(('.rar', '.zip', '.7z')):
            continue
        if any(os.path.exists(f + ext) for ext in ('.processed', '.failed', '.review')):
            continue
        # completion: sibling .done from the downloader
        stem = os.path.splitext(f)[0]
        if os.path.exists(stem + '.done'):
            out.append(f)
    return sorted(out)

def main():
    log("=== PIPELINE START ===")
    idle = 0
    while True:
        ready = find_ready()
        if not ready:
            idle += 1
            if idle > 2880:  # ~24h of 30s ticks with no work
                log("=== idle 24h, exit ===")
                break
            time.sleep(30)
            continue
        idle = 0
        for a in ready:
            try:
                process_archive(a)
            except Exception as e:
                log(f"  ERROR {os.path.basename(a)}: {e}")
                open(a + '.failed', 'w').close()

if __name__ == "__main__":
    main()
