"""Cost & Credit Tracker — instruments LLM token usage and Snowflake credits.

Provides per-run cost accumulation for:
  - LLM API calls (Anthropic Claude, Snowflake Cortex)
  - Snowflake warehouse credit consumption

Designed for integration into the Enertia platform workflow to produce
real cost data for agent pricing estimates.
"""

from __future__ import annotations

import datetime as dt
import logging
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────────────────────
# Pricing tables (USD per 1 M tokens)
# ──────────────────────────────────────────────────────────────────────────────

LLM_PRICING: Dict[str, Dict[str, float]] = {
    # Anthropic models
    "claude-sonnet-4-5-20250929": {"input": 3.00, "output": 15.00},
    "claude-haiku-4-5-20251001":  {"input": 0.80, "output": 4.00},
    "claude-opus-4-6":            {"input": 15.00, "output": 75.00},
    # Snowflake Cortex models (no per-token charge — credit-based)
    "mistral-large":  {"input": 0.0, "output": 0.0},
    "llama3-70b":     {"input": 0.0, "output": 0.0},
    "snowflake-arctic": {"input": 0.0, "output": 0.0},
}


# ──────────────────────────────────────────────────────────────────────────────
# Record types
# ──────────────────────────────────────────────────────────────────────────────

@dataclass
class LLMCallRecord:
    """Single LLM API call with token counts and cost."""
    timestamp: str
    model: str
    input_tokens: int
    output_tokens: int
    estimated_cost_usd: float
    context: str = ""


@dataclass
class SnowflakeCreditRecord:
    """Snapshot of Snowflake warehouse credit usage for a time window."""
    warehouse_name: str
    credits_used: float
    estimated_cost_usd: float
    query_count: int
    window_start: str
    window_end: str


# ──────────────────────────────────────────────────────────────────────────────
# CostTracker
# ──────────────────────────────────────────────────────────────────────────────

class CostTracker:
    """Accumulates LLM and Snowflake costs for a single workflow run."""

    def __init__(self, run_id: str, snowflake_credit_rate: float = 3.00):
        self.run_id = run_id
        self.snowflake_credit_rate = snowflake_credit_rate
        self.start_time = dt.datetime.now(dt.UTC).isoformat()

        self._llm_calls: List[LLMCallRecord] = []
        self._sf_credits: List[SnowflakeCreditRecord] = []

    # ── LLM tracking ─────────────────────────────────────────────────────

    def record_llm_call(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        context: str = "",
    ) -> LLMCallRecord:
        """Record a single LLM API call and calculate its cost."""
        pricing = LLM_PRICING.get(model, {"input": 0.0, "output": 0.0})
        cost = (
            (input_tokens / 1_000_000) * pricing["input"]
            + (output_tokens / 1_000_000) * pricing["output"]
        )

        record = LLMCallRecord(
            timestamp=dt.datetime.now(dt.UTC).isoformat(),
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            estimated_cost_usd=round(cost, 6),
            context=context,
        )
        self._llm_calls.append(record)
        logger.debug(
            "LLM cost recorded: model=%s in=%d out=%d cost=$%.6f ctx=%s",
            model, input_tokens, output_tokens, cost, context,
        )
        return record

    # ── Snowflake credit tracking ────────────────────────────────────────

    def record_snowflake_credits(
        self,
        warehouse: str,
        credits: float,
        query_count: int,
        window_start: str,
        window_end: str,
    ) -> SnowflakeCreditRecord:
        """Record a Snowflake credit consumption snapshot."""
        cost = credits * self.snowflake_credit_rate
        record = SnowflakeCreditRecord(
            warehouse_name=warehouse,
            credits_used=round(credits, 4),
            estimated_cost_usd=round(cost, 4),
            query_count=query_count,
            window_start=window_start,
            window_end=window_end,
        )
        self._sf_credits.append(record)
        logger.debug(
            "SF credits recorded: wh=%s credits=%.4f cost=$%.4f queries=%d",
            warehouse, credits, cost, query_count,
        )
        return record

    @staticmethod
    def query_warehouse_credits(
        cursor,
        warehouse: str,
        start_time: str,
        end_time: str,
    ) -> Tuple[float, int]:
        """Query Snowflake QUERY_HISTORY for credit/query usage.

        Uses INFORMATION_SCHEMA.QUERY_HISTORY (real-time, no lag) rather
        than WAREHOUSE_METERING_HISTORY (up to 3-hour delay).

        Returns:
            (estimated_credits, query_count)
        """
        sql = """
            SELECT
                COUNT(*) AS query_count,
                SUM(CREDITS_USED_CLOUD_SERVICES) AS credits_cloud
            FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY(
                DATE_RANGE_START => %s::TIMESTAMP_LTZ,
                DATE_RANGE_END   => %s::TIMESTAMP_LTZ
            ))
            WHERE WAREHOUSE_NAME = %s
        """
        cursor.execute(sql, (start_time, end_time, warehouse.upper()))
        row = cursor.fetchone()
        if row:
            query_count = row[0] or 0
            credits = float(row[1] or 0.0)
            return credits, query_count
        return 0.0, 0

    # ── Summary ──────────────────────────────────────────────────────────

    def get_summary(self) -> Dict[str, Any]:
        """Return a JSON-serializable cost summary for this run."""
        total_input = sum(r.input_tokens for r in self._llm_calls)
        total_output = sum(r.output_tokens for r in self._llm_calls)
        llm_cost = sum(r.estimated_cost_usd for r in self._llm_calls)

        sf_credits = sum(r.credits_used for r in self._sf_credits)
        sf_cost = sum(r.estimated_cost_usd for r in self._sf_credits)

        return {
            "run_id": self.run_id,
            "start_time": self.start_time,
            "end_time": dt.datetime.now(dt.UTC).isoformat(),
            "llm": {
                "call_count": len(self._llm_calls),
                "total_input_tokens": total_input,
                "total_output_tokens": total_output,
                "estimated_cost_usd": round(llm_cost, 6),
                "calls": [asdict(r) for r in self._llm_calls],
            },
            "snowflake": {
                "total_credits": round(sf_credits, 4),
                "estimated_cost_usd": round(sf_cost, 4),
                "credit_rate_usd": self.snowflake_credit_rate,
                "records": [asdict(r) for r in self._sf_credits],
            },
            "total_estimated_cost_usd": round(llm_cost + sf_cost, 4),
        }

    # ── Metrics integration ──────────────────────────────────────────────

    def record_to_metrics_store(self, store) -> int:
        """Push summary metrics into the observability MetricsStore.

        Records 7 gauge metrics tagged with the run_id.

        Args:
            store: A MetricsStore instance (from src.observability)

        Returns:
            Number of metrics recorded
        """
        from .types import Metric, MetricType

        summary = self.get_summary()
        tags = {"run_id": self.run_id, "project": "enertia"}
        metrics = [
            ("cost.llm.call_count",        float(summary["llm"]["call_count"])),
            ("cost.llm.input_tokens",      float(summary["llm"]["total_input_tokens"])),
            ("cost.llm.output_tokens",     float(summary["llm"]["total_output_tokens"])),
            ("cost.llm.estimated_usd",     summary["llm"]["estimated_cost_usd"]),
            ("cost.snowflake.credits",     summary["snowflake"]["total_credits"]),
            ("cost.snowflake.estimated_usd", summary["snowflake"]["estimated_cost_usd"]),
            ("cost.total_estimated_usd",   summary["total_estimated_cost_usd"]),
        ]

        count = 0
        for name, value in metrics:
            store.record(Metric(
                name=name,
                type=MetricType.GAUGE,
                value=float(value),
                tags=tags,
                unit="usd" if "usd" in name else "count",
            ))
            count += 1

        logger.info("Recorded %d cost metrics for run %s", count, self.run_id)
        return count
