import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile


FILES = ('server.py', 'benchmark.py', 'requirements.txt')
CUDA_INDEX = 'https://download.pytorch.org/whl/cu126'
MODEL_PRELOAD_TIMEOUT = 900


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def python_path(target):
    return target / ('Scripts/python.exe' if os.name == 'nt' else 'bin/python')


def ensure_runtime(target, dry_run):
    if dry_run:
        return python_path(target / '.venv')
    venv = target / '.venv'
    if not venv.exists():
        subprocess.run([sys.executable, '-m', 'venv', str(venv)], check=True, timeout=180)
    python = python_path(venv)
    if not python.is_file():
        raise ValueError('Semantic retrieval virtual environment is incomplete.')
    arguments = [str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')]
    if shutil.which('nvidia-smi'):
        try:
            subprocess.run([*arguments, '--extra-index-url', CUDA_INDEX], check=True, timeout=900)
        except subprocess.CalledProcessError:
            print('Semantic retrieval: CUDA runtime installation failed; falling back to CPU.', file=sys.stderr)
            subprocess.run(arguments, check=True, timeout=900)
    else:
        subprocess.run(arguments, check=True, timeout=900)
    return python


def ensure_models(python, target):
    subprocess.run([str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')],
                   check=True, timeout=MODEL_PRELOAD_TIMEOUT)


def load_state(path):
    if not path.exists():
        return {'version': 1, 'clients': [], 'files': {}}
    state = json.loads(path.read_text(encoding='utf-8'))
    if state.get('version') != 1 or not isinstance(state.get('clients'), list) or not isinstance(state.get('files'), dict):
        raise ValueError('Invalid semantic retrieval ownership state.')
    if any(client not in ('codex', 'claude') for client in state['clients']):
        raise ValueError('Invalid semantic retrieval client state.')
    if any(name not in FILES or not isinstance(value, str) or len(value) != 64 or any(character not in '0123456789abcdef' for character in value) for name, value in state['files'].items()):
        raise ValueError('Invalid semantic retrieval file state.')
    return state


def save_state(path, state, dry_run):
    if not dry_run:
        path.parent.mkdir(parents=True, exist_ok=True)
        temporary = None
        try:
            with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', dir=path.parent, delete=False) as stream:
                temporary = Path(stream.name)
                json.dump(state, stream, indent=2)
                stream.write('\n')
                stream.flush()
                os.fsync(stream.fileno())
            os.replace(temporary, path)
        finally:
            if temporary is not None:
                temporary.unlink(missing_ok=True)


def check_runtime(target, state):
    if target.is_symlink():
        raise ValueError('Semantic retrieval directory must not be a symbolic link.')
    if not target.exists():
        return
    if not target.is_dir():
        raise ValueError('Semantic retrieval target is not a directory.')
    for name, expected in state['files'].items():
        path = target / name
        if path.is_symlink() or not path.is_file() or digest(path) != expected:
            raise ValueError(f'Owned semantic retrieval file changed: {name}')
    allowed = set(state['files']) | {'.venv', 'model-cache', '__pycache__', 'data'}
    for path in target.iterdir():
        if path.is_symlink() or path.name not in allowed:
            raise ValueError(f'Foreign or linked semantic retrieval path: {path.name}')
        if path.name == 'data' and (not path.is_dir() or any(
                child.is_symlink() or not (
                    child.is_file() and (child.name in ('daemon.key', 'daemon.starting') or re.fullmatch(
                        r'[0-9a-f]{64}\.sqlite3(?:-(?:wal|shm|journal))?|daemon-[0-9a-f]{12}\.starting', child.name))
                    or child.is_socket() and re.fullmatch(r'daemon-[0-9a-f]{16}\.sock', child.name))
                for child in path.iterdir())):
            raise ValueError('Foreign or linked semantic retrieval path: data')
        if path.name == '__pycache__' and (not path.is_dir() or any(
                child.is_symlink() or not child.is_file() or not re.fullmatch(
                    r'(?:server|benchmark)\.[A-Za-z0-9_-]+(?:\.opt-\d+)?\.pyc', child.name)
                for child in path.iterdir())):
            raise ValueError('Foreign or linked semantic retrieval path: __pycache__')


def sync(root, home, clients, enabled, dry_run, update, summary=False):
    base = home / '.my-ai-configuration'
    target = base / 'semantic-retrieval'
    state_path = base / 'semantic-retrieval.json'
    state = load_state(state_path)
    clients = list(dict.fromkeys(clients))
    if base.exists() and base.is_symlink():
        raise ValueError('Semantic retrieval state directory must not be a symbolic link.')
    if not enabled:
        owned = [client for client in state['clients'] if client in clients]
        if dry_run:
            if not summary:
                print('DRYRUN disable semantic retrieval: ' + (', '.join(owned) if owned else 'nothing owned'))
            return
        state['clients'] = [client for client in state['clients'] if client not in clients]
        if state['clients']:
            save_state(state_path, state, False)
            return
        if not state['files']:
            state_path.unlink(missing_ok=True)
            return
        check_runtime(target, state)
        if target.exists() and target.is_dir():
            shutil.rmtree(target)
        state_path.unlink(missing_ok=True)
        return

    check_runtime(target, state)
    if target.exists() and not state['files']:
        raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.')
    if dry_run:
        if not summary:
            print('DRYRUN enable semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-0.6B')
        return
    target.mkdir(parents=True, exist_ok=True)
    source = root / 'shared' / 'retrieval'
    for name in FILES:
        shutil.copy2(source / name, target / name)
        state['files'][name] = digest(target / name)
    save_state(state_path, state, False)
    python = ensure_runtime(target, False)
    ensure_models(python, target)
    state['clients'] = list(dict.fromkeys(state['clients']))
    save_state(state_path, state, False)
    for client in clients:
        if client not in state['clients']:
            state['clients'].append(client)
            save_state(state_path, state, False)
    save_state(state_path, state, False)


def benchmark(root, home, dry_run, summary=False):
    base = home / '.my-ai-configuration'
    target = base / 'semantic-retrieval'
    state_path = base / 'semantic-retrieval.json'
    if dry_run:
        if not summary:
            print('DRYRUN benchmark semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-0.6B')
        return True
    state = load_state(state_path)
    source = root / 'shared' / 'retrieval'
    if state['files']:
        check_runtime(target, state)
    elif target.is_symlink():
        raise ValueError('Semantic retrieval directory must not be a symbolic link.')
    if target.exists() and not state['files']:
        if not all((target / name).is_file() and digest(target / name) == digest(source / name) for name in FILES):
            raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.')
        state['files'] = {name: digest(target / name) for name in FILES}
        check_runtime(target, state)
    target.mkdir(parents=True, exist_ok=True)
    for name in FILES:
        shutil.copy2(source / name, target / name)
        state['files'][name] = digest(target / name)
    python = ensure_runtime(target, False)
    result = subprocess.run([str(python), str(target / 'benchmark.py'), '--model-cache', str(target / 'model-cache')], check=True,
                            stdout=subprocess.PIPE, text=True, timeout=300)
    outcome = json.loads(result.stdout)
    if not isinstance(outcome, dict) or not isinstance(outcome.get('recommended'), bool):
        raise ValueError('Semantic retrieval benchmark returned an invalid result.')
    if outcome['recommended'] or state['clients']:
        save_state(state_path, state, False)
    else:
        check_runtime(target, state)
        shutil.rmtree(target)
        state_path.unlink(missing_ok=True)
    if not summary:
        print('Semantic retrieval warm-model check: ' + ('passed' if outcome['recommended'] else 'failed') + f" (embedding {outcome['embeddingSeconds']}s, reranking {outcome['rerankingSeconds']}s; excludes model startup and indexing)")
    return outcome['recommended']


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=('sync', 'benchmark'))
    parser.add_argument('--root', required=True, type=Path)
    parser.add_argument('--home', required=True, type=Path)
    parser.add_argument('--client', choices=('codex', 'claude', 'both'))
    parser.add_argument('--enabled', choices=('true', 'false'))
    parser.add_argument('--dry-run', action='store_true')
    parser.add_argument('--update', action='store_true')
    parser.add_argument('--summary', action='store_true')
    args = parser.parse_args()
    if args.action == 'benchmark':
        recommended = benchmark(args.root.resolve(), args.home.resolve(), args.dry_run, args.summary)
        raise SystemExit(0 if recommended else 2)
    if args.client is None or args.enabled is None:
        parser.error('sync requires --client and --enabled')
    clients = ('codex', 'claude') if args.client == 'both' else (args.client,)
    sync(args.root.resolve(), args.home.resolve(), clients, args.enabled == 'true', args.dry_run, args.update, args.summary)
    print('SEMANTIC RETRIEVAL: reconciliation complete' + (' (dry run)' if args.dry_run else ''))


if __name__ == '__main__':
    try:
        main()
    except (OSError, ValueError, json.JSONDecodeError, subprocess.SubprocessError) as error:
        print(f'Semantic retrieval: {error}', file=sys.stderr)
        raise SystemExit(1)