from __future__ import annotations

import json
import os
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

from fastapi import FastAPI
from fastapi.responses import PlainTextResponse

SESSIONS_DIR = Path(os.environ.get("OPENCLAW_SESSIONS_DIR", "/root/.openclaw/agents/main/sessions"))
TIMEZONE = ZoneInfo(os.environ.get("OPENCLAW_EXPORTER_TZ", "Europe/Berlin"))

app = FastAPI(title="openclaw-daily-tokens-exporter")


@dataclass
class Totals:
    total: int = 0
    input: int = 0
    output: int = 0
    cache_read: int = 0
    cache_write: int = 0
    reasoning: int = 0
    entries: int = 0


def _as_int(value: object) -> int:
    if value is None:
        return 0
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, (int, float)):
        return int(value)
    return 0


def _today_in_tz() -> datetime.date:
    return datetime.now(TIMEZONE).date()


def collect_daily_totals() -> Totals:
    today = _today_in_tz()
    totals = Totals()

    for path in sorted(SESSIONS_DIR.glob("*.jsonl*")):
        try:
            with path.open("r", encoding="utf-8") as handle:
                for line in handle:
                    if not line.strip():
                        continue
                    try:
                        record = json.loads(line)
                    except json.JSONDecodeError:
                        continue

                    if record.get("type") != "message":
                        continue

                    message = record.get("message") or {}
                    usage = message.get("usage")
                    if not isinstance(usage, dict):
                        continue

                    timestamp_ms = message.get("timestamp")
                    if not isinstance(timestamp_ms, (int, float)):
                        continue

                    usage_date = datetime.fromtimestamp(timestamp_ms / 1000, TIMEZONE).date()
                    if usage_date != today:
                        continue

                    totals.entries += 1
                    totals.total += _as_int(usage.get("totalTokens"))
                    totals.input += _as_int(usage.get("input"))
                    totals.output += _as_int(usage.get("output"))
                    totals.cache_read += _as_int(usage.get("cacheRead"))
                    totals.cache_write += _as_int(usage.get("cacheWrite"))
                    totals.reasoning += _as_int(usage.get("reasoningTokens"))
        except OSError:
            continue

    return totals


@app.get("/metrics", response_class=PlainTextResponse)
def metrics() -> str:
    totals = collect_daily_totals()
    today = _today_in_tz().isoformat()

    lines = [
        "# HELP openclaw_daily_total_tokens Total OpenClaw tokens recorded today.",
        "# TYPE openclaw_daily_total_tokens gauge",
        f'openclaw_daily_total_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.total}',
        "# HELP openclaw_daily_input_tokens OpenClaw input tokens recorded today.",
        "# TYPE openclaw_daily_input_tokens gauge",
        f'openclaw_daily_input_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.input}',
        "# HELP openclaw_daily_output_tokens OpenClaw output tokens recorded today.",
        "# TYPE openclaw_daily_output_tokens gauge",
        f'openclaw_daily_output_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.output}',
        "# HELP openclaw_daily_cache_read_tokens OpenClaw cache-read tokens recorded today.",
        "# TYPE openclaw_daily_cache_read_tokens gauge",
        f'openclaw_daily_cache_read_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.cache_read}',
        "# HELP openclaw_daily_cache_write_tokens OpenClaw cache-write tokens recorded today.",
        "# TYPE openclaw_daily_cache_write_tokens gauge",
        f'openclaw_daily_cache_write_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.cache_write}',
        "# HELP openclaw_daily_reasoning_tokens OpenClaw reasoning tokens recorded today.",
        "# TYPE openclaw_daily_reasoning_tokens gauge",
        f'openclaw_daily_reasoning_tokens{{date="{today}",tz="{TIMEZONE.key}"}} {totals.reasoning}',
        "# HELP openclaw_daily_usage_entries OpenClaw usage entries recorded today.",
        "# TYPE openclaw_daily_usage_entries gauge",
        f'openclaw_daily_usage_entries{{date="{today}",tz="{TIMEZONE.key}"}} {totals.entries}',
        "",
    ]
    return "\n".join(lines)