from pathlib import Path
from typing import Any

from app.core.logging import logger


async def test_model(model_id: str, prompt: str, max_new_tokens: int = 128, temperature: float = 0.8, top_p: float = 0.95) -> dict[str, Any]:
    """加载已缓存模型并生成测试响应。"""
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        from app.services.model_service import resolve_model_path
        model_path = await resolve_model_path(model_id)
        if not model_path:
            return {"error": f"Model not found in cache: {model_id}"}

        model_dir = Path(model_path)
        if not (model_dir / "config.json").exists():
            return {"error": f"Model directory not found: {model_dir}"}

        tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        model.eval()

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=temperature > 0,
                pad_token_id=tokenizer.eos_token_id,
            )

        generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

        return {
            "model_id": model_id,
            "prompt": prompt,
            "generated_text": generated_text,
        }

    except Exception as e:
        logger.error(f"Model test failed: {e}")
        return {"error": str(e)}