import ast
import hashlib
import os
import re
from urllib.parse import urlparse

import yaml

# ==========================================
# 🛠️ 配置
# ==========================================
CONFIG_FILE = "config.yaml"
ENV_FILE = ".env"
RUNTIME_FILE = "volc_limit_sync.py"


def load_env(env_path):
    """读取 .env 文件并返回字典"""
    env_vars = {}
    if not os.path.exists(env_path):
        print(f"⚠️ Warning: {env_path} not found. Env var substitution might fail.")
        return env_vars

    with open(env_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            try:
                key, value = line.split("=", 1)
                env_vars[key.strip()] = value.strip().strip('"').strip("'")
            except ValueError:
                continue
    return env_vars


def resolve_env_var(value, env_vars):
    """
    处理环境变量替换，支持以下两种格式：
    1. ${VAR_NAME} - 标准格式
    2. os.environ/VAR_NAME - LiteLLM 特有格式
    """
    if not isinstance(value, str):
        return value

    if value.startswith("os.environ/"):
        var_name = value.replace("os.environ/", "")
        return env_vars.get(var_name, os.getenv(var_name, value))

    pattern = re.compile(r"\$\{?([A-Z0-9_]+)\}?")

    def replace_match(match):
        var_name = match.group(1)
        return env_vars.get(var_name, os.getenv(var_name, match.group(0)))

    return pattern.sub(replace_match, value)


def clean_model_name(name):
    if not name:
        return ""
    return (
        name.replace("openai/", "")
        .replace("azure/", "")
        .replace("anthropic/", "")
        .replace("volcengine/", "")
    )


def generate_api_base_hash(api_base):
    """
    [同步更新] 生成 api_base 的 Hash，强制忽略 URL 的路径部分。
    保持与 volc_limit_sync.py 逻辑一致。
    """
    if not api_base:
        return "default"

    try:
        s = api_base.strip()
        parsed = urlparse(s)
        if parsed.scheme and parsed.netloc:
            clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
        else:
            clean_base = s.rstrip("/").lower()
    except Exception:
        clean_base = api_base.strip().rstrip("/").lower()

    return hashlib.md5(clean_base.encode()).hexdigest()[:8]


def strip_model_version_suffix(model_name):
    """Convert foo-bar-260215 -> foo-bar for common versioned upstream names."""
    if not model_name:
        return ""

    parts = model_name.rsplit("-", 1)
    if len(parts) == 2 and len(parts[1]) == 6 and parts[1].isdigit():
        return parts[0]

    return model_name


def load_python_constant(file_path, constant_name, default=None):
    """Safely read a literal constant from volc_limit_sync.py without importing it."""
    if not os.path.exists(file_path):
        return default

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            source = f.read()
        module = ast.parse(source, filename=file_path)
    except Exception as e:
        print(f"⚠️ Warning: failed to parse {file_path}: {e}")
        return default

    for node in module.body:
        if not isinstance(node, ast.Assign):
            continue
        for target in node.targets:
            if isinstance(target, ast.Name) and target.id == constant_name:
                try:
                    return ast.literal_eval(node.value)
                except Exception as e:
                    print(
                        f"⚠️ Warning: failed to read constant {constant_name} from {file_path}: {e}"
                    )
                    return default

    return default


def build_monitored_model_aliases(monitored_models, manual_aliases):
    """Mirror the monitored-model alias logic used by volc_limit_sync.py."""
    alias_map = {}

    def add_alias(alias_name, target_model):
        clean_alias = clean_model_name(alias_name)
        clean_target = clean_model_name(target_model)
        if not clean_alias or not clean_target:
            return
        alias_map.setdefault(clean_alias, [])
        if clean_target not in alias_map[clean_alias]:
            alias_map[clean_alias].append(clean_target)

    for target_model in sorted(set(monitored_models)):
        add_alias(target_model, target_model)

        base_alias = strip_model_version_suffix(target_model)
        if base_alias != target_model:
            add_alias(base_alias, target_model)

    for alias_name, targets in (manual_aliases or {}).items():
        if isinstance(targets, list):
            for target_model in targets:
                add_alias(alias_name, target_model)
        else:
            add_alias(alias_name, targets)

    return alias_map


def is_monitored_volc_entry(api_base_raw, api_base_resolved):
    if not isinstance(api_base_raw, str):
        return False
    resolved_text = str(api_base_resolved or "").lower()
    return "HUOSHAN_BASE_URL" in api_base_raw or "volcengineapi" in resolved_text


def collect_monitored_entries(config, env_vars):
    entries = {}
    model_list = config.get("model_list", [])

    for model in model_list:
        litellm_params = model.get("litellm_params", {})
        api_base_raw = litellm_params.get("api_base", "")
        api_base_resolved = resolve_env_var(api_base_raw, env_vars)

        if not is_monitored_volc_entry(api_base_raw, api_base_resolved):
            continue

        model_alias = model.get("model_name", "unknown")
        upstream_model = clean_model_name(litellm_params.get("model", ""))
        if not upstream_model:
            continue

        base_hash = generate_api_base_hash(api_base_resolved)
        final_key = f"{upstream_model}:{base_hash}"

        entry = entries.setdefault(
            final_key,
            {
                "config_aliases": [],
                "upstream_model": upstream_model,
                "api_base_raw": api_base_raw,
                "api_base_resolved": api_base_resolved,
                "base_hash": base_hash,
                "final_key": final_key,
            },
        )

        if model_alias not in entry["config_aliases"]:
            entry["config_aliases"].append(model_alias)

    return list(entries.values())


def get_limit_rule(clean_model, base_hash, hard_limits):
    specific_key = f"{clean_model}:{base_hash}"
    if specific_key in hard_limits:
        return specific_key, hard_limits[specific_key]

    if clean_model in hard_limits:
        return clean_model, hard_limits[clean_model]

    for quota_model, limit_val in hard_limits.items():
        if quota_model == "__DEFAULT__":
            continue
        if ":" not in quota_model and quota_model in clean_model:
            return quota_model, limit_val

    if "__DEFAULT__" in hard_limits:
        return "__DEFAULT__", hard_limits["__DEFAULT__"]

    return None, None


def build_reverse_alias_map(alias_map):
    reverse_map = {}
    for alias_name, targets in alias_map.items():
        for target_model in targets:
            reverse_map.setdefault(target_model, [])
            if alias_name not in reverse_map[target_model]:
                reverse_map[target_model].append(alias_name)
    return reverse_map


def format_limit_source(rule_name, limit_value, final_key, upstream_model):
    if rule_name is None:
        return "未配置"
    if rule_name == final_key:
        return f"自定义键 {rule_name} => {int(limit_value):,}"
    if rule_name == upstream_model:
        return f"模型级规则 {rule_name} => {int(limit_value):,}"
    if rule_name == "__DEFAULT__":
        return f"默认规则 __DEFAULT__ => {int(limit_value):,}"
    return f"模糊规则 {rule_name} => {int(limit_value):,}"


def print_monitored_audit(entries, reverse_alias_map, hard_limits):
    print("=== 监控模型审计 / config.yaml -> volc_limit_sync.py ===")
    print()

    for entry in entries:
        monitor_aliases = reverse_alias_map.get(
            entry["upstream_model"], [entry["upstream_model"]]
        )
        rule_name, limit_value = get_limit_rule(
            entry["upstream_model"], entry["base_hash"], hard_limits
        )

        print(f"- Config alias(es): {', '.join(entry['config_aliases'])}")
        print(f"  Upstream model : {entry['upstream_model']}")
        print(f"  监控别名       : {', '.join(monitor_aliases)}")
        print(f"  Generated key  : {entry['final_key']}")
        print(
            f"  限额来源       : {format_limit_source(rule_name, limit_value, entry['final_key'], entry['upstream_model'])}"
        )
        print()


def print_maintenance_guide(entries, manual_aliases, hard_limits):
    default_limit = hard_limits.get("__DEFAULT__")

    print("=== 维护入口 ===")
    print("1. 在 config.yaml 的 model_list 中增删火山模型，这是“纳入监控模型”的主入口。")
    print("2. 改完后运行 python generate_keys.py，先看上面的“监控别名”和“限额来源”。")
    print("3. 常规版本号模型（如 xxx-260215）会自动推导成监控别名 xxx，通常不用再改 volc_limit_sync.py。")
    print("4. 只有两种情况才需要改 volc_limit_sync.py 的 MANUAL_MONITORED_MODEL_ALIASES：")
    print("   - 资源包名和上游模型名不一致。")
    print("   - 一个资源包要映射到多个上游模型。")
    if manual_aliases:
        print(
            f"5. 当前手工别名覆盖: {', '.join(sorted(manual_aliases.keys()))}"
        )
    print()

    defaulted_entries = []
    for entry in entries:
        rule_name, _ = get_limit_rule(
            entry["upstream_model"], entry["base_hash"], hard_limits
        )
        if rule_name == "__DEFAULT__":
            defaulted_entries.append(entry)

    if defaulted_entries and default_limit is not None:
        print("📋 可选：以下监控模型当前走 __DEFAULT__，如需单模型硬限额可复制到 HARD_DAILY_LIMITS")
        print("HARD_DAILY_LIMITS = {")
        for entry in defaulted_entries:
            print(f'    "{entry["final_key"]}": {int(default_limit)},')
        print("}")
        print()


def main():
    print("🔍 Auditing monitored Volcengine models from config.yaml...\n")

    env_vars = load_env(ENV_FILE)

    if not os.path.exists(CONFIG_FILE):
        print(f"❌ Error: {CONFIG_FILE} not found!")
        return

    try:
        with open(CONFIG_FILE, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
    except Exception as e:
        print(f"❌ Error parsing yaml: {e}")
        return

    entries = collect_monitored_entries(config or {}, env_vars)
    if not entries:
        print("⚠️ No monitored Volcengine models found in config.yaml")
        return

    entries.sort(key=lambda item: item["upstream_model"])

    manual_aliases = load_python_constant(
        RUNTIME_FILE, "MANUAL_MONITORED_MODEL_ALIASES", default={}
    ) or {}
    hard_limits = load_python_constant(RUNTIME_FILE, "HARD_DAILY_LIMITS", default={}) or {}

    monitored_upstreams = [entry["upstream_model"] for entry in entries]
    alias_map = build_monitored_model_aliases(monitored_upstreams, manual_aliases)
    reverse_alias_map = build_reverse_alias_map(alias_map)

    print_monitored_audit(entries, reverse_alias_map, hard_limits)
    print_maintenance_guide(entries, manual_aliases, hard_limits)


if __name__ == "__main__":
    main()