"""Single in-place student -> teacher takeover at a DSH request boundary.

The budget proxy calls ``TakeoverController.plan`` before forwarding each
model request. At that point every tool call of the previous decision has
returned, so the request messages are exactly what the next model will see.
Triggers use only that visible history plus the container workspace state:

* L: a call repeats an earlier call with the same complete observation
  (sha256 + length) and no relevant file-state change in between;
* S: 20 decisions completed and the workspace never differed from base;
* W: a read/edit/write call returned the tool ``Error:`` envelope, or a visible
  test run failed after a successful modification.

The student prefix is counted with the student's real chat template and
tokenizer; above the limit the trajectory stops and nothing goes to a teacher.
History is never truncated or rewritten here.
"""
from __future__ import annotations

from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
import subprocess
import sys
import threading

EVAL = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(EVAL), str(EVAL.parent / 'swefactory_self/results')]
from evalkit.request_view import request_view  # noqa: E402

STUDENT_PREFIX_LIMIT = 15000
STALL_DECISIONS = 20
TRAJECTORY_KINDS = ('P1', 'P2', 'P3')
FILE_TOOLS = {'read', 'edit', 'write'}
TEST_COMMAND = re.compile(
    r'\bpy\.?test\b|-m\s+(?:pytest|unittest)\b|\btox\b|\bnox\b|runtests\.py\b|\bgo\s+test\b'
    r'|\bcargo\s+test\b|\b(?:npm|yarn|pnpm)\s+(?:run\s+)?test\b|\bmake\s+(?:test|check)\b|\bbin/test\b')
TEST_FAILED = re.compile(
    r'^=+ .*\b[1-9]\d* (?:failed|errors?)\b.*=+\s*$'          # pytest summary
    r'|^[1-9]\d* (?:failed|errors?)\b.* in [\d.]+s'           # pytest -q summary
    r'|^FAILED \((?:failures|errors)=[1-9]'                   # unittest / django
    r'|^ERROR collecting\b|^Interrupted: [1-9]\d* errors? during collection'
    r'|^--- FAIL: |^FAIL\s+\S+\s+[\d.]+s$'                    # go test
    r'|^test result: FAILED\.', re.M)
EXIT_CODE = re.compile(r'\[exit code: (-?\d+)\]\s*$')
CHAINED = re.compile(r'\||;|&&|\n')

WORKSPACE_STATE = r'''import hashlib,json,os,subprocess,sys
root=sys.argv[1]
def git(*a):
    return subprocess.run(['git','--no-optional-locks','-C',root,*a],capture_output=True,check=True).stdout
h=hashlib.sha256()
diff=git('diff','HEAD','--binary','--no-ext-diff','--no-color','--no-renames')
h.update(b'tracked\0'+hashlib.sha256(diff).digest())
untracked=sorted(p for p in git('ls-files','--others','--exclude-standard','-z').split(b'\0') if p)
for name in untracked:
    path=os.path.join(root.encode(),name)
    if os.path.islink(path):
        data=b'link\0'+os.readlink(path)
    else:
        try:
            with open(path,'rb') as stream: data=stream.read()
        except OSError as error: data=b'unreadable\0'+str(error.errno).encode()
    h.update(b'untracked\0'+name+b'\0'+hashlib.sha256(data).digest())
print(json.dumps(dict(sha256=h.hexdigest(),tracked_diff_bytes=len(diff),untracked_files=len(untracked))))
'''


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def digest(value) -> str:
    """Same canonical JSON digest as the token-alignment receipts."""
    return hashlib.sha256(json.dumps(value, ensure_ascii=False, sort_keys=True,
                                     separators=(',', ':')).encode()).hexdigest()


def workspace_state(container: str, workdir: str = '/testbed', docker: str = 'docker') -> dict:
    """Hash tracked changes and untracked, non-ignored files inside the task container."""
    result = subprocess.run([docker, 'exec', '-i', '-w', workdir, container, 'python', '-', workdir],
                            input=WORKSPACE_STATE, capture_output=True, encoding='utf-8',
                            errors='replace', timeout=120)
    if result.returncode:
        raise RuntimeError('workspace state probe failed: ' + result.stderr[-800:])
    return json.loads(result.stdout)


def _arguments(call: dict) -> dict:
    raw = (call.get('function') or {}).get('arguments')
    if isinstance(raw, dict):
        return raw
    try:
        value = json.loads(raw or '{}')
    except ValueError:
        return {'__unparsed__': raw}
    return value if isinstance(value, dict) else {'__value__': value}


def call_key(name: str, arguments: dict) -> str:
    # bash `description` is UI text, not part of what the command does.
    if name == 'bash':
        arguments = {key: arguments.get(key) for key in ('command', 'workdir')}
    return digest([name, arguments])


def decisions(messages: list[dict]) -> list[dict]:
    """Each assistant message with its calls and the tool results that answer them."""
    result = []
    for message in messages:
        if message.get('role') == 'assistant':
            calls = []
            for call in message.get('tool_calls') or []:
                name = (call.get('function') or {}).get('name') or ''
                arguments = _arguments(call)
                calls.append({'id': call.get('id'), 'name': name, 'arguments': arguments,
                              'key': call_key(name, arguments)})
            result.append({'calls': calls, 'results': {}})
        elif message.get('role') == 'tool' and result:
            content = message.get('content')
            if not isinstance(content, str):
                content = json.dumps(content, ensure_ascii=False)
            result[-1]['results'][message.get('tool_call_id')] = content
    return result


def observation_key(content: str) -> tuple[str, int]:
    return hashlib.sha256(content.encode('utf-8')).hexdigest(), len(content)


def tool_error(name: str, content: str | None) -> bool:
    return name in FILE_TOOLS and content is not None and content.lstrip().startswith('Error:')


def test_failed(command: str, output: str) -> bool:
    """A visible runner failure. Exit codes count only for a single unchained command."""
    if not TEST_COMMAND.search(command or ''):
        return False
    if TEST_FAILED.search(output or ''):
        return True
    code = EXIT_CODE.search(output or '')
    return bool(code and int(code.group(1)) != 0 and not CHAINED.search(command.strip()))


def detect(history: list[dict], states: dict[int, str], boundary: int) -> list[dict]:
    """Triggers visible at request ``boundary`` (decisions 1..boundary-1 completed).

    ``states[b]`` is the workspace sha256 observed before request ``b``; ``states[1]`` is base.
    """
    if len(history) != boundary - 1 or boundary < 2 or 1 not in states or boundary not in states:
        return []
    base, current = states[1], states[boundary]
    last_index = boundary - 1
    last = history[-1]
    fired = []
    for call in last['calls']:
        content = last['results'].get(call['id'])
        if content is None:
            continue
        seen = observation_key(content)
        for index, earlier in enumerate(history[:-1], start=1):
            match = next((c for c in earlier['calls'] if c['key'] == call['key']
                          and earlier['results'].get(c['id']) is not None
                          and observation_key(earlier['results'][c['id']]) == seen), None)
            if match is not None and all(states.get(b) == current for b in range(index + 1, boundary + 1)):
                fired.append({'kind': 'L', 'natural': True, 'evidence': {
                    'first': {'request_index': index, 'call_id': match['id']},
                    'second': {'request_index': last_index, 'call_id': call['id']},
                    'tool': call['name'], 'call_sha256': call['key'],
                    'observation_sha256': seen[0], 'observation_chars': seen[1],
                    'workspace_state_sha256': current}})
                break
        if fired:
            break
    modified_before = states.get(last_index, base) != base
    edited_here = False
    for call in last['calls']:
        content = last['results'].get(call['id'])
        if tool_error(call['name'], content):
            fired.append({'kind': 'W', 'natural': True, 'evidence': {
                'reason': 'tool_contract_error', 'request_index': last_index, 'call_id': call['id'],
                'tool': call['name'], 'observation_prefix': content.strip()[:240]}})
            break
        if call['name'] in {'edit', 'write'} and content is not None:
            edited_here = True
        if call['name'] == 'bash' and content is not None and test_failed(
                str(call['arguments'].get('command') or ''), content):
            if modified_before or (edited_here and current != base):
                fired.append({'kind': 'W', 'natural': True, 'evidence': {
                    'reason': 'visible_test_failure_after_modification', 'request_index': last_index,
                    'call_id': call['id'], 'command': str(call['arguments'].get('command'))[:400],
                    'observation_tail': content[-400:], 'workspace_state_sha256': current}})
                break
    if last_index == STALL_DECISIONS and all(states.get(b) == base for b in range(1, boundary + 1)):
        fired.append({'kind': 'S', 'natural': True, 'evidence': {
            'decisions_without_file_change': STALL_DECISIONS, 'workspace_state_sha256': base}})
    order = {'L': 0, 'W': 1, 'S': 2}
    return sorted(fired, key=lambda item: order[item['kind']])


class PrefixCounter:
    """Student prefix tokens: real Qwen chat template + tokenizer on CPU."""

    def __init__(self, model_dir: Path):
        from to_sft_qwen import Renderer
        self.renderer = Renderer(model_dir)
        self.assets = {name: hashlib.sha256((Path(model_dir) / name).read_bytes()).hexdigest()
                       for name in ('tokenizer.json', 'tokenizer_config.json', 'chat_template.jinja')}

    def __call__(self, body: dict) -> list[int]:
        view = request_view(body)
        prompt = self.renderer.render_request(view['messages'], view['tools'])
        return self.renderer.tokenizer.encode(prompt, add_special_tokens=False).ids


class TakeoverController:
    """Decides the phase of each logical request and keeps verifiable receipts."""

    def __init__(self, kind: str, run_dir: Path, *, prefix_counter=None, state_fn=None,
                 identity_fn=None, prefix_limit: int = STUDENT_PREFIX_LIMIT,
                 force_at: int | None = None):
        if kind not in TRAJECTORY_KINDS:
            raise ValueError('trajectory_kind must be one of ' + ', '.join(TRAJECTORY_KINDS))
        if kind != 'P1' and (prefix_counter is None or state_fn is None):
            raise ValueError(kind + ' needs the student prefix counter and workspace state probe')
        if force_at is not None and (kind != 'P2' or force_at < 2):
            raise ValueError('forced takeover is only defined for P2 at request >= 2')
        self.kind, self.run_dir = kind, run_dir
        self.count, self.state_fn, self.identity_fn = prefix_counter, state_fn, identity_fn or (lambda: {})
        self.limit, self.force_at = prefix_limit, force_at
        self.phase = 'teacher' if kind == 'P1' else 'student'
        self.states: dict[int, str] = {}
        self.handoff: dict | None = None
        self.stop: dict | None = None
        self.terminal: dict | None = None
        self.lock = threading.Lock()
        self.write_state()

    def write_state(self) -> None:
        state = {'trajectory_kind': self.kind, 'phase': self.phase, 'student_prefix_limit': self.limit,
                 'forced_takeover_at': self.force_at,
                 'handoff_request_index': self.handoff and self.handoff['boundary_request_index'],
                 'stop': self.stop, 'terminal': self.terminal, 'updated_at': utc_now()}
        (self.run_dir / 'takeover-state.json').write_text(
            json.dumps(state, ensure_ascii=False, indent=2) + '\n', encoding='utf-8', newline='\n')

    def record_terminal(self, reason: str, **detail) -> None:
        with self.lock:
            if self.terminal is None:
                self.terminal = {'reason': reason, 'phase': self.phase, **detail, 'at': utc_now()}
                self.write_state()

    def plan(self, index: int, body: dict) -> dict:
        """Return {'phase': 'student'|'teacher'} or {'stop': {...}} for request ``index``."""
        with self.lock:
            if self.stop is not None:
                return {'stop': self.stop}
            if self.phase == 'teacher':
                return {'phase': 'teacher'}
            state = self.state_fn()
            self.states[index] = state['sha256']
            ids = self.count(body)
            view = request_view(body)
            triggers = detect(decisions(view['messages']), self.states, index) if self.kind != 'P1' else []
            forced = self.kind == 'P2' and self.force_at == index
            row = {'request_index': index, 'prefix_tokens': len(ids), 'prefix_limit': self.limit,
                   'workspace_state': state, 'triggers': triggers, 'forced_smoke': forced, 'at': utc_now()}
            if len(ids) > self.limit:
                # Diagnostic only: never truncate history and never hand an oversized prefix to a teacher.
                self.stop = {'reason': 'student_prefix_limit', 'request_index': index,
                             'prefix_tokens': len(ids), 'prefix_limit': self.limit,
                             'pending_triggers': [t['kind'] for t in triggers]}
                row['action'] = 'stop_prefix_limit'
            elif self.kind == 'P2' and (triggers or forced):
                trigger = {
                    'kind': 'forced_smoke', 'natural': False,
                    'evidence': {'note': 'acceptance-smoke switch; not a natural symptom'}} if forced else triggers[0]
                self.handoff = {
                    'schema_version': 1, 'trajectory_kind': self.kind, 'trigger': trigger,
                    'other_triggers': [t['kind'] for t in (triggers if forced else triggers[1:])],
                    'boundary_request_index': index, 'student_decisions': index - 1,
                    'prefix_tokens': len(ids), 'prefix_ids_sha256': digest(ids),
                    'prefix_view_sha256': digest(view), 'workspace_state_sha256': state['sha256'],
                    'workspace_state': state, **self.identity_fn(), 'at': utc_now()}
                (self.run_dir / 'handoff.json').write_text(
                    json.dumps(self.handoff, ensure_ascii=False, indent=2) + '\n', encoding='utf-8', newline='\n')
                self.phase = 'teacher'
                row['action'] = 'handoff'
            else:
                row['action'] = 'student'
            with (self.run_dir / 'takeover-evaluations.jsonl').open('a', encoding='utf-8', newline='\n') as stream:
                stream.write(json.dumps(row, ensure_ascii=False) + '\n')
            self.write_state()
            return {'stop': self.stop} if self.stop else {'phase': self.phase}