"""BudgetMixin — budget tracking, pressure, spend recording."""
from __future__ import annotations

from dataclasses import asdict
from fnmatch import fnmatch
from typing import Any

from logos_adk.budgeting import BudgetAccount, BudgetStatus, budget_status, ensure_budget_account, record_budget_usage
from logos_adk.context import RunContext, get_run_context
from logos_adk.cost import DEFAULT_RATES
from logos_adk.errors import BudgetExhaustedError
from logos_adk.providers import Provider
from logos_adk.spend_ledger import InMemorySpendLedger, PostgreSQLSpendLedger, SQLiteSpendLedger, SpendLedger, SpendRecord
from logos_adk.types import Response, Usage

from logos_adk.agent._constants import _MAX_TOOL_ROUNDS


class BudgetMixin:
    """Mixin providing budget and spend tracking methods for Agent."""

    def _resolve_spend_ledger(self) -> SpendLedger | None:
        if self._spend_ledger is not None:
            return self._spend_ledger
        if not self._config.routing.enabled:
            return None
        backend = str(self._config.routing.ledger_backend or "sqlite").strip().lower()
        if backend in {"none", "disabled", "off"}:
            return None
        if backend in {"memory", "in_memory", "inmemory"}:
            self._spend_ledger = InMemorySpendLedger()
            return self._spend_ledger
        if backend in {"postgres", "postgresql"}:
            self._spend_ledger = PostgreSQLSpendLedger(dsn=self._config.routing.ledger_dsn)
            return self._spend_ledger
        self._spend_ledger = SQLiteSpendLedger(path=self._config.routing.ledger_path)
        return self._spend_ledger

    def _runtime_budget_status(self, run_context: RunContext | None) -> dict[str, BudgetStatus]:
        if run_context is None:
            return {}
        payload: dict[str, BudgetStatus] = {}
        for scope, account in (
            ("task", run_context.task_budget),
            ("team", run_context.team_budget),
        ):
            status = budget_status(account)
            if status is not None:
                payload[scope] = status
        return payload

    def _spend_budget_status(self, run_context: RunContext | None) -> dict[str, Any]:
        ledger = self._resolve_spend_ledger()
        if ledger is None or run_context is None:
            return {}
        try:
            return ledger.budget_status(
                agent_name=self.name,
                workspace_id=run_context.workspace_id,
                session_id=run_context.session_id,
                workspace_soft_limit_usd=self._config.routing.workspace_soft_limit_usd,
                session_soft_limit_usd=self._config.routing.session_soft_limit_usd,
                threshold_ratio=self._config.routing.threshold_ratio,
            )
        except Exception:
            return {}

    def _budget_status_snapshot(self, run_context: RunContext | None) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        for scope, status in self._runtime_budget_status(run_context).items():
            payload[scope] = status.to_dict()
        for scope, status in self._spend_budget_status(run_context).items():
            payload[scope] = asdict(status)
        return payload

    def _budget_pressure(self, run_context: RunContext | None) -> dict[str, Any]:
        runtime_status = self._runtime_budget_status(run_context)
        ledger_status = self._spend_budget_status(run_context)
        budget_ratio = max(
            [item.ratio for item in runtime_status.values()]
            + [float(item.ratio) for item in ledger_status.values()],
            default=0.0,
        )
        near_limit_scopes = [
            scope for scope, item in runtime_status.items() if item.near_limit
        ] + [
            scope for scope, item in ledger_status.items() if bool(item.near_limit)
        ]
        exhausted_scopes = [scope for scope, item in runtime_status.items() if item.exhausted]
        return {
            "budget_ratio": budget_ratio,
            "near_limit": bool(near_limit_scopes),
            "exhausted": bool(exhausted_scopes),
            "near_limit_scopes": near_limit_scopes,
            "exhausted_scopes": exhausted_scopes,
            "statuses": self._budget_status_snapshot(run_context),
        }

    def _budget_system_prompt(self) -> str | None:
        run_context = get_run_context()
        budget_snapshot = self._budget_status_snapshot(run_context)
        if not budget_snapshot:
            return None
        lines = [
            "Budget guidance:",
            "Optimize for quality within the remaining budget.",
        ]
        for scope in ("task", "team", "session", "workspace"):
            status = budget_snapshot.get(scope)
            if not isinstance(status, dict):
                continue
            usd_remaining = status.get("remaining_usd")
            usd_limit = status.get("limit_usd") or status.get("soft_limit_usd")
            token_remaining = status.get("remaining_tokens")
            token_limit = status.get("limit_tokens")
            line = f"- {scope}: ratio={status.get('ratio', 0.0)}"
            if usd_limit is not None:
                line += f", usd_remaining={usd_remaining}"
            if token_limit is not None:
                line += f", token_remaining={token_remaining}"
            if status.get("near_limit"):
                line += ", near_limit=true"
            if status.get("exhausted"):
                line += ", exhausted=true"
            lines.append(line)
        lines.append("Be concise when the budget is tight. Avoid unnecessary tool use or retries.")
        return "\n".join(lines)

    def _should_fail_on_budget_exhausted(self, run_context: RunContext | None) -> bool:
        if run_context is None:
            return bool(self._config.routing.fail_on_budget_exhausted)
        return bool(run_context.fail_on_budget_exhausted or self._config.routing.fail_on_budget_exhausted)

    def _strict_budget_failure_response(
        self,
        *,
        run_context: RunContext | None,
        stage: str,
        usage: Usage | None = None,
        raw: dict[str, Any] | None = None,
    ) -> Response | None:
        if not self._should_fail_on_budget_exhausted(run_context):
            return None
        runtime_status = self._runtime_budget_status(run_context)
        exhausted = [status for status in runtime_status.values() if status.exhausted]
        if not exhausted:
            return None
        scopes = ", ".join(status.scope for status in exhausted)
        limit_fragments: list[str] = []
        for status in exhausted:
            if status.limit_usd is not None:
                limit_fragments.append(
                    f"{status.scope} usd {status.spent_usd:.6f}/{status.limit_usd:.6f}"
                )
            if status.limit_tokens is not None:
                limit_fragments.append(
                    f"{status.scope} tokens {status.spent_tokens}/{status.limit_tokens}"
                )
        message = f"Budget exhausted during {stage}: {scopes}."
        if limit_fragments:
            message += " " + "; ".join(limit_fragments)
        failure_raw = dict(raw or {})
        failure_raw["budget_failure"] = {
            "stage": stage,
            "scopes": [status.scope for status in exhausted],
            "message": message,
            "error_type": BudgetExhaustedError.__name__,
        }
        return self._attach_budget_metadata(
            Response(
                content=message,
                content_parts=[],
                tool_calls=[],
                usage=usage or Usage(input_tokens=0, output_tokens=0),
                raw=failure_raw,
                status="failed",
            ),
            run_context=run_context,
        )

    def _effective_self_qa_max_retries(self, run_context: RunContext | None) -> int:
        configured = max(0, int(self._config.self_qa.max_retries))
        pressure = self._budget_pressure(run_context)
        if pressure["exhausted"]:
            return 0
        if pressure["near_limit"]:
            return min(configured, 0)
        return configured

    def _effective_max_tool_rounds(self, run_context: RunContext | None) -> int:
        pressure = self._budget_pressure(run_context)
        if pressure["exhausted"]:
            return 1
        if pressure["near_limit"]:
            return min(_MAX_TOOL_ROUNDS, 2)
        return _MAX_TOOL_ROUNDS

    def _attach_budget_metadata(
        self,
        response: Response,
        *,
        run_context: RunContext | None,
        self_qa_max_retries: int | None = None,
        max_tool_rounds: int | None = None,
    ) -> Response:
        raw = dict(response.raw or {})
        budget_payload = self._budget_status_snapshot(run_context)
        if budget_payload:
            budget_payload["budget_preference"] = self._routing_budget_preference(run_context)
            pressure = self._budget_pressure(run_context)
            budget_payload["near_limit"] = pressure["near_limit"]
            budget_payload["exhausted"] = pressure["exhausted"]
            budget_payload["budget_ratio"] = round(float(pressure["budget_ratio"]), 6)
            if self_qa_max_retries is not None:
                budget_payload["self_qa_max_retries"] = int(self_qa_max_retries)
            if max_tool_rounds is not None:
                budget_payload["max_tool_rounds"] = int(max_tool_rounds)
            raw["budget"] = budget_payload
        return Response(
            content=response.content,
            content_parts=response.content_parts,
            tool_calls=response.tool_calls,
            usage=response.usage,
            raw=raw,
            reasoning=response.reasoning,
            status=response.status,
        )

    def _resolve_runtime_budget_accounts(
        self,
        *,
        task_budget: Any,
        team_budget: Any,
        task_budget_usd: float | None,
        task_budget_tokens: int | None,
        team_budget_usd: float | None,
        team_budget_tokens: int | None,
        threshold_ratio: float | None,
    ) -> tuple[BudgetAccount | None, BudgetAccount | None]:
        effective_threshold = (
            self._default_budget_threshold_ratio
            if threshold_ratio is None
            else max(0.0, float(threshold_ratio))
        )
        task_account = ensure_budget_account(
            task_budget,
            scope="task",
            name=self.name,
            limit_usd=(
                task_budget_usd
                if task_budget_usd is not None
                else (
                    self._default_task_budget_usd
                    if self._default_task_budget_usd is not None
                    else self._config.routing.task_budget_usd
                )
            ),
            limit_tokens=(
                task_budget_tokens
                if task_budget_tokens is not None
                else (
                    self._default_task_budget_tokens
                    if self._default_task_budget_tokens is not None
                    else self._config.routing.task_budget_tokens
                )
            ),
            threshold_ratio=effective_threshold,
        )
        team_account = ensure_budget_account(
            team_budget,
            scope="team",
            name=self._team.name if self._team is not None else self.name,
            limit_usd=team_budget_usd,
            limit_tokens=team_budget_tokens,
            threshold_ratio=effective_threshold,
        )
        return task_account, team_account

    def _record_spend(
        self,
        *,
        provider: Provider,
        usage: Usage,
        run_context: RunContext | None,
        trace_id: str | None,
    ) -> None:
        amount_usd = self._estimate_response_cost_usd(self._provider_model_name(provider), usage)
        if run_context is not None:
            record_budget_usage(run_context.task_budget, usage=usage, amount_usd=amount_usd)
            record_budget_usage(run_context.team_budget, usage=usage, amount_usd=amount_usd)
        ledger = self._resolve_spend_ledger()
        if ledger is None or amount_usd <= 0:
            return
        try:
            ledger.record(
                SpendRecord(
                    agent_name=self.name,
                    workspace_id=run_context.workspace_id if run_context else None,
                    session_id=run_context.session_id if run_context else None,
                    trace_id=trace_id,
                    model=self._provider_model_name(provider),
                    amount_usd=amount_usd,
                )
            )
        except Exception:
            return

    def _routing_budget_preference(self, run_context: RunContext | None) -> str:
        preference = (
            run_context.budget_preference
            if run_context is not None and run_context.budget_preference
            else self._config.routing.budget_preference
        )
        return preference if preference in {"balanced", "cheap", "quality"} else "balanced"

    def _estimate_response_cost_usd(self, model_name: str, usage: Usage) -> float:
        input_rate = 0.0
        output_rate = 0.0
        for pattern, rates in DEFAULT_RATES.items():
            if pattern == model_name or ("*" in pattern and fnmatch(model_name, pattern)):
                input_rate, output_rate = rates
                break
        return (
            (usage.input_tokens * input_rate) / 1_000_000
            + (usage.output_tokens * output_rate) / 1_000_000
        )
