#!/usr/bin/env python3
"""Run repeated hard-limit tool-use turns with real weather/build tools."""
from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import statistics
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from pathlib import Path

MIB = 1024 * 1024
GIB = 1024 ** 3

SYSTEM = (
    'You are a function-calling agent. Available functions: '
    '{"name":"get_weather","parameters":{"city":"string"}}, '
    '{"name":"compile_project","parameters":{"project":"string"}}. '
    'For this request, call ONLY the named function and use exactly the argument value '
    'specified by the user. Output the function call, then stop. Do not answer in prose. '
    'Use this syntax: <tool_call>NAME<arg_key>ARG</arg_key>'
    '<arg_value>VALUE</arg_value></tool_call>. detailed thinking off'
)


def facts() -> dict:
    cpu = ""
    for line in Path('/proc/cpuinfo').open():
        if line.startswith('model name'):
            cpu = line.split(':', 1)[1].strip()
            break
    cg = {}
    for key in ('memory.max', 'memory.swap.max', 'memory.high'):
        try:
            cg[key] = Path('/sys/fs/cgroup', key).read_text().strip()
        except OSError:
            cg[key] = None
    return {'cpu_model': cpu, 'cpu_threads': os.cpu_count(), 'cgroup': cg}


def sample(pid: int) -> tuple[int, int, int]:
    rss = anon = swap = 0
    try:
        for line in Path(f'/proc/{pid}/status').read_text().splitlines():
            if line.startswith('VmRSS:'):
                rss = int(line.split()[1]) * 1024
            elif line.startswith('VmSwap:'):
                swap = int(line.split()[1]) * 1024
        for line in Path(f'/proc/{pid}/smaps_rollup').read_text().splitlines():
            if line.startswith('Anonymous:'):
                anon = int(line.split()[1]) * 1024
            elif line.startswith('AnonHWM:'):
                anon = max(anon, int(line.split()[1]) * 1024)
    except OSError:
        pass
    return rss, anon, swap


def model_call(args: argparse.Namespace, level_mb: int, prompt: str,
               prefix: str, token_cap: int) -> dict:
    log_path = Path(args.output_dir) / f'{prefix}.log'
    stats_path = log_path.with_suffix('.jsonl')
    log_path.parent.mkdir(parents=True, exist_ok=True)
    stats_path.unlink(missing_ok=True)
    if args.evict_helper:
        subprocess.run(['sync'], check=False)
        subprocess.run([args.evict_helper, '--evict', str(Path(args.model).resolve())], check=False)
    env = dict(os.environ)
    env.update({
        'SDQ_MODEL_PATH': str(Path(args.model).resolve()),
        'SDQ_RAM_BUDGET_MB': str(level_mb),
        'SDQ_RESERVE_MB': str(args.reserve_map[level_mb]),
        'SDQ_PREFETCH': '0', 'SDQ_FADV_DONTNEED': '1',
        'SDQ_STATS_FILE': str(stats_path),
    })
    cmd = [
        'prlimit', f'--data={level_mb * MIB}', '--nofile=65536',
        str(Path(args.bin).resolve()), '-m', str(Path(args.model).resolve()),
        '-p', prompt, '-n', str(token_cap), '-c', str(args.ctx),
        '-t', str(args.threads), '-b', str(args.ubatch), '--temp', '0',
        '--no-chat', '--no-interactive', '--ram-budget-mb', str(level_mb),
    ]
    with log_path.open('wb') as err_file:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=err_file, env=env)
        peak_rss = peak_anon = peak_swap = 0
        start = time.monotonic()
        while proc.poll() is None:
            rss, anon, swap = sample(proc.pid)
            peak_rss = max(peak_rss, rss)
            peak_anon = max(peak_anon, anon)
            peak_swap = max(peak_swap, swap)
            if time.monotonic() - start > args.timeout:
                proc.kill()
                proc.wait()
                raise TimeoutError(f'model call exceeded {args.timeout}s')
            time.sleep(0.03)
        stdout = proc.stdout.read().decode('utf-8', 'replace') if proc.stdout else ''
    err = log_path.read_text(errors='replace')
    text = stdout.split('assistant:', 1)[-1].strip() if 'assistant:' in stdout else stdout.strip()
    text = re.split(r'\[sdq\] decode:', text)[0].strip()
    if '</tool_call>' in text:
        text = text[:text.index('</tool_call>') + len('</tool_call>')]
    else:
        text = re.split(r'<\|role_end\|>|<\|endoftext\|>', text, maxsplit=1)[0].strip()
    rates = re.findall(r'decode:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s\s*\|\s*SSD 讀\s*([\d.]+) MB\s*\(([\d.]+) MB/tok\)', err)
    stats = [json.loads(line) for line in stats_path.read_text().splitlines()] if stats_path.exists() else []
    phases = [x for x in stats if x.get('phase') == 'decode']
    return {
        'exit_code': proc.returncode, 'text': text, 'peak_rss_mb': round(peak_rss / MIB, 1),
        'peak_anon_mb': round(peak_anon / MIB, 1), 'peak_swap_mb': round(peak_swap / MIB, 1),
        'anon_within_limit': peak_anon <= level_mb * MIB, 'swap_free': peak_swap == 0,
        'decode': ({'tokens': int(rates[-1][0]), 'seconds': float(rates[-1][1]),
                    'tok_s': float(rates[-1][2]), 'ssd_mb': float(rates[-1][3]),
                    'ssd_mb_per_token': float(rates[-1][4])} if rates else None),
        'pager': (phases[-1] if phases else {}), 'stderr_path': str(log_path),
    }


def weather(city: str, timeout: int) -> dict:
    q = urllib.parse.urlencode({
        'name': city, 'count': 1, 'language': 'en', 'format': 'json',
    })
    with urllib.request.urlopen('https://geocoding-api.open-meteo.com/v1/search?' + q,
                                timeout=timeout) as r:
        loc = json.load(r)['results'][0]
    q = urllib.parse.urlencode({
        'latitude': loc['latitude'], 'longitude': loc['longitude'],
        'current': 'temperature_2m,relative_humidity_2m,precipitation',
        'timezone': 'auto',
    })
    with urllib.request.urlopen('https://api.open-meteo.com/v1/forecast?' + q,
                                timeout=timeout) as r:
        cur = json.load(r)['current']
    return {'city': loc['name'], 'country': loc.get('country'), **cur}


def execute_tool(call: str, args: argparse.Namespace) -> tuple[str, dict]:
    call = call.replace('\\n', '\n')
    match = re.search(r'(?:<tool_call>\s*)?(get_weather|compile_project)\b', call)
    if not match:
        raise ValueError(f'no supported tool call in model output: {call!r}')
    name = match.group(1)
    if name == 'get_weather':
        city = re.search(r'<arg_key>city</arg_key>\s*<arg_value>\s*([^<]+)', call)
        if not city:
            city = re.search(r'<arg_key>city</arg_key>\s*<arg_value>\s*([^<]+)', call.replace('\\n', '\n'))
        if not city:
            city = re.search(r'<city</arg_key>\s*=\s*["\']([^"\']+)', call)
        if not city and re.search(r'<tool_call>\s*get_weather\b', call):
            city = re.search(r'\bTokyo\b', call, re.I)
        if not city:
            raise ValueError(f'weather call missing city: {call!r}')
        result = weather(city.group(1).strip() if city.lastindex else city.group(0), args.tool_timeout)
        return name, result

    project = re.search(r'<arg_key>project</arg_key>\s*<arg_value>\s*([^<]+)', call)
    if not project:
        project = re.search(r'<project</arg_key>\s*=\s*["\']([^"\']+)', call)
    project_name = project.group(1).strip() if project else 'tiny_fixture'
    if project_name != 'tiny_fixture':
        raise ValueError(f'unknown permitted project {project_name!r}')
    build_dir = Path(args.fixture).resolve() / 'build'
    if build_dir.exists():
        shutil.rmtree(build_dir)
    build_dir.mkdir(parents=True, exist_ok=True)
    p = subprocess.run(['cmake', '-S', str(Path(args.fixture).resolve()), '-B', str(build_dir),
                        '-DCMAKE_BUILD_TYPE=Release'], capture_output=True, text=True,
                       timeout=args.tool_timeout, check=False)
    if p.returncode == 0:
        p = subprocess.run(['cmake', '--build', str(build_dir), '-j', '1'],
                           capture_output=True, text=True, timeout=args.tool_timeout,
                           check=False)
    result = {'returncode': p.returncode, 'tool': 'cmake --build',
              'tail': (p.stdout + '\n' + p.stderr)[-1500:]}
    if p.returncode:
        raise RuntimeError(f'fixture compile failed: {result}')
    return name, result


def run_once(args: argparse.Namespace, level_mb: int, repeat: int) -> dict:
    outdir = Path(args.output_dir)
    outdir.mkdir(parents=True, exist_ok=True)
    prefix = f'level_{level_mb}_r{repeat + 1}'
    log_path = outdir / f'{prefix}.log'
    if args.evict_helper:
        subprocess.run(['sync'], check=False)
        subprocess.run([args.evict_helper, '--evict', str(Path(args.model).resolve())], check=False)
    t0 = time.perf_counter()
    tool_records = []
    try:
        session_peaks = []
        for tool_name, user, arg, value in (
            ('get_weather', 'Get the current weather in Tokyo.', 'city', 'Tokyo'),
            ('compile_project', 'Build the tiny_fixture CMake project.', 'project', 'tiny_fixture'),
        ):
            prompt = (f'<role>SYSTEM</role>{SYSTEM}<|role_end|>'
                      f'<role>HUMAN</role>{user} You MUST call {tool_name} with {arg}={value}. '
                      f'Output only this exact tool call: <tool_call>{tool_name}<arg_key>{arg}</arg_key>'
            f'<arg_value>{value}</arg_value></tool_call><|role_end|>'
                      '<role>ASSISTANT</role><think></think>')
            generated = model_call(args, level_mb, prompt, f'{prefix}_{tool_name}', args.max_tokens)
            session_peaks.append(generated)
            if generated['exit_code'] != 0:
                raise RuntimeError(f'model call failed for {tool_name}: {generated}')
            tool_started = time.perf_counter()
            called, result = execute_tool(generated['text'], args)
            tool_seconds = time.perf_counter() - tool_started
            if called != tool_name:
                raise ValueError(f'expected {tool_name}, model called {called}: {generated["text"]!r}')
            tool_records.append({'tool': called, 'call_output': generated['text'],
                                 'result': result, 'tool_seconds': round(tool_seconds, 3),
                                 'model_call': generated})
            returned = (f'<role>SYSTEM</role>Summarize this tool result in one short sentence. '
                        f'detailed thinking off<|role_end|><role>HUMAN</role>{json.dumps(result, ensure_ascii=False)}'
                        '<|role_end|><role>ASSISTANT</role>\n<think></think>')
            summary = model_call(args, level_mb, returned, f'{prefix}_{tool_name}_return', args.max_tokens)
            session_peaks.append(summary)
            tool_records[-1]['tool_result_summary'] = summary
        wall = time.perf_counter() - t0
        all_decodes = [x['decode'] for x in session_peaks if x['decode']]
        pagers = [x['pager'] for x in session_peaks if x['pager']]
        peak_anon = max(x['peak_anon_mb'] for x in session_peaks)
        peak_rss = max(x['peak_rss_mb'] for x in session_peaks)
        peak_swap = max(x['peak_swap_mb'] for x in session_peaks)
        res = {
            'level_mb': level_mb, 'repeat': repeat + 1, 'status': 'OK', 'wall_s': round(wall, 2),
            'tool_calls': tool_records, 'tools_passed': len(tool_records),
            'peak_rss_mb': peak_rss, 'peak_anon_mb': peak_anon, 'peak_swap_mb': peak_swap,
            'anon_within_limit': peak_anon <= level_mb, 'swap_free': peak_swap == 0,
            'decode_segments': all_decodes,
            'mean_decode_tok_s': statistics.mean(x['tok_s'] for x in all_decodes),
            'mean_ssd_mb_per_token': statistics.mean(x['ssd_mb_per_token'] for x in all_decodes),
            'mean_hot_rate_pct': statistics.mean(x['hit_rate'] for x in pagers if 'hit_rate' in x),
            'max_expert_requests': max(x.get('expert_requests', 0) for x in pagers),
        }
        return res
    except Exception as exc:
        return {'level_mb': level_mb, 'repeat': repeat + 1, 'status': 'FAIL',
                'error': repr(exc)}


def parse_map(raw: str) -> dict[int, int]:
    return {int(k): int(v) for k, v in (part.split(':') for part in raw.split(','))}


def write_summary(path: Path, runs: list[dict], repeat: int) -> None:
    rows = []
    for level in sorted({r['level_mb'] for r in runs}):
        group = [r for r in runs if r['level_mb'] == level and r.get('status') == 'OK']
        if not group:
            rows.append(f'| {level // 1024} GiB | {len(group)}/{repeat} | FAIL | - | - | - | - |')
            continue
        tok = [r['mean_decode_tok_s'] for r in group]
        ssd = [r['mean_ssd_mb_per_token'] for r in group]
        hot = [r['mean_hot_rate_pct'] for r in group]
        anon = [r['peak_anon_mb'] for r in group]
        rss = [r['peak_rss_mb'] for r in group]
        wall = [r['wall_s'] for r in group]
        rows.append(
            f"| {level // 1024} GiB | {len(group)}/{repeat} | "
            f"{statistics.mean(tok):.2f} ± {statistics.stdev(tok) if len(tok)>1 else 0:.2f} | "
            f"{statistics.mean(ssd):.1f} ± {statistics.stdev(ssd) if len(ssd)>1 else 0:.1f} | "
            f"{statistics.mean(hot):.1f} ± {statistics.stdev(hot) if len(hot)>1 else 0:.1f} | "
            f"{statistics.mean(anon):.0f} | {statistics.mean(rss):.0f} | "
            f"{statistics.mean(wall):.1f}"
        )
    md = """# Tool-use benchmark: Ling-3.0-tiny Q4_K_M

Each independent repetition performs two real tool cycles: model tool-call generation, actual Open-Meteo geocoding/current weather request for Tokyo, result summarization; then model tool-call generation, clean CMake configure+compile of the local `tiny_fixture`, and result summarization. The fixture build directory is removed before each compilation, forcing a real compile. Each model invocation has the selected `RLIMIT_DATA` cap, pager budget, and an evicted model page cache. Ten repetitions are run per memory level.

`decode tok/s`, SSD MiB/token, hot rate and wall seconds are per-repetition means across the four model calls / pager phases. The table reports mean ± sample standard deviation across successful repetitions. Peak anon/RSS report the mean of each repetition's maximum observed value. Weather readings vary with time and are preserved per run in JSON. RSS can include file-backed mapped pages beyond RLIMIT_DATA.

| RAM | Runs | Decode tok/s mean ± SD | SSD MiB/token mean ± SD | Hot rate % mean ± SD | Peak anon MB mean | Peak RSS MB mean | Wall s mean |
|---:|---:|---:|---:|---:|---:|---:|---:|
""" + "\n".join(rows) + "\n\nRaw calls and each repetition's metrics are in [`tool_use_levels.json`](tool_use_levels.json). The reproducible harness is [`../sdq-ling/tool_use_bench.py`](../sdq-ling/tool_use_bench.py). Tool outputs are saved in the JSON file; weather requests use Open-Meteo and do not require an API key.\n"
    path.write_text(md)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--bin', required=True)
    ap.add_argument('--fixture', default='tools/tiny_fixture')
    ap.add_argument('--levels-mb', nargs='+', type=int, default=[1024, 2048, 4096, 8192, 16384])
    ap.add_argument('--repeat', type=int, default=10)
    ap.add_argument('--ctx', type=int, default=256)
    ap.add_argument('--ubatch', type=int, default=32)
    ap.add_argument('--threads', type=int, default=4)
    ap.add_argument('--max-tokens', type=int, default=80)
    ap.add_argument('--timeout', type=int, default=1800)
    ap.add_argument('--tool-timeout', type=int, default=120)
    ap.add_argument('--evict-helper', default='./pagecache_resident')
    ap.add_argument('--reserve-map', default='1024:700,2048:700,4096:700,8192:1000,16384:1400')
    ap.add_argument('--output-dir', default='results/tool_use_runs')
    ap.add_argument('--out-json', default='results/tool_use_levels.json')
    ap.add_argument('--retry-failed', action='store_true')
    args = ap.parse_args()
    args.reserve_map = parse_map(args.reserve_map)
    Path(args.fixture).mkdir(parents=True, exist_ok=True)
    out_path = Path(args.out_json)
    if args.retry_failed and out_path.exists():
        runs = json.loads(out_path.read_text()).get('runs', [])
    else:
        runs = []
    for level in args.levels_mb:
        for repeat in range(args.repeat):
            old = next((i for i, r in enumerate(runs)
                        if r.get('level_mb') == level and r.get('repeat') == repeat + 1), None)
            if old is not None and (not args.retry_failed or runs[old].get('status') == 'OK'):
                continue
            print(f'level={level} repeat={repeat + 1}/{args.repeat}', file=sys.stderr, flush=True)
            rec = run_once(args, level, repeat)
            if old is None:
                runs.append(rec)
            else:
                runs[old] = rec
            print(json.dumps(rec, ensure_ascii=False)[:500], file=sys.stderr, flush=True)
            out_path.write_text(json.dumps({'host': facts(), 'config': vars(args), 'runs': runs},
                                           ensure_ascii=False, indent=2))
    write_summary(out_path.with_suffix('.md'), runs, args.repeat)
    print('DONE', file=sys.stderr)


if __name__ == '__main__':
    main()