"""AwoooP Token Budget Hard Kill Service
=======================================
ADR-120: 三層 Hard Kill 防護架構
2026-05-04 ogt + Claude Sonnet 4.6（Phase 2.6）

防線：
1. Pre-call check（呼叫前）— Layer 1 Tenant + Layer 2 Platform + Layer 3 Emergency Kill
2. Post-call accounting（呼叫後）— 寫 budget_ledger + 更新 Redis cache
3. 告警閾值通知（80% / 95% Telegram 告警）

注意：Layer 0 Run budget 需要 awooop_run_state（Phase 3 SAGA 實作後補加）
"""
from __future__ import annotations

import time
from decimal import Decimal

import structlog

from src.core.config import settings

logger = structlog.get_logger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# 告警閾值（ADR-120 D4）
# ─────────────────────────────────────────────────────────────────────────────
BUDGET_ALERT_THRESHOLDS = {
    "warn": Decimal("0.80"),
    "critical": Decimal("0.95"),
    "hard_kill": Decimal("1.00"),
}

# Redis key 前綴
_EMERGENCY_KILL_KEY = "platform:budget:emergency_kill"
_TENANT_BUDGET_KEY_PREFIX = "budget:tenant:"   # {project_id}:daily_used_usd
_PLATFORM_BUDGET_KEY = "budget:platform:daily_used_usd"
_BUDGET_CACHE_TTL = 300  # 5 分鐘，每次寫入後 refresh


class BudgetExhaustedError(Exception):
    """LLM call 被 hard kill 攔截"""

    def __init__(self, error_code: str, message: str) -> None:
        self.error_code = error_code
        super().__init__(f"[{error_code}] {message}")


# ─────────────────────────────────────────────────────────────────────────────
# 費用計算（按模型定價估算）
# ─────────────────────────────────────────────────────────────────────────────

# USD per 1M tokens（in + out）
_COST_PER_MILLION_TOKENS: dict[str, tuple[float, float]] = {
    # (prompt_per_M, completion_per_M)
    "claude-opus-4-7":        (15.0, 75.0),
    "claude-sonnet-4-6":      (3.0, 15.0),
    "claude-haiku-4-5":       (0.8, 4.0),
    "gpt-4o":                 (5.0, 15.0),
    "gpt-4o-mini":            (0.15, 0.6),
    "gemini-2.0-flash":       (0.075, 0.3),
    "deepseek-r1:14b":        (0.0, 0.0),   # local Ollama — 無費用
    "qwen3:8b":               (0.0, 0.0),   # local Ollama — 無費用
}
_DEFAULT_COST_PER_M = (3.0, 15.0)  # fallback → claude-sonnet


def estimate_cost(
    prompt_tokens: int,
    completion_tokens: int,
    model: str,
) -> Decimal:
    """估算一次 LLM call 的費用（USD）"""
    prompt_rate, completion_rate = _COST_PER_MILLION_TOKENS.get(
        model, _DEFAULT_COST_PER_M
    )
    cost = (prompt_tokens / 1_000_000 * prompt_rate +
            completion_tokens / 1_000_000 * completion_rate)
    return Decimal(str(round(cost, 6)))


# ─────────────────────────────────────────────────────────────────────────────
# Pre-call Budget Check（ADR-120 D2 防線 1）
# ─────────────────────────────────────────────────────────────────────────────

async def check_budget_before_llm_call(
    project_id: str,
    model: str,
    estimated_prompt_tokens: int = 4000,
    *,
    agent_id: str | None = None,
) -> None:
    """
    LLM call 前的三層 budget check。

    超出任一層預算 → 拋出 BudgetExhaustedError，阻止 API call。
    Redis 不可用時 fail-open（不阻擋呼叫，但記 warning）。

    Args:
        project_id: 租戶 ID
        model: 模型名稱（用於費用估算）
        estimated_prompt_tokens: 預估 prompt token 數（保守估計 × 1.5 已含在外）
    """
    # Layer 3：Emergency Kill Switch（最優先）
    await check_emergency_kill()

    # Local Ollama 模型無費用，跳過 Layer 1/2
    if model in {"deepseek-r1:14b", "qwen3:8b"} or model.startswith("ollama/"):
        return

    estimated_cost = estimate_cost(estimated_prompt_tokens, 0, model)

    # Layer 2：Tenant Budget
    await _check_tenant_budget(project_id, estimated_cost)

    # Layer 1：Platform Budget
    await _check_platform_budget(estimated_cost)


async def check_emergency_kill() -> None:
    """Layer 3: Emergency Kill Switch — Redis key platform:budget:emergency_kill"""
    try:
        from src.core.redis_client import get_redis
        redis = get_redis()
        if await redis.exists(_EMERGENCY_KILL_KEY):
            raise BudgetExhaustedError(
                "E-BUDGET-004",
                "Emergency kill switch activated — contact platform admin",
            )
    except BudgetExhaustedError:
        raise
    except Exception as exc:
        logger.warning("budget_emergency_kill_check_failed", error=str(exc))


async def _check_tenant_budget(project_id: str, estimated_cost: Decimal) -> None:
    """Layer 2: Tenant Budget（Redis 快取 + awooop_projects.budget_limit_usd）"""
    try:
        from src.core.redis_client import get_redis
        redis = get_redis()

        # 讀取 Tenant 每日已用金額
        cache_key = f"{_TENANT_BUDGET_KEY_PREFIX}{project_id}"
        used_raw = await redis.get(cache_key)
        used_usd = Decimal(used_raw.decode() if isinstance(used_raw, bytes) else used_raw or "0")

        # 讀取 Tenant 預算上限（從 awooop_projects 表）
        limit_usd = await _get_tenant_budget_limit(project_id)
        if limit_usd is None:
            return  # 無上限 → 放行

        if used_usd + estimated_cost > limit_usd:
            raise BudgetExhaustedError(
                "E-BUDGET-002",
                f"Tenant {project_id} budget exhausted: "
                f"used ${used_usd:.4f} / ${limit_usd:.4f}",
            )

        # 告警閾值
        usage_pct = (used_usd + estimated_cost) / limit_usd
        if usage_pct >= BUDGET_ALERT_THRESHOLDS["critical"]:
            logger.warning(
                "budget_tenant_critical",
                project_id=project_id,
                usage_pct=float(usage_pct),
                used_usd=float(used_usd),
                limit_usd=float(limit_usd),
            )
        elif usage_pct >= BUDGET_ALERT_THRESHOLDS["warn"]:
            logger.warning(
                "budget_tenant_warn",
                project_id=project_id,
                usage_pct=float(usage_pct),
                used_usd=float(used_usd),
                limit_usd=float(limit_usd),
            )

    except BudgetExhaustedError:
        raise
    except Exception as exc:
        logger.warning("budget_tenant_check_failed", project_id=project_id, error=str(exc))


async def _check_platform_budget(estimated_cost: Decimal) -> None:
    """Layer 1: Platform Budget（config 靜態上限 + Redis 累計）"""
    platform_limit = getattr(settings, "PLATFORM_DAILY_BUDGET_USD", None)
    if not platform_limit:
        return  # 未設定 → 放行

    try:
        from src.core.redis_client import get_redis
        redis = get_redis()
        used_raw = await redis.get(_PLATFORM_BUDGET_KEY)
        used_usd = Decimal(used_raw.decode() if isinstance(used_raw, bytes) else used_raw or "0")
        limit_usd = Decimal(str(platform_limit))

        if used_usd + estimated_cost > limit_usd:
            raise BudgetExhaustedError(
                "E-BUDGET-003",
                f"Platform budget exhausted: used ${used_usd:.4f} / ${limit_usd:.4f} — "
                "all LLM calls suspended",
            )
    except BudgetExhaustedError:
        raise
    except Exception as exc:
        logger.warning("budget_platform_check_failed", error=str(exc))


# ─────────────────────────────────────────────────────────────────────────────
# Post-call Accounting（ADR-120 D2 防線 2）
# ─────────────────────────────────────────────────────────────────────────────

async def record_token_usage(
    *,
    project_id: str,
    model: str,
    provider: str,
    prompt_tokens: int,
    completion_tokens: int,
    agent_id: str | None = None,
    run_id: str | None = None,
) -> Decimal:
    """
    LLM call 完成後記帳。

    1. 計算實際費用
    2. INSERT budget_ledger
    3. 更新 Redis budget cache（async，不阻擋回傳）
    4. 觸發告警閾值通知

    Returns:
        actual_cost_usd
    """
    import asyncio
    from uuid import UUID

    actual_cost = estimate_cost(prompt_tokens, completion_tokens, model)

    # 寫入 budget_ledger（非阻擋）
    asyncio.create_task(
        _write_budget_ledger(
            project_id=project_id,
            agent_id=agent_id,
            run_id=UUID(run_id) if run_id else None,
            model=model,
            provider=provider,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            cost_usd=actual_cost,
        ),
        name="budget_ledger_write",
    )

    # 更新 Redis cache（非阻擋）
    asyncio.create_task(
        _update_budget_cache(project_id, actual_cost),
        name="budget_cache_update",
    )

    logger.info(
        "token_usage_recorded",
        project_id=project_id,
        model=model,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        cost_usd=float(actual_cost),
    )
    return actual_cost


async def _write_budget_ledger(
    *,
    project_id: str,
    agent_id: str | None,
    run_id,  # UUID | None
    model: str,
    provider: str,
    prompt_tokens: int,
    completion_tokens: int,
    cost_usd: Decimal,
) -> None:
    """INSERT budget_ledger（leWOOOgo: DB 寫入在 Service 層，非 Router）"""
    try:
        from sqlalchemy import text
        from src.db.base import get_db_context
        async with get_db_context(project_id) as db:
            await db.execute(
                text("""
                    INSERT INTO budget_ledger
                        (project_id, agent_id, run_id, model, provider,
                         prompt_tokens, completion_tokens, cost_usd)
                    VALUES
                        (:project_id, :agent_id, :run_id, :model, :provider,
                         :prompt_tokens, :completion_tokens, :cost_usd)
                """),
                {
                    "project_id": project_id,
                    "agent_id": agent_id,
                    "run_id": run_id,
                    "model": model,
                    "provider": provider,
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "cost_usd": cost_usd,
                },
            )
    except Exception as exc:
        logger.warning("budget_ledger_write_failed", project_id=project_id, error=str(exc))


async def _update_budget_cache(project_id: str, cost: Decimal) -> None:
    """用 Redis INCRBYFLOAT 更新 Tenant + Platform daily budget cache"""
    try:
        from src.core.redis_client import get_redis
        redis = get_redis()
        cost_f = float(cost)

        # Tenant daily budget
        tenant_key = f"{_TENANT_BUDGET_KEY_PREFIX}{project_id}"
        await redis.incrbyfloat(tenant_key, cost_f)
        await redis.expire(tenant_key, 86400)  # 24h TTL（每日重置）

        # Platform daily budget
        await redis.incrbyfloat(_PLATFORM_BUDGET_KEY, cost_f)
        await redis.expire(_PLATFORM_BUDGET_KEY, 86400)

    except Exception as exc:
        logger.warning("budget_cache_update_failed", project_id=project_id, error=str(exc))


# ─────────────────────────────────────────────────────────────────────────────
# Helper：從 DB 讀取 Tenant budget limit
# ─────────────────────────────────────────────────────────────────────────────

async def _get_tenant_budget_limit(project_id: str) -> Decimal | None:
    """從 awooop_projects.budget_limit_usd 讀取 Tenant 每日上限（允許 NULL = 無上限）"""
    try:
        from sqlalchemy import text
        from src.db.base import get_db_context
        async with get_db_context(project_id) as db:
            row = await db.execute(
                text("SELECT budget_limit_usd FROM awooop_projects WHERE project_id = :pid"),
                {"pid": project_id},
            )
            result = row.scalar_one_or_none()
            return Decimal(str(result)) if result is not None else None
    except Exception as exc:
        logger.warning("get_tenant_budget_limit_failed", project_id=project_id, error=str(exc))
        return None


# ─────────────────────────────────────────────────────────────────────────────
# Emergency Kill Switch 管理（Admin 工具）
# ─────────────────────────────────────────────────────────────────────────────

async def activate_emergency_kill(reason: str = "") -> None:
    """啟動緊急停機 — SET platform:budget:emergency_kill"""
    from src.core.redis_client import get_redis
    redis = get_redis()
    await redis.set(_EMERGENCY_KILL_KEY, reason or "activated", ex=86400 * 7)
    logger.warning("budget_emergency_kill_activated", reason=reason)


async def deactivate_emergency_kill() -> None:
    """解除緊急停機"""
    from src.core.redis_client import get_redis
    redis = get_redis()
    await redis.delete(_EMERGENCY_KILL_KEY)
    logger.info("budget_emergency_kill_deactivated")


async def is_emergency_kill_active() -> bool:
    """查詢緊急停機狀態"""
    try:
        from src.core.redis_client import get_redis
        redis = get_redis()
        return bool(await redis.exists(_EMERGENCY_KILL_KEY))
    except Exception:
        return False