"""Internal XML bridge with explicit, verifiable workflow requirements."""
import json
import re
import time
import uuid
import hashlib
import torch


def schema(name,description,properties,required):
    return {"type":"function","function":{"name":name,"description":description,
        "parameters":{"type":"object","properties":properties,"required":required,
                      "additionalProperties":False}}}

STRING={"type":"string"}
TOOLS=[
    schema("calc","Evaluate exact arithmetic on the internal computer. + - * / // % ** gcd and parentheses.",
           {"expression":STRING},["expression"]),
    schema("put","Save a text file; creates a new revision.",{"name":STRING,"content":STRING},["name","content"]),
    schema("get","Read a text file.",{"name":STRING},["name"]),
    schema("list_files","List saved files.",{},[]),
    schema("delete","Delete a saved file.",{"name":STRING},["name"]),
    schema("rollback","Restore a retained file revision, creating a new revision.",
           {"name":STRING,"revision":STRING},["name","revision"]),
]


def parse_call(raw):
    block=re.search(r"<tool_call>(.*?)</tool_call>",raw,re.S)
    if not block:
        if "<tool_call>" in raw:raise ValueError("Truncated tool call")
        return None
    body=block.group(1).strip()
    if body.startswith('<function'):
        function=re.fullmatch(r"<function=([A-Za-z_]\w*)>(.*?)</function>\s*",body,re.S)
        if function is None:raise ValueError('Malformed function envelope')
        inner=function.group(2)
        pattern=r"<parameter=([A-Za-z_]\w*)>(.*?)</parameter>"
        pairs=re.findall(pattern,inner,re.S)
        if re.sub(pattern,'',inner,flags=re.S).strip():raise ValueError('Malformed parameter envelope')
        if len(dict(pairs))!=len(pairs):raise ValueError('Duplicate parameter')
        return {'name':function.group(1),'arguments':{k:(v if k in ('content','text','query','forbidden','name') else v.strip()) for k,v in pairs}}
    def unique(pairs):
        value={}
        for key,item in pairs:
            if key in value:raise ValueError('Duplicate JSON key')
            value[key]=item
        return value
    def bad_constant(x):raise ValueError('Non-finite JSON number')
    value=json.loads(body,object_pairs_hook=unique,parse_constant=bad_constant)
    if not isinstance(value,dict) or set(value)-{'name','arguments'} or 'name' not in value:
        raise ValueError('Invalid tool-call object')
    if not isinstance(value['name'],str) or not re.fullmatch(r'[A-Za-z_]\w*',value['name']):
        raise ValueError('Invalid tool name')
    args=value.get('arguments',{})
    if isinstance(args,str):args=json.loads(args,object_pairs_hook=unique,parse_constant=bad_constant)
    if not isinstance(args,dict):raise ValueError('Arguments must be an object')
    return {'name':value['name'],'arguments':args}


def execute(model,session,call):
    name,args=call["name"],call["arguments"]
    choices={t["function"]["name"]:t["function"]["parameters"] for t in TOOLS}
    if name not in choices: raise ValueError(f"Unknown tool: {name}")
    spec=choices[name]
    if not isinstance(args,dict) or set(args)-set(spec["properties"]) or set(spec["required"])-set(args):
        raise ValueError("Invalid parameter names")
    if any(not isinstance(v,str) for v in args.values()): raise TypeError("String parameters required")
    if name=="calc": return model.calculate(args["expression"])
    if name=="put": return model(mode="memory",memory_state=session.memory,operation="put_many",files={args["name"]:args["content"]})
    if name=="get": return model(mode="memory",memory_state=session.memory,operation="get",**args)
    if name=="list_files": return model(mode="memory",memory_state=session.memory,operation="list")
    if name=="delete": return model(mode="memory",memory_state=session.memory,operation="put_many",files={args["name"]:None})
    if name=="rollback": return model(mode="memory",memory_state=session.memory,operation="rollback",**args)
    raise AssertionError(name)


def fulfilled(events,requirements):
    cursor=0
    for event in events:
        if cursor<len(requirements) and event["response"]["ok"] and event["call"]["name"]==requirements[cursor]:
            cursor+=1
    return cursor==len(requirements)


@torch.inference_mode()
def generate_text(model,session,max_new_tokens):
    from transformers import StoppingCriteria,StoppingCriteriaList
    tok=model.tokenizer
    close=tok.encode("</tool_call>",add_special_tokens=False)
    class Stop(StoppingCriteria):
        def __call__(self,input_ids,scores,**kwargs):
            return input_ids.shape[1]>=len(close) and input_ids[0,-len(close):].tolist()==close
    batch=tok.apply_chat_template(session.messages,tools=TOOLS,add_generation_prompt=True,
        enable_thinking=False,tokenize=True,return_dict=True,return_tensors="pt")
    if batch["input_ids"].shape[1]>16384: raise RuntimeError("Session context exceeds 16K notebook limit")
    batch=batch.to(next(model.language_model.parameters()).device)
    out=model.language_model.generate(**batch,max_new_tokens=max_new_tokens,do_sample=False,
        repetition_penalty=1.1,stopping_criteria=StoppingCriteriaList([Stop()]),
        pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id)
    return tok.decode(out[0,batch["input_ids"].shape[1]:],skip_special_tokens=False)


def chat(model,session,prompt,required_tools=(),required_files=None,max_calls=12,max_new_tokens=1536,verbose=True):
    if model.tokenizer is None: raise RuntimeError("Tokenizer is not attached")
    if max_calls<1: raise ValueError("max_calls must be positive")
    session.messages.append({"role":"user","content":prompt})
    events=[];started=time.perf_counter()
    for attempt in range(max_calls+3):
        raw=generate_text(model,session,max_new_tokens)
        call=parse_call(raw)
        if call is None:
            missing=[]
            if not fulfilled(events,required_tools): missing.append("required ordered calls: "+", ".join(required_tools))
            for name,expected in (required_files or {}).items():
                try:
                    actual=model.files(session.memory,"get",name=name)["content"]
                    if expected not in actual: missing.append("required file content: "+name)
                except KeyError: missing.append("required file: "+name)
            if missing:
                if attempt>=max_calls+1 or len(events)>=max_calls:
                    raise RuntimeError("Workflow incomplete: "+"; ".join(missing))
                session.messages.append({"role":"assistant","content":raw})
                session.messages.append({"role":"user","content":
                    "Execution verification failed. Complete the missing operations before answering: "+"; ".join(missing)})
                continue
            answer=re.sub(r"<think>.*?</think>","",raw,flags=re.S)
            if "<think>" in answer: raise RuntimeError("Incomplete reasoning; increase token budget")
            answer=re.sub(r"<\|(?:im_end|im_start|endoftext|startoftext)\|>","",answer).strip()
            session.messages.append({"role":"assistant","content":answer})
            return {"answer":answer,"events":events,"workflow_verified":bool(required_tools or required_files),
                    "seconds":time.perf_counter()-started}
        if len(events)>=max_calls: raise RuntimeError("Tool-call budget exhausted")
        identifier="call_"+uuid.uuid4().hex
        content=re.sub(r"<tool_call>.*?</tool_call>","",raw,flags=re.S).strip()
        session.messages.append({"role":"assistant","content":content,
            "tool_calls":[{"id":identifier,"type":"function","function":call}]})
        try: response={"ok":True,"result":execute(model,session,call)}
        except (ValueError,KeyError,TypeError,MemoryError,TimeoutError,RuntimeError) as error:
            response={"ok":False,"error":f"{type(error).__name__}: {error}"}
        event={"call":call,"response":response}
        events.append(event)
        previous=session.audit[-1]["hash"] if session.audit else "0"*64
        digest=hashlib.sha256((previous+json.dumps(event,sort_keys=True)).encode()).hexdigest()
        session.audit.append({**event,"previous":previous,"hash":digest})
        if verbose: print("→",call["name"],json.dumps(response,ensure_ascii=False)[:500],flush=True)
        session.messages.append({"role":"tool","tool_call_id":identifier,"name":call["name"],
                                 "content":json.dumps(response,ensure_ascii=False)})
    raise RuntimeError("Workflow did not finish")