from __future__ import annotations

import argparse
import sys
import urllib.error

from common import ROOT, call_chat_completion, read_json


REGISTRY_INDEX = ROOT / "registry" / "index.json"
DEFAULT_BASE_URL = "http://docker-gpu.cin.su:8000"
PLUGIN_TASKS = {
    "text": {"text", "chat", "summarization", "code", "tool-use"},
    "translation": {"translation"},
    "audio": {"speech-to-text", "speech-translation"},
    "video": {"video", "image-understanding", "document-understanding", "visual-question-answering"},
    "1c": {"1c", "1c-rag", "bsl-code", "metadata-safety", "1c-query"},
}
DEFAULT_PROMPTS = {
    "text": "Кратко объясни, зачем нужен реестр локальных моделей.",
    "translation": "Переведи на английский: Проверяем локальную модель перевода.",
    "audio": "Составь чек-лист проверки качества распознавания речи.",
    "video": "Какие тесты нужны для проверки модели анализа видео?",
    "1c": "Какие метаданные 1С нужно запросить перед изменением документа?",
}
STATUS_RANK = {
    "production": 0,
    "staging": 1,
    "candidate": 2,
    "draft": 3,
    "archived": 9,
}


def model_matches_plugin(model: dict, plugin: str) -> bool:
    tasks = set(model.get("task") or [])
    return bool(tasks & PLUGIN_TASKS.get(plugin, set()))


def model_score(model: dict, plugin: str) -> tuple[int, int, str]:
    tasks = set(model.get("task") or [])
    status = str(model.get("status") or "")
    score = 0
    if plugin == "text":
        if tasks & {"text", "chat"}:
            score -= 5
        if tasks & {"1c", "1c-rag", "bsl-code"}:
            score += 6
    elif plugin == "1c":
        if tasks & PLUGIN_TASKS["1c"]:
            score -= 6
    else:
        if tasks & PLUGIN_TASKS.get(plugin, set()):
            score -= 5
    return (score, STATUS_RANK.get(status, 8), str(model.get("id") or ""))


def select_model(plugin: str, model_id: str | None) -> dict:
    registry = read_json(REGISTRY_INDEX)
    models = registry.get("models") or []
    candidates = [model for model in models if model_matches_plugin(model, plugin)]
    if model_id:
        candidates = [model for model in candidates if model.get("id") == model_id]
    candidates.sort(key=lambda model: model_score(model, plugin))
    if not candidates:
        raise ValueError(f"No model found for plugin={plugin!r} model={model_id!r}")
    return candidates[0]


def main() -> int:
    parser = argparse.ArgumentParser(description="Chat with a registry model through an OpenAI-compatible endpoint.")
    parser.add_argument("--plugin", choices=sorted(PLUGIN_TASKS), default="text")
    parser.add_argument("--model-id", help="Registry model id. Defaults to the first model for the plugin.")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--prompt")
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--max-tokens", type=int, default=1000)
    parser.add_argument("--system", default="Отвечай по-русски, кратко и по делу.")
    args = parser.parse_args()

    try:
        model = select_model(args.plugin, args.model_id)
    except ValueError as exc:
        print(exc, file=sys.stderr)
        return 1

    served_model = model.get("served_model_name") or model.get("id")
    prompt = args.prompt or DEFAULT_PROMPTS[args.plugin]
    messages = [
        {"role": "system", "content": args.system},
        {"role": "user", "content": prompt},
    ]

    print(f"Plugin: {args.plugin}")
    print(f"Model: {model.get('id')} -> {served_model}")
    print(f"Endpoint: {args.base_url}")
    print("")

    try:
        answer = call_chat_completion(
            base_url=args.base_url,
            model=served_model,
            messages=messages,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
        )
    except (urllib.error.URLError, ValueError) as exc:
        print(f"Chat request failed: {exc}", file=sys.stderr)
        return 1

    print(answer)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())