#!/usr/bin/env python3
"""Loopback-only CPU BGE v2-m3 service matching MemOS HTTP reranking."""
import argparse
import os
import threading
import time


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-path', required=True)
    parser.add_argument('--port', type=int, default=18086)
    parser.add_argument('--threads', type=int, default=8)
    parser.add_argument('--batch-size', type=int, default=8)
    parser.add_argument('--max-length', type=int, default=1024)
    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
    import torch
    import uvicorn
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel, Field
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    torch.set_num_threads(args.threads)
    torch.set_num_interop_threads(1)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, local_files_only=True)
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_path, local_files_only=True).to('cpu').eval()
    lock = threading.Lock()
    stats = {'requests': 0, 'documents': 0, 'seconds': 0.0}
    app = FastAPI(title='MemOS CPU BGE reranker')

    class Request(BaseModel):
        model: str = 'bge-reranker-v2-m3'
        query: str
        documents: list[str] = Field(max_length=1000)

    @app.get('/health')
    def health():
        return {'status': 'ok', 'model': 'bge-reranker-v2-m3', 'device': 'cpu',
                'score': 'sigmoid(raw_sequence_classification_logit)',
                'max_length': args.max_length, 'batch_size': args.batch_size,
                'threads': args.threads, **stats}

    @app.post('/rerank')
    def rerank(request: Request):
        if request.model not in ('bge-reranker-v2-m3', 'BAAI/bge-reranker-v2-m3'):
            raise HTTPException(400, 'Unsupported reranker model')
        if not request.documents:
            return {'results': []}
        started = time.monotonic()
        scores = []
        with lock, torch.inference_mode():
            try:
                for offset in range(0, len(request.documents), args.batch_size):
                    pairs = [[request.query, doc] for doc in request.documents[offset:offset+args.batch_size]]
                    encoded = tokenizer(pairs, padding=True, truncation=True,
                                        return_tensors='pt', max_length=args.max_length)
                    logits = model(**encoded, return_dict=True).logits.view(-1).float()
                    scores.extend(torch.sigmoid(logits).tolist())
                stats['requests'] += 1
                stats['documents'] += len(request.documents)
                stats['seconds'] += time.monotonic() - started
            except Exception as error:
                raise HTTPException(500, f'Reranking failed: {type(error).__name__}') from None
        if len(scores) != len(request.documents):
            raise HTTPException(500, 'Reranker output length mismatch')
        return {'results': sorted(
            [{'index': i, 'relevance_score': score} for i, score in enumerate(scores)],
            key=lambda row: (-row['relevance_score'], row['index']))}

    uvicorn.run(app, host='127.0.0.1', port=args.port, access_log=False, log_level='warning')


if __name__ == '__main__':
    main()