import argparse
import json
import os
import sys
from typing import Dict

# 兼容：支持直接运行脚本或用 -m 模块方式运行
if __name__ == "__main__" and __package__ is None:
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 统一使用绝对导入，避免相对导入在脚本直跑时失败
from exp_multimodal.labels import build_labels
from exp_multimodal.vlm_classifier import classify_image, classify_image_openset
from exp_multimodal.ollama_client import OLLAMA_URL, DEFAULT_MODEL
from exp_multimodal.text_embedder import OllamaEmbedder
from exp_multimodal.vector_matcher import DishNameMatcher
from exp_multimodal.vlm_providers.ollama_provider import OllamaVLMProvider


FEWSHOT_HINTS: Dict[str, str] = {
    # 可选：仅对菜品模式提供少量文字提示，帮助区分相似菜
    # "西红柿鸡蛋": "红黄对比明显，蛋块与番茄块同炒，汤汁偏红",
    # "麻婆豆腐": "豆腐块+红油辣椒+花椒，肉末点缀",
}



def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["dish", "whole", "processed", "openset_dish"], default="openset_dish")
    ap.add_argument("--image", default=r"D:\MyProjects\PythonProjects\FoodClassifier\dataset\DishClassification\test\红烧肉\img04.png")
    ap.add_argument("--alias_map", default=None)
    
    # 开放式识别专用参数
    ap.add_argument("--openset-top-k", type=int, default=3, help="开放式识别：向量检索Top-K候选数")
    ap.add_argument("--openset-min-score", type=float, default=0.5, help="开放式识别：最低匹配分数阈值")
    ap.add_argument("--openset-index-path", default="exp_multimodal/dish_name_index", help="开放式识别：向量索引路径")

    args = ap.parse_args()

    print(
        f"[Main] mode={args.mode} image={args.image} alias_map={args.alias_map} "
        f"env_OLLAMA_URL={os.getenv('OLLAMA_URL')} env_VLM_MODEL={os.getenv('VLM_MODEL')} "
        f"defaults url={OLLAMA_URL} model={DEFAULT_MODEL}"
    )
    print(f"[Main] Image exists={os.path.exists(args.image)} size={os.path.getsize(args.image) if os.path.exists(args.image) else 'N/A'}")

    # 开放式识别分支
    if args.mode == "openset_dish":
        print(f"[Main] Openset mode: top_k={args.openset_top_k} min_score={args.openset_min_score} index_path={args.openset_index_path}")
        
        # 初始化Provider、Embedder和Matcher
        provider = OllamaVLMProvider(base_url=OLLAMA_URL, model_name=DEFAULT_MODEL)
        embedder = OllamaEmbedder(base_url=OLLAMA_URL, model_name="bge-large-zh-v1.5")
        matcher = DishNameMatcher(index_dir=args.openset_index_path)
        
        print("[Main] Loading openset components...")
        matcher.load()
        print(f"[Main] Loaded {len(matcher.dish_names)} dish names from index")
        
        res = classify_image_openset(
            image_path=args.image,
            provider=provider,
            embedder=embedder,
            matcher=matcher,
            top_k=args.openset_top_k,
            min_match_score=args.openset_min_score,
        )
        
        print(
            json.dumps(
                {"mode": args.mode, "result": res},
                ensure_ascii=False,
                indent=2,
            )
        )
    
    # 封闭式识别分支（原有逻辑）
    else:
        labels = build_labels(args.mode, args.alias_map)
        ingredient_only = args.mode in {"whole", "processed"}
        fewshot = FEWSHOT_HINTS if args.mode == "dish" else None
        print(f"[Main] Built labels count={len(labels)} ingredient_only={ingredient_only} fewshot={bool(fewshot)}")

        provider = OllamaVLMProvider(base_url=OLLAMA_URL, model_name=DEFAULT_MODEL)
        
        res = classify_image(
            args.image,
            labels,
            provider=provider,
            fewshot_hints=fewshot,
            ingredient_only=ingredient_only,
        )

        print(
            json.dumps(
                {"mode": args.mode, "result": res, "num_labels": len(labels)},
                ensure_ascii=False,
            )
        )


if __name__ == "__main__":
    main()