"""Submodule-aware clean setup and lossless checkpoint transport for retries."""
import base64
import hashlib
import json
from pathlib import Path
import shlex
import time

from swe_skyrl_recipe import edit_credit as ec
from swe_skyrl_recipe.agents.mini_swe_tito import MiniSweTitoAgent
from .exact_patch_agent import BytePatchApplyAgent

MAX_PATCH_BYTES = 2_000_000


class ProjectCleanMixin:
    async def setup(self, environment):
        await super().setup(environment)
        await environment.upload_file(Path(ec.__file__).parent / 'agents/sanitize_git.py', '/tmp/harbor-sanitize-git.py')
        await environment.upload_file(Path(__file__).with_name('sanitize_submodules.py'), '/tmp/jev-sanitize-submodules.py')
        result = await environment.exec(command='python /tmp/jev-sanitize-submodules.py --repo /testbed --disposable-sandbox', timeout_sec=240)
        if result.return_code != 0: raise RuntimeError('Project Git sanitization failed: ' + (result.stderr or result.stdout or '')[-2000:])
        self._clean_git_manifest = json.loads(result.stdout.strip().splitlines()[-1])
        (self.logs_dir / 'clean_git_manifest.json').write_text(json.dumps(self._clean_git_manifest, indent=2))

    async def run(self, instruction, environment, context):
        mode = getattr(environment.network_policy.network_mode, 'value', environment.network_policy.network_mode)
        if mode != 'no-network': raise RuntimeError('Clean collector requires no-network agent phase')
        check = await environment.exec(command='ls -1 /sys/class/net', timeout_sec=30)
        if check.return_code != 0 or check.stdout.strip().splitlines() != ['lo']:
            raise RuntimeError('Agent sandbox exposes a non-loopback interface')
        return await super().run(instruction, environment, context)


def diff_bytes_command(initial, target):
    if any(len(t) != 40 or any(c not in '0123456789abcdef' for c in t) for t in (initial, target)):
        raise ValueError('Snapshot trees must be SHA1 object IDs')
    program = (
        'import subprocess,base64,sys; '
        f'p=subprocess.Popen(["git","diff","--binary","--no-renames","{initial}","{target}"],cwd="/testbed",stdout=subprocess.PIPE,stderr=subprocess.PIPE); '
        f'raw=p.stdout.read({MAX_PATCH_BYTES + 1}); '
        f'p.kill() if len(raw)>{MAX_PATCH_BYTES} else None; '
        '_,err=p.communicate(); sys.stderr.buffer.write(err); '
        'print(base64.b64encode(raw).decode("ascii")); '
        f'sys.exit(2 if len(raw)>{MAX_PATCH_BYTES} else p.returncode)'
    )
    return 'python -c ' + shlex.quote(program)


class CheckpointBytesAgent(ProjectCleanMixin, MiniSweTitoAgent):
    def __init__(self, *args, **kwargs):
        kwargs['edit_snapshots'] = True
        super().__init__(*args, **kwargs)
        self._ckpt_dir = self.logs_dir / 'checkpoints'; self._ckpt_dir.mkdir(parents=True, exist_ok=True)
        self._started = time.monotonic(); self._last_byte_snapshot = None; self._byte_capture_ok = True

    def version(self): return 'jev-checkpoint-bytes-submodules-v3'

    async def _snapshot_step(self, environment, step, snaps):
        ok = await super()._snapshot_step(environment, step, snaps)
        if not ok or len(snaps) < 2 or snaps[-1].tree == snaps[-2].tree: return ok
        record = {'step': step, 'tree': snaps[-1].tree, 'changed': snaps[-1].changed,
                  'elapsed_seconds': time.monotonic() - self._started, 'transport': 'base64-bytes-v3'}
        try:
            rc, encoded = await self._exec_raw(environment, diff_bytes_command(snaps[0].tree, snaps[-1].tree), timeout_sec=120)
            if rc: raise RuntimeError(f'Byte capture exit {rc}')
            raw = base64.b64decode(encoded.strip(), validate=True)
            if len(raw) > MAX_PATCH_BYTES: raise ValueError('Checkpoint exceeds byte limit')
            name = f'{step:04d}.patch'; (self._ckpt_dir / name).write_bytes(raw)
            record.update({'patch_ok': True, 'patch_file': name, 'patch_bytes': len(raw),
                           'patch_sha256': hashlib.sha256(raw).hexdigest()})
            self._last_byte_snapshot = record
        except Exception as e:
            record.update({'patch_ok': False, 'error': repr(e)}); self._byte_capture_ok = False
        with (self._ckpt_dir / 'index.jsonl').open('a') as f: f.write(json.dumps(record) + '\n')
        return ok

    async def _build_edit_credit(self, environment, snaps, snap_ok, snap_err):
        info = await super()._build_edit_credit(environment, snaps, snap_ok, snap_err)
        final = snaps[-1].tree if snaps else None
        last = self._last_byte_snapshot
        empty = bool(snaps and last is None and final == snaps[0].tree)
        complete = self._byte_capture_ok and snap_ok and (empty or bool(last and last['tree'] == final))
        info['byte_capture'] = {'schema': 'base64-bytes-v3', 'complete': complete,
             'final_tree': final, 'final_patch_sha256': hashlib.sha256(b'').hexdigest() if empty else (last or {}).get('patch_sha256')}
        return info


class SubmoduleExactPatchAgent(ProjectCleanMixin, BytePatchApplyAgent):
    def version(self): return 'jev-exact-byte-tree-submodules-v3'