"""
title: LiteLLM Dashboard
author: Kagamine Rin
author_url: https://rin.captcouch.xyz
description: Check LiteLLM proxy health + auth status and get cost breakdowns with configurable timerange. Two functions, that's it.
required_open_webui_version: 0.4.0
requirements: requests
version: 1.3.2
licence: MIT
"""

import requests
from datetime import datetime, timedelta
from zoneinfo import ZoneInfo
from pydantic import BaseModel, Field
from typing import Optional


class Tools:
    def __init__(self):
        self.valves = self.Valves()

    # TASK-004: Hardcore FinOps Audit Registry (Standard Tier Pricing)
    MODEL_PRICING = {
        "gemini/gemini-3.5-flash": {
            "input_uncached": 1.50 / 1_000_000,
            "input_cached": 0.15 / 1_000_000,
            "output": 9.00 / 1_000_000,
        },
        "gemini/gemini-3-flash-preview": {
            "input_uncached": 0.50 / 1_000_000,
            "input_cached": 0.05 / 1_000_000,
            "output": 3.00 / 1_000_000,
        },
        "gemini/gemini-2.5-flash": {
            "input_uncached": 0.30 / 1_000_000,
            "input_cached": 0.03 / 1_000_000,
            "output": 2.50 / 1_000_000,
        },
        "gemini/gemini-3.1-flash-lite": {
            "input_uncached": 0.25 / 1_000_000,
            "input_cached": 0.025 / 1_000_000,
            "output": 1.50 / 1_000_000,
        },
        "gemini/gemini-3.1-pro-preview": {
            "input_uncached": 2.00 / 1_000_000,
            "input_cached": 0.20 / 1_000_000,
            "output": 12.00 / 1_000_000,
        }
    }

    def _audit_cost(self, model: str, prompt: int, cached: int, completion: int, reported_cost: float) -> dict:
        """Audit reported cost against our pricing registry, factoring in context caching discounts."""
        pricing = self.MODEL_PRICING.get(model)
        if not pricing:
            return {"audited": reported_cost, "delta": 0.0, "status": "unsupported", "display": ""}
        
        # Subtract cached tokens from prompt before applying standard rate
        uncached = max(0, prompt - cached)
        calculated = (cached * pricing["input_cached"]) + (uncached * pricing["input_uncached"]) + (completion * pricing["output"])
        delta = reported_cost - calculated
        
        # Flag discrepancies greater than half a cent
        if abs(delta) > 0.005:
            status = "mismatch"
            display = f" [calc: ${calculated:.4f} ⚠️]"
        else:
            status = "ok"
            display = f" [calc: ${calculated:.4f} ✓]"
            
        return {"audited": calculated, "delta": delta, "status": status, "display": display}

    class Valves(BaseModel):
        """LiteLLM proxy connection settings."""
        litellm_url: str = Field(
            default="http://localhost:4000",
            description="LiteLLM proxy base URL (e.g. http://localhost:4000)",
        )
        litellm_master_key: str = Field(
            default="",
            description="LiteLLM master API key for authentication",
            json_schema_extra={"input": {"type": "password"}},
        )
        default_hide_zero_spend: bool = Field(
            default=True,
            description="Default behavior for zero-spend model filtering. When True, models with sub-$0.001 spend are hidden by default and summarized. The per-call hide_zero_spend parameter can override this.",
        )
        timezone: str = Field(
            default="UTC",
            description="IANA Timezone name (e.g. America/New_York, Europe/London, Asia/Tokyo). Used for natural language date resolution (today, yesterday).",
        )

    def _headers(self) -> dict:
        """Build auth headers from valves."""
        if not self.valves.litellm_master_key:
            return {}
        return {"Authorization": f"Bearer {self.valves.litellm_master_key}"}

    def _base_url(self) -> str:
        """Get base URL, strip trailing slash."""
        return self.valves.litellm_url.rstrip("/")

    def _resolve_date(self, date_str: str) -> str:
        """Resolve natural language date shortcuts to YYYY-MM-DD.

        Supports: today, yesterday, tomorrow, N days ago,
        monday-sunday (this week), last monday-last sunday,
        this week, last week, this month, last month.
        Falls back to passthrough if already YYYY-MM-DD or unrecognized.
        """
        date_str = date_str.strip().lower()
        # Resolve "today" in the configured timezone
        try:
            today = datetime.now(ZoneInfo(self.valves.timezone)).date()
        except Exception:
            # Fallback to UTC if timezone is invalid
            today = datetime.utcnow().date()

        # Direct shortcuts
        if date_str == "today":
            return today.isoformat()
        if date_str == "yesterday":
            return (today - timedelta(days=1)).isoformat()
        if date_str == "tomorrow":
            return (today + timedelta(days=1)).isoformat()

        # "N days ago"
        if "days ago" in date_str:
            try:
                n = int(date_str.split()[0])
                return (today - timedelta(days=n)).isoformat()
            except (ValueError, IndexError):
                pass

        # Day names — most recent occurrence (including today if it matches)
        days = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]
        if date_str in days:
            target_day = days.index(date_str)
            days_back = (today.weekday() - target_day) % 7
            return (today - timedelta(days=days_back)).isoformat()

        # "last monday", "last tuesday", etc.
        if date_str.startswith("last "):
            day_part = date_str[5:].strip()
            if day_part in days:
                target_day = days.index(day_part)
                days_back = (today.weekday() - target_day) % 7 + 7
                return (today - timedelta(days=days_back)).isoformat()

        # "this week" → Monday of this week
        if date_str == "this week":
            return (today - timedelta(days=today.weekday())).isoformat()

        # "last week" → Monday of last week
        if date_str == "last week":
            return (today - timedelta(days=today.weekday() + 7)).isoformat()

        # "this month" → 1st of this month
        if date_str == "this month":
            return today.replace(day=1).isoformat()

        # "last month" → 1st of last month
        if date_str == "last month":
            if today.month == 1:
                return today.replace(year=today.year - 1, month=12, day=1).isoformat()
            return today.replace(month=today.month - 1, day=1).isoformat()

        # Already YYYY-MM-DD — passthrough
        try:
            datetime.strptime(date_str, "%Y-%m-%d")
            return date_str
        except ValueError:
            pass

        # Unknown — return as-is (API will reject with clear error)
        return date_str

    def litellm_health(self) -> str:
        """
        Check LiteLLM proxy health and authentication status.
        Returns which models are healthy/unhealthy, their rate limits, and any auth errors.
        Use this to verify your LiteLLM proxy is up and your master key works.
        """
        url = f"{self._base_url()}/health"
        headers = self._headers()

        try:
            resp = requests.get(url, headers=headers, timeout=15)
        except requests.exceptions.ConnectionError:
            return f"❌ Connection failed — LiteLLM proxy not reachable at {self._base_url()}"
        except requests.exceptions.Timeout:
            return f"❌ Connection timed out after 15s — proxy at {self._base_url()} may be overloaded"

        if resp.status_code == 401:
            return "❌ Authentication failed — check your litellm_master_key valve. The proxy rejected the key."
        if resp.status_code == 403:
            return "❌ Forbidden — your key doesn't have permission to access /health"
        if resp.status_code != 200:
            return f"❌ Unexpected status {resp.status_code}: {resp.text[:500]}"

        data = resp.json()
        healthy = data.get("healthy_endpoints", [])
        unhealthy = data.get("unhealthy_endpoints", [])
        healthy_count = data.get("healthy_count", len(healthy))
        unhealthy_count = data.get("unhealthy_count", len(unhealthy))

        # Strip litellm_metadata — it's ~2KB per model of auth context we don't need
        # Also strip other internal fields that bloat the response
        STRIP_KEYS = {
            "litellm_metadata", "safety_settings", "cache_control_injection_points",
            "use_in_pass_through", "use_litellm_proxy", "merge_reasoning_content_in_choices",
        }

        lines = []
        lines.append(f"LiteLLM Health Check — {self._base_url()}")
        lines.append("=" * 50)

        if not self.valves.litellm_master_key:
            lines.append("⚠️  No master key configured — showing public health only")
        else:
            lines.append("✅ Authenticated with master key")

        lines.append(f"\nHealthy: {healthy_count}  |  Unhealthy: {unhealthy_count}")

        if healthy:
            lines.append(f"\n--- Healthy Models ({healthy_count}) ---")
            for ep in healthy:
                model = ep.get("model", "unknown")
                # Strip gemini/ prefix for readability
                display = model.replace("gemini/", "") if model.startswith("gemini/") else model
                rpm = ep.get("rpm", "?")
                timeout = ep.get("timeout", "?")
                extras = []
                if ep.get("reasoning_effort") and ep["reasoning_effort"] != "none":
                    extras.append(f"reasoning: {ep['reasoning_effort']}")
                if ep.get("service_tier"):
                    extras.append(f"tier: {ep['service_tier']}")
                extra_str = f"  ({', '.join(extras)})" if extras else ""
                lines.append(f"  ✅ {display}  — rpm: {rpm}, timeout: {timeout}s{extra_str}")

        if unhealthy:
            lines.append(f"\n--- Unhealthy Models ({unhealthy_count}) ---")
            for ep in unhealthy:
                model = ep.get("model", "unknown")
                display = model.replace("gemini/", "") if model.startswith("gemini/") else model
                lines.append(f"  ❌ {display}")

        if not healthy and not unhealthy:
            lines.append("\n(no endpoint data returned)")

        return "\n".join(lines)

    def litellm_cost(
        self,
        start_date: str,
        end_date: str,
        hide_zero_spend: bool = False,
    ) -> str:
        """
        Get LiteLLM cost breakdown for a date range.
        :param start_date: Start date — YYYY-MM-DD or natural language shortcut (e.g. "2026-06-01", "today", "yesterday", "monday", "last monday", "two days ago", "this week", "last month")
        :param end_date: End date — YYYY-MM-DD or natural language shortcut
        :param hide_zero_spend: When True, force-hide models with sub-$0.001 spend. When False (default), use the valve-configured default behavior (default_hide_zero_spend). This two-level design lets the human set a default via valves while allowing the SLM to explicitly override on a per-call basis.
        Returns formatted cost report with per-model and per-day breakdowns, token counts, cache hit rate, and request stats.
        """
        if not self.valves.litellm_master_key:
            return "❌ Master key required for cost data. Set the litellm_master_key valve."

        # Resolve effective hide behavior: explicit param overrides valve default
        # hide_zero_spend=True  → ALWAYS hide (explicit override)
        # hide_zero_spend=False → use valve default (default_hide_zero_spend)
        effective_hide = hide_zero_spend or self.valves.default_hide_zero_spend

        resolved_start = self._resolve_date(start_date)
        resolved_end = self._resolve_date(end_date)

        url = f"{self._base_url()}/user/daily/activity"
        headers = self._headers()
        params = {"start_date": resolved_start, "end_date": resolved_end}

        try:
            resp = requests.get(url, headers=headers, params=params, timeout=30)
        except requests.exceptions.ConnectionError:
            return f"❌ Connection failed — LiteLLM proxy not reachable at {self._base_url()}"
        except requests.exceptions.Timeout:
            return f"❌ Connection timed out after 30s — large date ranges may be slow"

        if resp.status_code == 401:
            return "❌ Authentication failed — check your litellm_master_key valve."
        if resp.status_code == 403:
            return "❌ Forbidden — your key doesn't have permission to access cost data"
        if resp.status_code != 200:
            return f"❌ Unexpected status {resp.status_code}: {resp.text[:500]}"

        data = resp.json()
        results = data.get("results", [])
        metadata = data.get("metadata", {})

        if not results:
            return f"No cost data found for {resolved_start} to {resolved_end}."

        # Aggregate totals from metadata
        total_spend = metadata.get("total_spend", 0)
        total_prompt = metadata.get("total_prompt_tokens", 0)
        total_completion = metadata.get("total_completion_tokens", 0)
        total_tokens = metadata.get("total_tokens", 0)
        total_successful = metadata.get("total_successful_requests", 0)
        total_failed = metadata.get("total_failed_requests", 0)
        total_api = metadata.get("total_api_requests", 0)
        cache_read = metadata.get("total_cache_read_input_tokens", 0)
        cache_created = metadata.get("total_cache_creation_input_tokens", 0)

        # Cache hit rate
        total_input_tokens = total_prompt  # cache_read is a subset of prompt
        if total_input_tokens > 0 and cache_read > 0:
            cache_pct = (cache_read / total_input_tokens) * 100
            cache_str = f"{cache_pct:.1f}% ({cache_read:,} / {total_input_tokens:,} input tokens)"
        else:
            cache_str = "N/A"

        # Build per-model breakdown across all dates
        model_totals = {}
        for day in results:
            breakdown = day.get("breakdown", {})
            models = breakdown.get("models", {})
            for model_name, model_data in models.items():
                m_metrics = model_data.get("metrics", {})
                if model_name not in model_totals:
                    model_totals[model_name] = {
                        "spend": 0, "requests": 0, "prompt": 0,
                        "completion": 0, "total": 0, "failed": 0,
                        "cache_read": 0, "cache_created": 0
                    }
                model_totals[model_name]["spend"] += m_metrics.get("spend", 0)
                model_totals[model_name]["requests"] += m_metrics.get("successful_requests", 0)
                model_totals[model_name]["failed"] += m_metrics.get("failed_requests", 0)
                model_totals[model_name]["prompt"] += m_metrics.get("prompt_tokens", 0)
                model_totals[model_name]["completion"] += m_metrics.get("completion_tokens", 0)
                model_totals[model_name]["total"] += m_metrics.get("total_tokens", 0)
                model_totals[model_name]["cache_read"] += m_metrics.get("cache_read_input_tokens", 0)
                model_totals[model_name]["cache_created"] += m_metrics.get("cache_creation_input_tokens", 0)

        # Build per-date breakdown
        date_lines = []
        for day in results:
            date = day.get("date", "?")
            m = day.get("metrics", {})
            date_lines.append(
                f"  {date}: ${m.get('spend', 0):.2f} | "
                f"{m.get('successful_requests', 0)} ok, {m.get('failed_requests', 0)} fail | "
                f"{m.get('total_tokens', 0):,} tokens"
            )

        # Build model lines sorted by spend
        model_lines = []
        hidden_count = 0
        for model_name in sorted(model_totals, key=lambda k: model_totals[k]["spend"], reverse=True):
            mt = model_totals[model_name]
            if effective_hide and mt["spend"] < 0.001:
                hidden_count += 1
                continue
            # Strip common provider prefixes for readability
            display = model_name
            for prefix in ("gemini/", "openai/", "anthropic/", "mistral/", "cohere/"):
                if display.startswith(prefix):
                    display = display[len(prefix):]
                    break
            # Truncate long model names to fit table, appending ellipsis
            max_name = 24
            if len(display) > max_name:
                display = display[:max_name - 1] + "…"

            # Build token string with split and cache indicator
            p_tokens = mt["prompt"]
            c_tokens = mt["completion"]
            
            # Compact token numbers
            p_str = f"{p_tokens / 1000:.1f}k" if p_tokens >= 1000 else str(p_tokens)
            c_str = f"{c_tokens / 1000:.1f}k" if c_tokens >= 1000 else str(c_tokens)
            token_display = f"{p_str}→{c_str}"

            # Cache indicator
            cache_info = ""
            if mt["cache_read"] > 0:
                hit_rate = (mt["cache_read"] / p_tokens * 100) if p_tokens > 0 else 0
                cache_info = f" 💾{hit_rate:.0f}%"

            model_lines.append(
                f"  {display:25s} ${mt['spend']:7.2f} | "
                f"{mt['requests']:3d} ok, {mt['failed']:d} fail | "
                f"{token_display:>10s}{cache_info}"
            )

        # Assemble report
        lines = []
        lines.append(f"LiteLLM Cost Report: {resolved_start} to {resolved_end}")
        lines.append("=" * 70)
        lines.append(f"Total Spend:     ${total_spend:.2f}")
        lines.append(f"Total Tokens:    {total_tokens:,} ({total_prompt:,} prompt + {total_completion:,} completion)")
        lines.append(f"Cache Hit Rate:  {cache_str}")
        lines.append(f"Requests:        {total_successful} successful, {total_failed} failed ({total_api} total)")
        if cache_created > 0:
            lines.append(f"Cache Created:   {cache_created:,} tokens")
        lines.append("")
        lines.append("By Model:")
        lines.append(f'  {"Model":25s} {"Spend":>7s} | {"Requests":>14s} | {"Tokens":>10s}')
        lines.extend(model_lines)
        if hidden_count > 0:
            lines.append(f"  ... ({hidden_count} model(s) with sub-$0.001 spend hidden)")
        lines.append("")
        lines.append("By Date:")
        lines.extend(date_lines)

        return "\n".join(lines)

    def litellm_history(
        self,
        limit: int = 20,
        show_cost: bool = True,
        show_failed: bool = True,
    ) -> str:
        """
        Get recent LiteLLM proxy request history.
        :param limit: Number of recent requests to show (1-50, default 20)
        :param show_cost: Include per-request cost (default True)
        :param show_failed: Include failed requests in output (default True). Set False to hide auth failures, routing errors, etc.
        Returns a formatted table of recent LLM requests with model, status, tokens, duration, and optional cost.
        """
        if not self.valves.litellm_master_key:
            return "❌ Master key required for request history. Set the litellm_master_key valve."

        limit = max(1, min(50, limit))

        url = f"{self._base_url()}/spend/logs"
        headers = self._headers()

        try:
            resp = requests.get(url, headers=headers, timeout=30)
        except requests.exceptions.ConnectionError:
            return f"❌ Connection failed — LiteLLM proxy not reachable at {self._base_url()}"
        except requests.exceptions.Timeout:
            return f"❌ Connection timed out after 30s — request history may be large"

        if resp.status_code == 401:
            return "❌ Authentication failed — check your litellm_master_key valve."
        if resp.status_code == 403:
            return "❌ Forbidden — your key doesn't have permission to access request logs"
        if resp.status_code != 200:
            return f"❌ Unexpected status {resp.status_code}: {resp.text[:500]}"

        data = resp.json()
        if not isinstance(data, list) or len(data) == 0:
            return "No request history found."

        # Sort by startTime descending and take the most recent `limit` entries
        sorted_logs = sorted(data, key=lambda e: e.get("startTime", ""), reverse=True)
        
        # Filter out failed requests if show_failed is False
        if not show_failed:
            sorted_logs = [
                e for e in sorted_logs
                if not (
                    (e.get("metadata", {}) or {}).get("exception", "")
                    or (e.get("metadata", {}) or {}).get("status", "") == "failure"
                )
            ]
        
        recent = sorted_logs[:limit]
        total_count = len(data)
        filtered_note = " (failed requests hidden)" if not show_failed else ""

        lines = []
        lines.append(f"LiteLLM Request History — {self._base_url()}")
        lines.append(f"Showing {len(recent)} of {total_count} total requests{filtered_note} (most recent first)")
        lines.append("=" * 90)

        for i, entry in enumerate(recent, 1):
            # Timestamp
            start_time = entry.get("startTime", "?")[:19]  # YYYY-MM-DDTHH:MM:SS
            start_short = start_time[5:] if len(start_time) > 5 else start_time  # MM-DDTHH:MM:SS

            # Model — strip provider prefix for readability
            model = entry.get("model", "unknown")
            for prefix in ("gemini/", "openai/", "anthropic/", "mistral/", "cohere/"):
                if model.startswith(prefix):
                    model = model[len(prefix):]
                    break

            # Call type
            call_type = entry.get("call_type", "?")
            # Shorten common call types
            call_type = call_type.replace("acompletion", "chat").replace("aembedding", "embed").replace("aimage_generation", "image").replace("_", " ")

            # Status — check for errors
            metadata = entry.get("metadata", {}) or {}
            exception = metadata.get("exception", "")
            status_flag = metadata.get("status", "")
            is_failed = bool(exception) or status_flag == "failure"
            status = "❌ FAIL" if is_failed else "✅ OK"

            # Duration
            duration_ms = entry.get("request_duration_ms", 0) or 0
            if duration_ms >= 1000:
                duration_str = f"{duration_ms / 1000:.1f}s"
            else:
                duration_str = f"{duration_ms}ms"

            # Tokens
            prompt_tokens = entry.get("prompt_tokens", 0) or 0
            completion_tokens = entry.get("completion_tokens", 0) or 0
            total_tokens = entry.get("total_tokens", 0) or 0

            # Cache extraction (handles Anthropic/OpenAI top-level and Gemini/z.ai nested metadata structures)
            cache_read = entry.get("cache_read_input_tokens", 0) or 0
            if cache_read == 0:
                metadata = entry.get("metadata", {}) or {}
                usage_obj = metadata.get("usage_object", {}) or {}
                if isinstance(usage_obj, dict):
                    cache_read = usage_obj.get("cache_read_input_tokens", 0) or 0
                if cache_read == 0:
                    prompt_details = entry.get("prompt_tokens_details", {}) or {}
                    if isinstance(prompt_details, dict):
                        cache_read = prompt_details.get("cached_tokens", 0) or 0

            # Spend
            spend = entry.get("spend", 0) or 0

            # Build line
            base = f"  {i:2d}. {start_short} | {model[:28]:28s} | {call_type[:10]:10s} | {status} | {duration_str:>7s}"

            # Token info
            if total_tokens > 0:
                # Compact token display
                if prompt_tokens >= 1000:
                    p_str = f"{prompt_tokens / 1000:.1f}k"
                else:
                    p_str = str(prompt_tokens)
                if completion_tokens >= 1000:
                    c_str = f"{completion_tokens / 1000:.1f}k"
                else:
                    c_str = str(completion_tokens)
                token_str = f"{p_str}→{c_str}"

                # Cache indicator
                if cache_read > 0:
                    cache_pct = (cache_read / prompt_tokens * 100) if prompt_tokens > 0 else 0
                    token_str += f" 💾{cache_pct:.0f}%"

                base += f" | {token_str}"

            # Cost & Audit
            if show_cost and spend > 0:
                audit = self._audit_cost(entry.get("model", ""), prompt_tokens, cache_read, completion_tokens, spend)
                if audit["status"] == "mismatch":
                    base += f" | ${spend:.4f} (calc: ${audit['audited']:.4f} ⚠️)"
                elif audit["status"] == "ok":
                    base += f" | ${spend:.4f} ✓"
                else:
                    base += f" | ${spend:.4f}"

            # Error info
            if is_failed:
                # Priority: exception string > error_information.error_message > generic status
                if exception:
                    err_short = exception[:80] + "…" if len(exception) > 80 else exception
                else:
                    err_info = metadata.get("error_information", {}) or {}
                    err_msg = err_info.get("error_message", "")
                    if err_msg:
                        err_short = err_msg[:80] + "…" if len(err_msg) > 80 else err_msg
                    else:
                        err_short = "status: failure"
                base += f"\n      └─ {err_short}"

            lines.append(base)

        # Summary
        if len(recent) > 0:
            recent_spend = sum(e.get("spend", 0) or 0 for e in recent)
            recent_prompt = sum(e.get("prompt_tokens", 0) or 0 for e in recent)
            recent_completion = sum(e.get("completion_tokens", 0) or 0 for e in recent)
            recent_cache = sum(e.get("cache_read_input_tokens", 0) or 0 for e in recent)
            recent_failed = sum(1 for e in recent if ((e.get("metadata", {}) or {}).get("exception")) or ((e.get("metadata", {}) or {}).get("status", "")) == "failure")

            lines.append("")
            lines.append(f"  Summary ({len(recent)} shown): ${recent_spend:.4f} spent | "
                        f"{recent_prompt:,} prompt + {recent_completion:,} completion tokens | "
                        f"{recent_failed} failed | "
                        f"cache hit {recent_cache / recent_prompt * 100:.0f}%" if recent_prompt > 0 and recent_cache > 0
                        else f"  Summary ({len(recent)} shown): ${recent_spend:.4f} spent | "
                             f"{recent_prompt:,} prompt + {recent_completion:,} completion tokens | "
                             f"{recent_failed} failed")

        return "\n".join(lines)