"""
BuildShield MCP Tool Handlers

All offline tools (check_cost, recommend_model, get_pricing, optimize_prompt)
work with zero network calls using the embedded pricing table.

Online tools (check_budget, log_usage) require a BuildShield API key (bsk_...)
which users can obtain at https://thebuildshield.com.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Try to import the live pricing engine from the repo backend
# ---------------------------------------------------------------------------
_BACKEND_PATH = Path(__file__).parent.parent.parent.parent / "backend"
_PRICING_AVAILABLE = False

if _BACKEND_PATH.exists():
    sys.path.insert(0, str(_BACKEND_PATH))
    try:
        from app.modules.llm_cost.pricing import (  # type: ignore[import]
            get_pricing_table,
            resolve_model,
            calculate_cost,
        )
        _PRICING_AVAILABLE = True
    except Exception:
        pass

# ---------------------------------------------------------------------------
# Embedded fallback — used when running standalone outside the repo
# ---------------------------------------------------------------------------
_FALLBACK_MODELS: list[dict] = [
    {"provider": "openai",    "model": "gpt-4.1",           "display_name": "GPT-4.1",           "input_per_1m": 2.00,  "output_per_1m": 8.00,  "context": 1_000_000, "tier": "balanced"},
    {"provider": "openai",    "model": "gpt-4.1-mini",       "display_name": "GPT-4.1 Mini",       "input_per_1m": 0.40,  "output_per_1m": 1.60,  "context": 1_000_000, "tier": "cheap"},
    {"provider": "openai",    "model": "gpt-4o",             "display_name": "GPT-4o",             "input_per_1m": 2.50,  "output_per_1m": 10.00, "context": 128_000,   "tier": "balanced"},
    {"provider": "openai",    "model": "gpt-4o-mini",        "display_name": "GPT-4o Mini",        "input_per_1m": 0.15,  "output_per_1m": 0.60,  "context": 128_000,   "tier": "cheap"},
    {"provider": "openai",    "model": "o1",                 "display_name": "o1",                 "input_per_1m": 15.00, "output_per_1m": 60.00, "context": 200_000,   "tier": "premium"},
    {"provider": "openai",    "model": "o3-mini",            "display_name": "o3-mini",            "input_per_1m": 1.10,  "output_per_1m": 4.40,  "context": 200_000,   "tier": "balanced"},
    {"provider": "anthropic", "model": "claude-opus-4.6",    "display_name": "Claude Opus 4.6",    "input_per_1m": 15.00, "output_per_1m": 75.00, "context": 200_000,   "tier": "premium"},
    {"provider": "anthropic", "model": "claude-sonnet-4.6",  "display_name": "Claude Sonnet 4.6",  "input_per_1m": 3.00,  "output_per_1m": 15.00, "context": 200_000,   "tier": "balanced"},
    {"provider": "anthropic", "model": "claude-sonnet-4.5",  "display_name": "Claude Sonnet 4.5",  "input_per_1m": 3.00,  "output_per_1m": 15.00, "context": 200_000,   "tier": "balanced"},
    {"provider": "anthropic", "model": "claude-haiku-4.5",   "display_name": "Claude Haiku 4.5",   "input_per_1m": 1.00,  "output_per_1m": 5.00,  "context": 200_000,   "tier": "cheap"},
    {"provider": "google",    "model": "gemini-2.5-pro",     "display_name": "Gemini 2.5 Pro",     "input_per_1m": 1.25,  "output_per_1m": 10.00, "context": 1_000_000, "tier": "balanced"},
    {"provider": "google",    "model": "gemini-2.5-flash",   "display_name": "Gemini 2.5 Flash",   "input_per_1m": 0.30,  "output_per_1m": 2.50,  "context": 1_000_000, "tier": "cheap"},
    {"provider": "google",    "model": "gemini-2.0-flash",   "display_name": "Gemini 2.0 Flash",   "input_per_1m": 0.10,  "output_per_1m": 0.40,  "context": 1_000_000, "tier": "cheap"},
    {"provider": "deepseek",  "model": "deepseek-chat",      "display_name": "DeepSeek Chat",      "input_per_1m": 0.28,  "output_per_1m": 0.42,  "context": 128_000,   "tier": "cheap"},
    {"provider": "deepseek",  "model": "deepseek-r1",        "display_name": "DeepSeek R1",        "input_per_1m": 0.55,  "output_per_1m": 2.19,  "context": 128_000,   "tier": "cheap"},
]

_API_BASE = "https://thebuildshieldcom-production.up.railway.app"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _fmt_tokens(n: int) -> str:
    if n >= 1_000_000:
        return f"{n/1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n/1_000:.0f}K"
    return str(n)


def _get_fallback_table() -> list[dict]:
    return _FALLBACK_MODELS


def _fallback_resolve(model_name: str) -> dict | None:
    name_lower = model_name.lower().strip()
    for m in _FALLBACK_MODELS:
        if m["model"] == name_lower or m["display_name"].lower() == name_lower:
            return m
    # Partial match
    for m in _FALLBACK_MODELS:
        if name_lower in m["model"] or m["model"] in name_lower:
            return m
    return None


def _fallback_calculate(model_name: str, input_tokens: int, output_tokens: int) -> float | None:
    m = _fallback_resolve(model_name)
    if not m:
        return None
    return (input_tokens / 1_000_000) * m["input_per_1m"] + (output_tokens / 1_000_000) * m["output_per_1m"]


# ---------------------------------------------------------------------------
# Tool 1: buildshield_check_cost
# ---------------------------------------------------------------------------

def handle_check_cost(model: str, input_tokens: int, output_tokens: int = 0) -> str:
    """Calculate exact cost for a specific model and token count. Works offline."""
    if _PRICING_AVAILABLE:
        mp = resolve_model(model)
        if mp:
            input_cost = (input_tokens / 1_000_000) * mp.input_per_1m
            output_cost = (output_tokens / 1_000_000) * mp.output_per_1m
            total = input_cost + output_cost
            return (
                f"Model:  {mp.display_name} ({mp.provider})\n"
                f"Tier:   {mp.family}\n"
                f"Input:  {_fmt_tokens(input_tokens)} tokens × ${mp.input_per_1m:.2f}/1M = ${input_cost:.6f}\n"
                f"Output: {_fmt_tokens(output_tokens)} tokens × ${mp.output_per_1m:.2f}/1M = ${output_cost:.6f}\n"
                f"Total:  ${total:.6f}  (${total*100:.4f}¢)\n"
                f"\nContext window: {_fmt_tokens(mp.max_context_tokens)}"
            )
    else:
        m = _fallback_resolve(model)
        if m:
            input_cost = (input_tokens / 1_000_000) * m["input_per_1m"]
            output_cost = (output_tokens / 1_000_000) * m["output_per_1m"]
            total = input_cost + output_cost
            return (
                f"Model:  {m['display_name']} ({m['provider']})\n"
                f"Tier:   {m['tier']}\n"
                f"Input:  {_fmt_tokens(input_tokens)} tokens × ${m['input_per_1m']:.2f}/1M = ${input_cost:.6f}\n"
                f"Output: {_fmt_tokens(output_tokens)} tokens × ${m['output_per_1m']:.2f}/1M = ${output_cost:.6f}\n"
                f"Total:  ${total:.6f}  (${total*100:.4f}¢)\n"
                f"\nContext window: {_fmt_tokens(m['context'])}"
            )

    return (
        f"Model '{model}' not found in BuildShield pricing registry.\n"
        f"Try: gpt-4.1, gpt-4o, claude-sonnet-4.6, gemini-2.5-flash, deepseek-chat\n"
        f"Or use buildshield_get_pricing to browse all available models."
    )


# ---------------------------------------------------------------------------
# Tool 2: buildshield_recommend_model
# ---------------------------------------------------------------------------

def handle_recommend_model(
    task: str,
    max_cost_per_1m_input: float | None = None,
    min_context_tokens: int | None = None,
    provider: str | None = None,
) -> str:
    """Recommend the best models for a task based on cost and capability filters. Works offline."""
    if _PRICING_AVAILABLE:
        table = get_pricing_table()
        candidates = list(table.values())

        def to_dict(mp: Any) -> dict:
            return {
                "model": mp.model,
                "display_name": mp.display_name,
                "provider": mp.provider,
                "input_per_1m": mp.input_per_1m,
                "output_per_1m": mp.output_per_1m,
                "context": mp.max_context_tokens,
                "tier": mp.family,
                "status": mp.status,
            }

        models = [to_dict(m) for m in candidates if m.status != "legacy"]
    else:
        models = [m for m in _FALLBACK_MODELS]

    # Apply filters
    if max_cost_per_1m_input is not None:
        models = [m for m in models if m["input_per_1m"] <= max_cost_per_1m_input]
    if min_context_tokens is not None:
        models = [m for m in models if m["context"] >= min_context_tokens]
    if provider:
        models = [m for m in models if m["provider"].lower() == provider.lower()]

    # Sort by input cost
    models.sort(key=lambda m: m["input_per_1m"])

    if not models:
        return (
            "No models match your filters.\n"
            "Try relaxing: max_cost_per_1m_input or min_context_tokens.\n"
            "Use buildshield_get_pricing to see all available models."
        )

    # Build recommendation
    lines = [f"Top model recommendations for: {task}\n"]
    lines.append(f"{'Model':<30} {'Provider':<12} {'Input/1M':>10} {'Output/1M':>11} {'Context':>10} {'Tier':<10}")
    lines.append("─" * 85)

    for m in models[:6]:
        ctx = _fmt_tokens(m["context"])
        lines.append(
            f"{m['display_name']:<30} {m['provider']:<12} ${m['input_per_1m']:>8.2f} "
            f"${m['output_per_1m']:>9.2f} {ctx:>10}   {m['tier']:<10}"
        )

    # Add a recommendation note based on task keywords
    task_lower = task.lower()
    note = ""
    if any(w in task_lower for w in ["summar", "classif", "extract", "simple", "short"]):
        note = "\n💡 For simple extraction/classification tasks, consider the cheapest option — quality difference is minimal."
    elif any(w in task_lower for w in ["code", "debug", "program", "function"]):
        note = "\n💡 For coding tasks, balanced-tier models (Sonnet, GPT-4.1) offer the best quality/cost ratio."
    elif any(w in task_lower for w in ["reason", "complex", "analyz", "research"]):
        note = "\n💡 For complex reasoning, premium models (Opus, o1) justify their cost with significantly better output."
    elif any(w in task_lower for w in ["long", "document", "book", "context"]):
        note = "\n💡 For long-context tasks, Gemini 2.5 Pro (1M context) or Claude (200K) are the best choices."

    return "\n".join(lines) + note


# ---------------------------------------------------------------------------
# Tool 3: buildshield_get_pricing
# ---------------------------------------------------------------------------

def handle_get_pricing(model: str | None = None, provider: str | None = None) -> str:
    """Get pricing for one model or browse all models. Works offline."""
    if model:
        if _PRICING_AVAILABLE:
            mp = resolve_model(model)
            if mp:
                return (
                    f"Model:         {mp.display_name}\n"
                    f"Provider:      {mp.provider}\n"
                    f"Model ID:      {mp.model}\n"
                    f"Input price:   ${mp.input_per_1m:.2f} per 1M tokens\n"
                    f"Output price:  ${mp.output_per_1m:.2f} per 1M tokens\n"
                    f"Context:       {_fmt_tokens(mp.max_context_tokens)}\n"
                    f"Tier:          {mp.family}\n"
                    f"Status:        {mp.status}\n"
                    + (f"Alternatives:  {', '.join(mp.alternatives)}" if mp.alternatives else "")
                )
        else:
            m = _fallback_resolve(model)
            if m:
                return (
                    f"Model:         {m['display_name']}\n"
                    f"Provider:      {m['provider']}\n"
                    f"Model ID:      {m['model']}\n"
                    f"Input price:   ${m['input_per_1m']:.2f} per 1M tokens\n"
                    f"Output price:  ${m['output_per_1m']:.2f} per 1M tokens\n"
                    f"Context:       {_fmt_tokens(m['context'])}\n"
                    f"Tier:          {m['tier']}"
                )
        return f"Model '{model}' not found. Use buildshield_get_pricing (no model arg) to see all models."

    # Return full table
    if _PRICING_AVAILABLE:
        table = get_pricing_table()
        models = sorted(table.values(), key=lambda m: (m.provider, m.input_per_1m))
        rows = []
        for mp in models:
            if provider and mp.provider.lower() != provider.lower():
                continue
            rows.append({
                "display_name": mp.display_name,
                "provider": mp.provider,
                "input_per_1m": mp.input_per_1m,
                "output_per_1m": mp.output_per_1m,
                "context": mp.max_context_tokens,
                "tier": mp.family,
            })
    else:
        rows = sorted(
            [m for m in _FALLBACK_MODELS if not provider or m["provider"].lower() == provider.lower()],
            key=lambda m: (m["provider"], m["input_per_1m"])
        )
        rows = [dict(display_name=m["display_name"], provider=m["provider"],
                     input_per_1m=m["input_per_1m"], output_per_1m=m["output_per_1m"],
                     context=m["context"], tier=m["tier"]) for m in rows]

    if not rows:
        return f"No models found for provider '{provider}'."

    lines = [f"{'Model':<30} {'Provider':<12} {'Input/1M':>10} {'Output/1M':>11} {'Context':>10} {'Tier':<10}"]
    lines.append("─" * 85)
    for r in rows:
        ctx = _fmt_tokens(r["context"])
        lines.append(
            f"{r['display_name']:<30} {r['provider']:<12} ${r['input_per_1m']:>8.2f} "
            f"${r['output_per_1m']:>9.2f} {ctx:>10}   {r['tier']:<10}"
        )
    lines.append(f"\n{len(rows)} models listed.")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 4: buildshield_optimize_prompt
# ---------------------------------------------------------------------------

_FILLER_REPLACEMENTS = [
    (r"\bplease\s+", ""),
    (r"\bkindly\s+", ""),
    (r"\bI would like you to\s+", ""),
    (r"\bmake sure to\s+", "ensure "),
    (r"\bin order to\b", "to"),
    (r"\bas well as\b", "and"),
    (r"\bin the event that\b", "if"),
    (r"\bat this point in time\b", "now"),
    (r"\bdue to the fact that\b", "because"),
    (r"\bfor the purpose of\b", "for"),
    (r"\bit is important to note that\b", "Note:"),
    (r"\bplease note that\b", "Note:"),
    (r"\bit should be noted that\b", "Note:"),
    (r"\bfeel free to\b", ""),
    (r"\bdo not hesitate to\b", ""),
    (r"\bI would appreciate it if you could\b", ""),
    (r"\bwould you be able to\b", ""),
    (r"\bcould you please\b", ""),
]


def _word_overlap(a: str, b: str) -> float:
    wa, wb = set(a.lower().split()), set(b.lower().split())
    if not wa or not wb:
        return 0.0
    return len(wa & wb) / max(len(wa), len(wb))


def _dedup_paragraphs(text: str) -> tuple[str, int]:
    paras = [p.strip() for p in re.split(r"\n{2,}", text) if p.strip()]
    removed = 0
    result = []
    for i, para in enumerate(paras):
        is_dup = False
        for j in range(i + 1, len(paras)):
            if _word_overlap(para, paras[j]) > 0.60:
                is_dup = True
                removed += 1
                break
        if not is_dup:
            result.append(para)
    return "\n\n".join(result), removed


def handle_optimize_prompt(text: str, aggressiveness: str = "medium") -> str:
    """Rule-based prompt compression. No AI — works offline. Saves 20-40% on verbose prompts."""
    if not text or not text.strip():
        return "Empty prompt provided. Please pass the text you want to optimize."

    original_tokens = max(1, len(text) // 4)
    techniques: list[str] = []
    optimized = text

    # ---- Light: whitespace normalization ----
    before = optimized
    optimized = re.sub(r"[ \t]+", " ", optimized)          # collapse spaces/tabs
    optimized = re.sub(r"\n{3,}", "\n\n", optimized)        # max 2 blank lines
    optimized = "\n".join(line.rstrip() for line in optimized.splitlines())
    optimized = optimized.strip()
    if optimized != before:
        techniques.append("Whitespace normalization")

    if aggressiveness in ("medium", "heavy"):
        # ---- Medium: filler phrase removal ----
        before = optimized
        for pattern, replacement in _FILLER_REPLACEMENTS:
            optimized = re.sub(pattern, replacement, optimized, flags=re.IGNORECASE)
        # Clean up double spaces left by removals
        optimized = re.sub(r"  +", " ", optimized)
        optimized = re.sub(r"^ +", "", optimized, flags=re.MULTILINE)
        if optimized != before:
            techniques.append("Filler phrase removal")

    if aggressiveness == "heavy":
        # ---- Heavy: paragraph deduplication ----
        before = optimized
        optimized, removed_count = _dedup_paragraphs(optimized)
        if removed_count:
            techniques.append(f"Duplicate paragraph removal ({removed_count} removed)")

    optimized_tokens = max(1, len(optimized) // 4)
    savings_pct = round((1 - optimized_tokens / original_tokens) * 100, 1)

    if not techniques:
        return (
            f"Prompt analyzed — no optimizations needed.\n"
            f"Tokens: {original_tokens:,} | The prompt is already concise."
        )

    # Format output
    lines = [
        f"Optimization complete ({aggressiveness} mode)",
        f"",
        f"Original:  {original_tokens:,} tokens",
        f"Optimized: {optimized_tokens:,} tokens",
        f"Saved:     {original_tokens - optimized_tokens:,} tokens ({savings_pct}%)",
        f"",
        f"Techniques applied:",
    ]
    for t in techniques:
        lines.append(f"  • {t}")
    lines += [
        f"",
        f"--- OPTIMIZED PROMPT ---",
        optimized,
        f"--- END ---",
    ]
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Tool 5: buildshield_check_budget
# ---------------------------------------------------------------------------

async def handle_check_budget(api_key: str | None = None, period: str = "today") -> str:
    """Check your BuildShield budget. Requires a BuildShield API key (bsk_...)."""
    if not api_key:
        return (
            "This tool requires a BuildShield API key.\n\n"
            "Get your key at: https://thebuildshield.com/settings\n"
            "Then pass it as: api_key='bsk_...'\n\n"
            "Or configure it via the BUILDSHIELD_API_KEY environment variable."
        )

    if not api_key.startswith("bsk_"):
        return f"Invalid key format. BuildShield keys start with 'bsk_'. Got: '{api_key[:8]}...'"

    try:
        import httpx
        async with httpx.AsyncClient(timeout=10.0) as client:
            r = await client.get(
                f"{_API_BASE}/api/v1/agent/budget",
                params={"period": period},
                headers={"X-Api-Key": api_key},
            )
        if r.status_code == 401:
            return "Invalid or expired API key. Check your key at https://thebuildshield.com/settings"
        if r.status_code == 404:
            return "No budget configured for this key. Set a budget at https://thebuildshield.com/settings"
        r.raise_for_status()
        data = r.json()
        budget = data.get("budget", {})
        quota = data.get("quota", {})
        limit = budget.get("limit_usd") or 0
        used = budget.get("used_usd") or 0
        remaining = budget.get("remaining_usd")
        if remaining is None:
            remaining = limit - used
        resets = budget.get("reset_at") or "not set"
        calls_used = quota.get("calls_used_this_month", 0)
        calls_limit = quota.get("calls_limit")
        pct = round((used / limit * 100) if limit else 0, 1)
        bar_filled = int(pct / 5)
        bar = "█" * bar_filled + "░" * (20 - bar_filled)
        plan = data.get("plan", "unknown")
        quota_line = f"API calls: {calls_used}/{calls_limit if calls_limit else '∞'} this month"
        if budget.get("has_budget"):
            return (
                f"Budget Status ({period}) — Plan: {plan}\n"
                f"{'─'*30}\n"
                f"Used:      ${used:.4f} / ${limit:.2f}\n"
                f"Remaining: ${remaining:.4f}\n"
                f"Progress:  [{bar}] {pct}%\n"
                f"Resets at: {resets}\n"
                f"{quota_line}"
            )
        else:
            return (
                f"Budget Status — Plan: {plan}\n"
                f"{'─'*30}\n"
                f"No spending budget configured.\n"
                f"Set a budget at: https://thebuildshield.com/settings\n"
                f"{quota_line}"
            )
    except ImportError:
        return "httpx not installed. Run: pip install httpx"
    except Exception as e:
        return f"Could not reach BuildShield API: {e}\nCheck your connection or try again later."


# ---------------------------------------------------------------------------
# Tool 6: buildshield_log_usage
# ---------------------------------------------------------------------------

async def handle_log_usage(
    model: str,
    input_tokens: int,
    output_tokens: int,
    api_key: str | None = None,
    task_description: str | None = None,
) -> str:
    """Log AI usage to BuildShield. Always calculates cost locally; syncs to cloud if api_key provided."""
    # Calculate cost locally first (always works offline)
    if _PRICING_AVAILABLE:
        cost = calculate_cost(model, input_tokens, output_tokens)
        mp = resolve_model(model)
        model_display = mp.display_name if mp else model
    else:
        cost = _fallback_calculate(model, input_tokens, output_tokens)
        m = _fallback_resolve(model)
        model_display = m["display_name"] if m else model

    if cost is None:
        cost_str = "(model not found — cost unknown)"
        cost = 0.0
    else:
        cost_str = f"${cost:.6f}"

    local_summary = (
        f"Usage logged locally\n"
        f"Model:   {model_display}\n"
        f"Input:   {input_tokens:,} tokens\n"
        f"Output:  {output_tokens:,} tokens\n"
        f"Cost:    {cost_str}"
    )
    if task_description:
        local_summary += f"\nTask:    {task_description}"

    if not api_key:
        return local_summary + "\n\nNote: Pass api_key='bsk_...' to sync to your BuildShield dashboard."

    # Sync to API
    try:
        import httpx
        payload = {
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "estimated_cost_usd": cost,
            "source": "mcp_server",
        }
        if task_description:
            payload["task_description"] = task_description

        async with httpx.AsyncClient(timeout=10.0) as client:
            r = await client.post(
                f"{_API_BASE}/api/v1/agent/usage",
                json=payload,
                headers={"X-Api-Key": api_key},
            )
        if r.status_code in (200, 201):
            data = r.json()
            remaining = data.get("budget_remaining_usd")
            synced_note = "✓ Synced to BuildShield dashboard"
            if remaining is not None:
                synced_note += f" | Budget remaining: ${remaining:.4f}"
            return local_summary + f"\n\n{synced_note}"
        elif r.status_code == 429:
            return local_summary + "\n\n⚠ Budget limit reached — usage logged locally only."
        else:
            return local_summary + f"\n\nCloud sync failed (HTTP {r.status_code}) — usage logged locally."
    except ImportError:
        return local_summary + "\n\nhttpx not installed. Run: pip install httpx"
    except Exception as e:
        return local_summary + f"\n\nCloud sync failed: {e} — usage logged locally."
