"""Single-writer, no-follow toy workspace. Process separation is not an OS sandbox."""
import hashlib
import json
import os
from pathlib import Path
import stat
import time
import uuid
from tiny_agent_boundary import Denied,TASK,INITIAL,PUBLIC,HELDOUT,source_expressions,expression,wrapper,evaluate,parse_proposal


def digest(data):return hashlib.sha256(data).hexdigest()


def fresh_directory(path):
    path=Path(path).absolute()
    for ancestor in (path,*path.parents):
        if ancestor.is_symlink():raise Denied('unsafe_run_path')
        if (ancestor/'.git').exists() or (ancestor/'.git').is_symlink():raise Denied('run_inside_repository')
    if path.exists():raise Denied('run_exists')
    if not path.parent.is_dir():raise Denied('run_parent_missing')
    path.mkdir(mode=0o700)
    return path


class Workspace:
    def __init__(self,path,*,deadline=None,clock=time.monotonic,hook=lambda stage:None):
        if not hasattr(os,'O_NOFOLLOW') or os.open not in os.supports_dir_fd or os.stat not in os.supports_dir_fd:
            raise Denied('containment_unsupported')
        self.path=fresh_directory(path)
        self.clock=clock;self.deadline=deadline if deadline is not None else clock()+600
        self.cancelled=False;self.hook=hook;self.version=0;self.source_read=False
        self.test_binding=None;self.finish_binding=None;self.conflict=False;self.event_number=0
        self.proposals=0;self.model_attempts=0;self.invalid=0;self.observations=[];self.termination=None
        self.root_fd=os.open(self.path,os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW)
        os.mkdir('workspace',0o700,dir_fd=self.root_fd);os.mkdir('backups',0o700,dir_fd=self.root_fd)
        self.fd=os.open('workspace',os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW,dir_fd=self.root_fd)
        self.backup_fd=os.open('backups',os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW,dir_fd=self.root_fd)
        self.audit_fd=os.open('audit.jsonl',os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW,0o600,dir_fd=self.root_fd)
        self.expected={'TASK.txt':digest(TASK.encode()),'toy_math.py':digest(INITIAL.encode())}
        for name,data in (('TASK.txt',TASK.encode()),('toy_math.py',INITIAL.encode())):self.exclusive(self.fd,name,data)
        self.event('created',version=0,hashes=self.expected)

    def close(self):
        for name in ('audit_fd','backup_fd','fd','root_fd'):
            fd=getattr(self,name,None)
            if fd is not None:os.close(fd);setattr(self,name,None)

    def exclusive(self,dir_fd,name,data):
        fd=os.open(name,os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW,0o600,dir_fd=dir_fd)
        try:
            view=memoryview(data)
            while view:
                count=os.write(fd,view);view=view[count:]
            os.fsync(fd)
        finally:os.close(fd)
        os.fsync(dir_fd)

    def event(self,kind,**data):
        self.event_number+=1
        record={'event_id':self.event_number,'kind':kind,'monotonic':self.clock(),**data}
        encoded=(json.dumps(record,sort_keys=True,allow_nan=False)+'\n').encode()
        view=memoryview(encoded)
        while view:
            count=os.write(self.audit_fd,view);view=view[count:]
        os.fsync(self.audit_fd)
        return self.event_number

    def check(self):
        if self.cancelled:raise Denied('cancelled')
        if self.clock()>=self.deadline:raise Denied('time_limit')

    def file_bytes(self,name):
        try:
            before=os.stat(name,dir_fd=self.fd,follow_symlinks=False)
            if not stat.S_ISREG(before.st_mode) or before.st_nlink!=1:raise Denied('workspace_conflict')
            fd=os.open(name,os.O_RDONLY|os.O_NOFOLLOW,dir_fd=self.fd)
            try:
                opened=os.fstat(fd)
                if (before.st_dev,before.st_ino)!=(opened.st_dev,opened.st_ino) or opened.st_nlink!=1 or not stat.S_ISREG(opened.st_mode):
                    raise Denied('workspace_conflict')
                data=os.read(fd,4097)
                if len(data)>4096 or os.read(fd,1):raise Denied('workspace_conflict')
                after=os.fstat(fd)
                if (opened.st_size,opened.st_mtime_ns)!=(after.st_size,after.st_mtime_ns):raise Denied('workspace_conflict')
            finally:os.close(fd)
            visible=os.stat(name,dir_fd=self.fd,follow_symlinks=False)
            if (visible.st_dev,visible.st_ino,visible.st_nlink)!=(opened.st_dev,opened.st_ino,1):raise Denied('workspace_conflict')
            return data
        except (OSError,UnicodeError):raise Denied('workspace_conflict') from None

    def integrity(self):
        if self.conflict:raise Denied('workspace_conflict')
        snapshots={}
        for name,expected in self.expected.items():
            data=self.file_bytes(name)
            if digest(data)!=expected:self.conflict=True;raise Denied('workspace_conflict')
            snapshots[name]=data
        self._verified_snapshots=snapshots
        return snapshots['toy_math.py'].decode('utf-8')

    def bounded(self,observation):
        if len(json.dumps(observation,separators=(',',':'),allow_nan=False).encode())>2048:raise Denied('observation_limit')
        return observation

    def dispatch(self,text,*,source='scripted'):
        self.proposals+=1;action=None
        self.event('proposal',proposal_number=self.proposals,source=source,raw=text)
        try:
            self.check()
            if self.proposals>6:raise Denied('proposal_limit')
            p=parse_proposal(text);action=p['action'];current=self.integrity();self.check()
            if action=='read':
                name=p['file'];content=self._verified_snapshots[name].decode('utf-8')
                observation={'ok':True,'action':action,'file':name,'version':self.version,'sha256':self.expected[name],'content':content}
                if name=='toy_math.py':observation['expressions']=source_expressions(current);self.source_read=True
            elif action=='replace':observation=self.replace(p,current)
            elif action=='test':
                if not self.source_read:raise Denied('source_read_required')
                results=evaluate(current,PUBLIC,self.check)
                self.integrity()
                observation={'ok':True,'action':action,'version':self.version,'sha256':self.expected['toy_math.py'],**results}
                self.test_binding=(self.version,self.expected['toy_math.py']) if results['passed']==results['total'] else None
            elif action=='finish':
                if p['version']!=self.version:raise Denied('stale_version')
                binding=(self.version,self.expected['toy_math.py'])
                if self.test_binding!=binding:raise Denied('tests_not_passing')
                observation={'ok':True,'action':action,'status':'finish_accepted','version':self.version,'sha256':binding[1]}
                self.finish_binding=binding;self.termination='finish_accepted'
            observation=self.bounded(observation)
            self.invalid=0
        except Denied as error:
            self.invalid+=1
            action=getattr(error,'action',action)
            observation={'ok':False,'action':action,'error':error.code,'message':error.code.replace('_',' '),'version':self.version}
            if error.code in ('workspace_conflict','cancelled','time_limit','observation_limit','proposal_limit'):
                self.termination=error.code
                if error.code=='workspace_conflict':self.conflict=True
            elif self.invalid>=2:self.termination='invalid_proposal_limit'
        # Error responses have the same hard envelope bound as successful observations.
        try:self.bounded(observation)
        except Denied:
            self.termination='observation_limit'
            observation={'ok':False,'action':action if type(action) is str and len(action)<=32 else None,
                         'error':'observation_limit','message':'complete observation exceeds byte limit',
                         'version':self.version}
        if not observation['ok'] and type(observation.get('action')) is str and len(observation['action'])>32:
            observation['action']=None
        self.bounded(observation)
        event=self.event('observation',proposal_number=self.proposals,observation=observation)
        self.observations.append({'event_id':event,'observation':observation})
        return observation

    def replace(self,p,current):
        if not self.source_read:raise Denied('source_read_required')
        if p['version']!=self.version:raise Denied('stale_version')
        expressions=source_expressions(current);name=p['function']
        if expressions[name]!=p['old']:raise Denied('old_mismatch')
        expression(p['new'],name)
        updated={**expressions,name:p['new']};after=wrapper(updated);source_expressions(after)
        old_version=self.version
        if after==current:
            return {'ok':True,'action':'replace','status':'no_change','function':name,'old_version':old_version,
                    'version':self.version,'sha256':self.expected['toy_math.py'],'expressions':expressions}
        receipt={'ok':True,'action':'replace','status':'edit_committed','function':name,'old_version':old_version,
                 'version':self.version+1,'sha256':digest(after.encode()),'expressions':updated}
        self.bounded(receipt)  # Never commit an effect whose required receipt cannot be returned.
        backup=f'v{self.version}-{self.expected["toy_math.py"]}.py'
        self.exclusive(self.backup_fd,backup,current.encode())
        temporary='edit-'+uuid.uuid4().hex+'.tmp';self.exclusive(self.fd,temporary,after.encode())
        before_hash=self.expected['toy_math.py'];after_hash=receipt['sha256']
        self.event('edit_intention',before_sha256=before_hash,after_sha256=after_hash,old_version=old_version,
                   new_version=old_version+1,backup=backup,temporary=temporary,expressions=updated)
        self.hook('before_commit');self.check();self.integrity();self.check()
        try:
            os.replace(temporary,'toy_math.py',src_dir_fd=self.fd,dst_dir_fd=self.fd);os.fsync(self.fd)
            self.hook('after_commit')
        except BaseException as error:
            actual=digest(self.file_bytes('toy_math.py'))
            if actual==after_hash:
                self.expected['toy_math.py']=after_hash;self.version=old_version+1;self.test_binding=None
                self.event('edit_reconciled',outcome='committed',receipt=receipt,diagnostic_type=type(error).__name__)
                # Do not blindly retry; report the committed receipt after integrity reconciliation.
                self.integrity();return receipt
            if actual==before_hash:
                self.event('edit_reconciled',outcome='not_committed',diagnostic_type=type(error).__name__)
                raise Denied('edit_not_committed') from None
            self.conflict=True;raise Denied('workspace_conflict') from None
        if digest(self.file_bytes('toy_math.py'))!=after_hash:raise Denied('workspace_conflict')
        self.expected['toy_math.py']=after_hash;self.version=old_version+1;self.test_binding=None
        self.integrity();self.event('edit_committed',receipt=receipt)
        return receipt

    def final(self,termination=None,*,model_invoked=False):
        reason=termination or self.termination or 'proposal_limit'
        report={'termination':reason,'version':self.version,'sha256':self.expected['toy_math.py'],
                'proposals':self.proposals,'model_generation_attempts':self.model_attempts,'verified_success':False,'model_invoked':model_invoked}
        try:
            snapshot=self.integrity(); report['snapshot']=snapshot
            if self.conflict:raise Denied('workspace_conflict')
            self.check()
            if reason=='harness_error':raise Denied('harness_error')
            result=evaluate(snapshot,PUBLIC+HELDOUT,self.check)
            self.hook('after_final_evaluation');self.integrity();self.check()
            report['evaluation']={'status':'completed',**result}
            report['verified_success']=(reason=='finish_accepted' and result['passed']==12
                                       and self.finish_binding==(self.version,self.expected['toy_math.py']))
        except Denied as error:report['evaluation']={'status':error.code}
        self.event('final',report=report)
        self.exclusive(self.root_fd,'final.json',(json.dumps(report,indent=2)+'\n').encode())
        return report