import os
import time
import hashlib
from urllib.parse import urlparse  # <--- 新增导入
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
import redis
import traceback

# ==========================================
# ⚠️ 入口说明
# ==========================================
# 当前 config.yaml 的 callbacks 使用的是 volc_limit_sync.proxy_handler_instance，
# 不是本文件。除非你把 callbacks 切回 daily_limit.py，否则这里的 MODEL_QUOTAS
# 修改不会影响当前 Docker 容器的监控与限额行为。

# ==========================================
# 🛠️ 调试配置
# ==========================================
# 日志级别配置
# 可选值: "DEBUG", "INFO", "WARNING", "ERROR"
# 修改此变量以控制日志详细程度，或使用环境变量 DAILY_LIMIT_LOG_LEVEL
LOG_LEVEL = os.getenv("DAILY_LIMIT_LOG_LEVEL", "INFO").upper()


def get_log_level_value(level_name):
    levels = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
    return levels.get(level_name, 20)


CURRENT_LEVEL_VAL = get_log_level_value(LOG_LEVEL)
DEBUG_MODE = CURRENT_LEVEL_VAL <= 10


def log(msg, level="INFO"):
    msg_val = get_log_level_value(level)
    if msg_val >= CURRENT_LEVEL_VAL:
        icon = {"DEBUG": "🔍", "INFO": "ℹ️", "WARNING": "⚠️", "ERROR": "❌"}.get(level, "📝")
        print(f"{icon} [DailyLimit] {msg}", flush=True)


def dprint(msg):
    log(msg, level="INFO")


dprint("Loading daily_limit.py module v20 (legacy entry; current config uses volc_limit_sync.py)...")

# ==========================================
# Redis 与 限额配置
# ==========================================
REDIS_HOST = os.getenv("REDIS_HOST", "litellm_redis")
REDIS_PORT = os.getenv("REDIS_PORT", "6379")
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)

MODEL_QUOTAS = {
    "deepseek-v3-2-251201:ba0383d3": 1600000,
    "doubao-seed-1-8-251228:ba0383d3": 1800000,
    "glm-4-7-251222:ba0383d3": 1800000,
    "doubao-seed-2-0-pro-260215:ba0383d3": 1900000,
    "doubao-seed-2-0-lite-260215:ba0383d3": 1900000,
}


try:
    redis_host = os.getenv("REDIS_HOST", "redis")
    redis_port = int(os.getenv("REDIS_PORT", 6379))
    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0, decode_responses=True)
    log(f"✅ Redis connected successfully to {redis_host}:{redis_port}", level="INFO")
except Exception as e:
    log(f"❌ Redis connection failed: {str(e)}", level="ERROR")
    redis_client = None


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


def generate_api_base_hash(api_base):
    """
    生成 api_base 的 Hash，强制忽略 URL 的路径部分。
    例如:
    'https://api.com/v1/' -> 'https://api.com'
    'https://api.com/v3'  -> 'https://api.com'
    两者生成的 Hash 将完全一致。
    """
    if not api_base:
        return "default"

    try:
        # 1. 预处理：去空格
        s = api_base.strip()
        
        # 2. 解析 URL
        parsed = urlparse(s)
        
        # 3. 提取 Scheme 和 Netloc (域名+端口)
        # 如果解析成功且包含 scheme 和 netloc，则重组
        if parsed.scheme and parsed.netloc:
            clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
        else:
            # 兜底：如果不是标准 URL，或者解析失败，尝试简单的去斜杠和小写化
            # 这种情况比较少见，通常 litellm 会校验 URL 格式
            clean_base = s.rstrip('/').lower()
            
    except Exception:
        # 极端情况下的兜底
        clean_base = api_base.strip().rstrip('/').lower()

    # 4. 计算 Hash
    h = hashlib.md5(clean_base.encode()).hexdigest()[:8]
    
    if DEBUG_MODE:
        # Debugging hash mismatch
        # 打印原始值和清洗后的值，方便确认
        print(f"🔍 [Hash Debug] Base='{api_base}' -> Clean='{clean_base}' -> Hash='{h}'", flush=True)
        
    return h


# ==========================================
# ⏰ 时间逻辑：北京时间 9:00 (UTC 01:00) 刷新
# ==========================================
def get_redis_key(clean_name, api_base=None):
    now = time.time()
    # UTC 01:00 视为新的一天，所以减去 1小时 (3600秒)
    adjusted_time = now - 3600
    date_str = time.strftime("%Y-%m-%d", time.gmtime(adjusted_time))

    suffix = ""
    if api_base:
        suffix = f":{generate_api_base_hash(api_base)}"

    return f"daily_usage:{date_str}:{clean_name}{suffix}"


# ==========================================
# 🎯 限额查找逻辑 (支持精确匹配)
# ==========================================
def get_daily_limit(clean_model, api_base):
    # 1. 尝试精确匹配 (Model + Hash)
    # 用户可以在 MODEL_QUOTAS 中配置 "model_name:hash": limit
    current_hash = generate_api_base_hash(api_base)
    specific_key = f"{clean_model}:{current_hash}"

    limit = MODEL_QUOTAS.get(specific_key)
    if limit:
        return limit

    # 2. 尝试精确模型名匹配
    limit = MODEL_QUOTAS.get(clean_model)
    if limit:
        return limit

    # 3. 尝试模糊模型名匹配 (向下兼容)
    for quota_model, limit_val in MODEL_QUOTAS.items():
        # 排除包含冒号的特定规则，防止误匹配
        if ":" not in quota_model and quota_model in clean_model:
            return limit_val

    return None


# ==========================================
# 👻 核心逻辑：软更名踢出 & 自动恢复
# ==========================================
def soft_evict_from_router(target_upstream_model, target_api_base=None):
    """
    不删除节点，而是修改节点的 model_name，让 Router 找不到它。
    同时备份原始信息以便恢复。
    """
    try:
        from litellm.proxy.proxy_server import llm_router

        if not llm_router or not llm_router.model_list:
            return

        evicted = False
        target_upstream_clean = clean_model_name(target_upstream_model)

        # Iterate over the actual list in memory
        for i, model_entry in enumerate(llm_router.model_list):
            params = model_entry.get("litellm_params", {})
            upstream_model = params.get("model", "")
            current_api_base = params.get("api_base", "")

            # 匹配上游模型名称
            if target_upstream_clean not in clean_model_name(upstream_model):
                continue

            # Strict api_base matching if provided
            # ⚠️ 注意：这里比较的是原始字符串，为了确保精准匹配当前运行的实例
            if target_api_base and current_api_base != target_api_base:
                continue

            current_alias = model_entry.get("model_name")

            # 如果它还没有被标记为 blocked
            if current_alias and "blocked" not in current_alias:
                log(
                    f"👻 [ROUTER HIDE] Renaming alias '{current_alias}' -> '{current_alias}-blocked' for {target_upstream_model}",
                    level="INFO",
                )

                # Backup original state
                model_entry["_daily_limit_backup"] = {"model_name": current_alias, "api_base": current_api_base}

                # 1. 修改别名 - Modify IN PLACE
                llm_router.model_list[i]["model_name"] = f"{current_alias}-blocked"

                # 2. 破坏连接
                llm_router.model_list[i]["litellm_params"]["api_base"] = "http://0.0.0.0"

                evicted = True

        if evicted:
            # Re-initialize router to reflect changes immediately if possible
            # But changing model_list in place is usually enough for next lookup
            dprint("✅ Node effectively removed from routing pool.")

    except Exception as e:
        print(f"❌ Router eviction failed: {e}", flush=True)


last_restore_check = 0
RESTORE_CHECK_INTERVAL = 60  # Check every 60 seconds


def restore_blocked_models():
    """
    检查是否有 blocked 的模型配额已重置，如果是则恢复。
    """
    global last_restore_check
    now = time.time()
    if now - last_restore_check < RESTORE_CHECK_INTERVAL:
        return

    last_restore_check = now

    try:
        from litellm.proxy.proxy_server import llm_router

        if not llm_router or not llm_router.model_list:
            return

        restored_count = 0
        for i, model_entry in enumerate(llm_router.model_list):
            backup = model_entry.get("_daily_limit_backup")
            if not backup:
                continue

            # 获取必要信息检查配额
            params = model_entry.get("litellm_params", {})
            upstream_model = params.get("model", "")
            clean_model = clean_model_name(upstream_model)
            original_api_base = backup.get("api_base")

            # Check quota
            daily_limit = get_daily_limit(clean_model, original_api_base)

            should_restore = False
            if not daily_limit:
                should_restore = True
            else:
                key = get_redis_key(clean_model, original_api_base)
                current_usage = redis_client.get(key)
                current_val = float(current_usage) if current_usage else 0.0

                if current_val < float(daily_limit):
                    should_restore = True

            if should_restore:
                dprint(f"♻️ [RESTORE] Restoring {backup['model_name']} (Quota OK)")

                # Restore IN PLACE
                llm_router.model_list[i]["model_name"] = backup["model_name"]
                llm_router.model_list[i]["litellm_params"]["api_base"] = original_api_base
                del llm_router.model_list[i]["_daily_limit_backup"]
                restored_count += 1

        if restored_count > 0:
            dprint(f"✅ Restored {restored_count} models to routing pool.")

    except Exception as e:
        log(f"❌ Restore check failed: {e}", level="ERROR")


# ==========================================
# Input Callback
# ==========================================
def quota_interceptor(kwargs, user_api_key_dict=None):
    """
    Standard LiteLLM input callback signature:
    def input_callback(kwargs, user_api_key_dict):
    """
    try:
        restore_blocked_models()
        # check_and_hide(kwargs, source="input_callback")  # Removed to avoid rigid blocking
    except (ServiceUnavailableError, ContextWindowExceededError):
        raise
    except Exception:
        # Don't let other errors block the request
        pass


# ==========================================
# Handler & Logging
# ==========================================
class DailyLimitHandler(CustomLogger):
    def __init__(self):
        # Register the callback properly
        if quota_interceptor not in litellm.input_callback:
            litellm.input_callback.append(quota_interceptor)
            dprint("✅ quota_interceptor registered")
        if self not in litellm.callbacks:
            litellm.callbacks.append(self)
            dprint("✅ DailyLimitHandler registered")

    # 统计逻辑
    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        self._record_usage(kwargs, response_obj)

    def log_success_event(self, kwargs, response_obj, start_time, end_time):
        self._record_usage(kwargs, response_obj)

    def _record_usage(self, kwargs, response_obj):
        if not redis_client:
            return
        try:
            litellm_params = kwargs.get("litellm_params") or {}
            raw_model = litellm_params.get("model") or kwargs.get("model")
            clean_model = clean_model_name(raw_model)
            api_base = litellm_params.get("api_base")

            usage = None
            if hasattr(response_obj, "usage"):
                usage = response_obj.usage
            elif isinstance(response_obj, dict):
                usage = response_obj.get("usage")
            if hasattr(usage, "model_dump"):
                usage = usage.model_dump()
            elif hasattr(usage, "dict"):
                usage = usage.dict()

            if not usage:
                return

            total_tokens = usage.get("total_tokens", 0)
            if clean_model and total_tokens > 0:
                key = get_redis_key(clean_model, api_base)
                current_total = redis_client.incrby(key, int(total_tokens))
                redis_client.expire(key, 172800)

                # Get Limit for display
                daily_limit = get_daily_limit(clean_model, api_base)
                display_limit = f"/{int(daily_limit):,}" if daily_limit else ""

                # Show usage if DEBUG level is enabled
                key_display = key.split(":", 2)[-1]
                log(
                    f"📊 [Usage] {key_display}: +{total_tokens} (Today: {current_total:,}{display_limit})",
                    level="INFO",
                )

                # Enforcement Logic
                if daily_limit and current_total >= float(daily_limit):
                    log(f"🚫 [BLOCK] {clean_model} Limit: {current_total}/{daily_limit}", level="WARNING")
                    soft_evict_from_router(raw_model, api_base)
            else:
                if DEBUG_MODE:
                    if not clean_model:
                        log("⚠️ [Usage Skip] No clean model name found", level="DEBUG")
                    if total_tokens <= 0:
                        log(f"⚠️ [Usage Skip] Total tokens is {total_tokens}", level="DEBUG")

        except Exception as e:
            log(f"❌ _record_usage failed: {e}\n{traceback.format_exc()}", level="ERROR")


proxy_handler_instance = DailyLimitHandler()