"""Token statistics API routes"""
from datetime import date, timedelta, datetime, timezone
from flask import Blueprint, request, g
from sqlalchemy import func, extract
from backend.models import TokenUsage, Message, Conversation
from backend.utils.helpers import ok, err
from backend import db

bp = Blueprint("stats", __name__)


def _utc_today():
    """Get today's date in UTC to match stored timestamps."""
    return datetime.now(timezone.utc).date()


@bp.route("/api/stats/tokens", methods=["GET"])
def token_stats():
    """Get token usage statistics"""
    user = g.current_user
    period = request.args.get("period", "daily")

    today = _utc_today()

    if period == "daily":
        stats = TokenUsage.query.filter_by(user_id=user.id, date=today).all()
        # Hourly breakdown from Message table
        hourly = _build_hourly_stats(user.id, today)
        result = {
            "period": "daily",
            "date": today.isoformat(),
            "prompt_tokens": sum(s.prompt_tokens for s in stats),
            "completion_tokens": sum(s.completion_tokens for s in stats),
            "total_tokens": sum(s.total_tokens for s in stats),
            "hourly": hourly,
            "by_model": {
                s.model: {
                    "prompt": s.prompt_tokens,
                    "completion": s.completion_tokens,
                    "total": s.total_tokens
                } for s in stats
            }
        }
    elif period == "weekly":
        start_date = today - timedelta(days=6)
        stats = TokenUsage.query.filter(
            TokenUsage.user_id == user.id,
            TokenUsage.date >= start_date,
            TokenUsage.date <= today
        ).all()
        
        result = _build_period_result(stats, "weekly", start_date, today, 7)
    elif period == "monthly":
        start_date = today - timedelta(days=29)
        stats = TokenUsage.query.filter(
            TokenUsage.user_id == user.id,
            TokenUsage.date >= start_date,
            TokenUsage.date <= today
        ).all()
        
        result = _build_period_result(stats, "monthly", start_date, today, 30)
    else:
        return err(400, "invalid period")
    
    return ok(result)


def _build_period_result(stats, period, start_date, end_date, days):
    """Build result for period-based statistics"""
    daily_data = {}
    for s in stats:
        d = s.date.isoformat()
        if d not in daily_data:
            daily_data[d] = {"prompt": 0, "completion": 0, "total": 0}
        daily_data[d]["prompt"] += s.prompt_tokens
        daily_data[d]["completion"] += s.completion_tokens
        daily_data[d]["total"] += s.total_tokens
    
    # Fill missing dates
    for i in range(days):
        d = (end_date - timedelta(days=days - 1 - i)).isoformat()
        if d not in daily_data:
            daily_data[d] = {"prompt": 0, "completion": 0, "total": 0}
    
    # Aggregate by model
    by_model = {}
    for s in stats:
        if s.model not in by_model:
            by_model[s.model] = {"prompt": 0, "completion": 0, "total": 0}
        by_model[s.model]["prompt"] += s.prompt_tokens
        by_model[s.model]["completion"] += s.completion_tokens
        by_model[s.model]["total"] += s.total_tokens

    return {
        "period": period,
        "start_date": start_date.isoformat(),
        "end_date": end_date.isoformat(),
        "prompt_tokens": sum(s.prompt_tokens for s in stats),
        "completion_tokens": sum(s.completion_tokens for s in stats),
        "total_tokens": sum(s.total_tokens for s in stats),
        "daily": daily_data,
        "by_model": by_model,
    }


def _build_hourly_stats(user_id, day):
    """Build hourly token breakdown for a given day (UTC) from Message table."""
    day_start = datetime.combine(day, datetime.min.time()).replace(tzinfo=timezone.utc)
    day_end = datetime.combine(day, datetime.max.time()).replace(tzinfo=timezone.utc)

    conv_ids = (
        db.session.query(Conversation.id)
        .filter(Conversation.user_id == user_id)
        .subquery()
    )

    rows = (
        db.session.query(
            extract("hour", Message.created_at).label("hour"),
            func.sum(Message.token_count).label("total"),
        )
        .filter(
            Message.conversation_id.in_(conv_ids),
            Message.role == "assistant",
            Message.created_at >= day_start,
            Message.created_at <= day_end,
        )
        .group_by(extract("hour", Message.created_at))
        .all()
    )

    hourly = {}
    for r in rows:
        h = int(r.hour)
        hourly[str(h)] = {"total": r.total or 0}

    return hourly