from __future__ import annotations

from datetime import date, datetime, timezone
from typing import Any

from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.enums.upload_resource import UploadResourceTypeEnum
from app.models.generated_resource import GeneratedResource
from app.models.upload_resource import UploadResource
from app.models.user_resource_month_stat import UserResourceMonthStat
from app.models.user_resource_total_stat import UserResourceTotalStat
from app.utils.id_gen import generate_id


def resource_month_from_datetime(value: datetime | None = None) -> date:
    value = value or datetime.now(timezone.utc)
    return date(value.year, value.month, 1)


def _int(value: Any) -> int:
    return int(value or 0)


def _add_non_negative(obj: Any, field: str, delta: int) -> None:
    current = _int(getattr(obj, field, 0))
    setattr(obj, field, max(0, current + int(delta or 0)))


def _add_raw(obj: Any, field: str, delta: int) -> None:
    current = _int(getattr(obj, field, 0))
    setattr(obj, field, current + int(delta or 0))


async def get_or_create_month_stat(db: AsyncSession, user_id: str, stat_month: date) -> UserResourceMonthStat:
    result = await db.execute(
        select(UserResourceMonthStat).where(
            UserResourceMonthStat.user_id == user_id,
            UserResourceMonthStat.stat_month == stat_month,
        ).limit(1)
    )
    stat = result.scalar_one_or_none()
    if stat:
        return stat
    stat = UserResourceMonthStat(id=generate_id(), user_id=user_id, stat_month=stat_month)
    db.add(stat)
    await db.flush()
    return stat


async def get_or_create_total_stat(db: AsyncSession, user_id: str, *, for_update: bool = False) -> UserResourceTotalStat:
    stmt = select(UserResourceTotalStat).where(UserResourceTotalStat.user_id == user_id).limit(1)
    if for_update:
        stmt = stmt.with_for_update()
    result = await db.execute(stmt)
    stat = result.scalar_one_or_none()
    if stat:
        return stat
    stat = UserResourceTotalStat(id=generate_id(), user_id=user_id)
    db.add(stat)
    await db.flush()
    if for_update:
        result = await db.execute(
            select(UserResourceTotalStat)
            .where(UserResourceTotalStat.user_id == user_id)
            .with_for_update()
            .limit(1)
        )
        locked = result.scalar_one_or_none()
        if locked:
            return locked
    return stat


async def apply_upload_resource_stat_delta(
    db: AsyncSession,
    *,
    user_id: str,
    stat_month: date,
    resource_type: str,
    active_size_delta: int = 0,
    active_count_delta: int = 0,
    deleted_size_delta: int = 0,
    deleted_count_delta: int = 0,
    upload_size_delta: int = 0,
    upload_count_delta: int = 0,
) -> None:
    month_stat = await get_or_create_month_stat(db, user_id, stat_month)
    total_stat = await get_or_create_total_stat(db, user_id)
    now = datetime.now(timezone.utc)

    for stat in (month_stat, total_stat):
        _add_non_negative(stat, "active_size_bytes", active_size_delta)
        _add_non_negative(stat, "active_count", active_count_delta)
        _add_non_negative(stat, "deleted_size_bytes", deleted_size_delta)
        _add_non_negative(stat, "deleted_count", deleted_count_delta)
        _add_raw(stat, "upload_size_bytes", upload_size_delta)
        _add_raw(stat, "upload_count", upload_count_delta)

        if resource_type == UploadResourceTypeEnum.IMAGE.value:
            _add_non_negative(stat, "image_size_bytes", active_size_delta)
            _add_non_negative(stat, "image_count", active_count_delta)
        elif resource_type == UploadResourceTypeEnum.VIDEO.value:
            _add_non_negative(stat, "video_size_bytes", active_size_delta)
            _add_non_negative(stat, "video_count", active_count_delta)
        elif resource_type == UploadResourceTypeEnum.AUDIO.value:
            _add_non_negative(stat, "audio_size_bytes", active_size_delta)
            _add_non_negative(stat, "audio_count", active_count_delta)
        elif resource_type == UploadResourceTypeEnum.SHOT_SEGMENT.value:
            _add_non_negative(stat, "shot_segment_size_bytes", active_size_delta)
            _add_non_negative(stat, "shot_segment_count", active_count_delta)

        stat.last_recalculated_at = now


async def release_upload_resource_capacity(db: AsyncSession, resource: UploadResource, *, released_at: datetime | None = None) -> bool:
    if resource.capacity_released_at is not None:
        return False
    released_at = released_at or datetime.now(timezone.utc)
    stat_month = resource_month_from_datetime(resource.created_at or released_at)
    size = _int(resource.file_size_bytes)
    resource.capacity_released_at = released_at
    await apply_upload_resource_stat_delta(
        db,
        user_id=resource.user_id,
        stat_month=stat_month,
        resource_type=resource.resource_type,
        active_size_delta=-size,
        active_count_delta=-1,
        deleted_size_delta=size,
        deleted_count_delta=1,
    )
    return True


async def rebuild_user_resource_stats(db: AsyncSession, *, user_ids: list[str] | None = None) -> dict[str, int]:
    """按数据库真实资源账本重算统计。

    说明：这里只重置并重算 user_ids 范围内的统计。未传 user_ids 时重算所有在资源表中出现过的用户。
    """
    if user_ids is None:
        ids: set[str] = set()
        for model in (GeneratedResource, UploadResource):
            result = await db.execute(select(model.user_id).distinct())
            ids.update(v for v in result.scalars().all() if v)
        user_ids = sorted(ids)
    else:
        user_ids = sorted({v for v in user_ids if v})
    if not user_ids:
        return {"users": 0, "month_rows": 0, "total_rows": 0}

    await db.execute(update(UserResourceTotalStat).where(UserResourceTotalStat.user_id.in_(user_ids)).values(
        active_size_bytes=0,
        deleted_size_bytes=0,
        total_generated_size_bytes=0,
        upload_size_bytes=0,
        image_size_bytes=0,
        video_size_bytes=0,
        audio_size_bytes=0,
        shot_segment_size_bytes=0,
        active_count=0,
        deleted_count=0,
        upload_count=0,
        image_count=0,
        video_count=0,
        audio_count=0,
        shot_segment_count=0,
        last_recalculated_at=datetime.now(timezone.utc),
    ))
    await db.execute(update(UserResourceMonthStat).where(UserResourceMonthStat.user_id.in_(user_ids)).values(
        active_size_bytes=0,
        deleted_size_bytes=0,
        total_generated_size_bytes=0,
        upload_size_bytes=0,
        image_size_bytes=0,
        video_size_bytes=0,
        audio_size_bytes=0,
        shot_segment_size_bytes=0,
        active_count=0,
        deleted_count=0,
        upload_count=0,
        image_count=0,
        video_count=0,
        audio_count=0,
        shot_segment_count=0,
        last_recalculated_at=datetime.now(timezone.utc),
    ))

    month_rows = 0
    total_rows = 0

    gen_rows = await db.execute(
        select(
            GeneratedResource.user_id,
            GeneratedResource.resource_month,
            GeneratedResource.resource_type,
            GeneratedResource.deleted_at,
            func.count(GeneratedResource.id),
            func.coalesce(func.sum(GeneratedResource.file_size_bytes), 0),
        ).where(GeneratedResource.user_id.in_(user_ids)).group_by(
            GeneratedResource.user_id,
            GeneratedResource.resource_month,
            GeneratedResource.resource_type,
            GeneratedResource.deleted_at,
        )
    )
    for user_id, month, rtype, deleted_at, count, size in gen_rows.all():
        active = deleted_at is None
        await apply_upload_resource_stat_delta(
            db,
            user_id=user_id,
            stat_month=month,
            resource_type=rtype,
            active_size_delta=int(size or 0) if active else 0,
            active_count_delta=int(count or 0) if active else 0,
            deleted_size_delta=0 if active else int(size or 0),
            deleted_count_delta=0 if active else int(count or 0),
        )
        # 生成资源字段单独累加
        month_stat = await get_or_create_month_stat(db, user_id, month)
        total_stat = await get_or_create_total_stat(db, user_id)
        if active:
            _add_raw(month_stat, "total_generated_size_bytes", int(size or 0))
            _add_raw(total_stat, "total_generated_size_bytes", int(size or 0))

    upload_rows = await db.execute(
        select(
            UploadResource.user_id,
            func.date_trunc("month", UploadResource.created_at).label("month"),
            UploadResource.resource_type,
            UploadResource.deleted_at,
            func.count(UploadResource.id),
            func.coalesce(func.sum(UploadResource.file_size_bytes), 0),
        ).where(UploadResource.user_id.in_(user_ids)).group_by(
            UploadResource.user_id,
            "month",
            UploadResource.resource_type,
            UploadResource.deleted_at,
        )
    )
    for user_id, month_dt, rtype, deleted_at, count, size in upload_rows.all():
        month = resource_month_from_datetime(month_dt or datetime.now(timezone.utc))
        active = deleted_at is None
        await apply_upload_resource_stat_delta(
            db,
            user_id=user_id,
            stat_month=month,
            resource_type=rtype,
            active_size_delta=int(size or 0) if active else 0,
            active_count_delta=int(count or 0) if active else 0,
            deleted_size_delta=0 if active else int(size or 0),
            deleted_count_delta=0 if active else int(count or 0),
            upload_size_delta=int(size or 0) if active else 0,
            upload_count_delta=int(count or 0) if active else 0,
        )

    await db.flush()
    total_rows = len(user_ids)
    month_count = await db.execute(select(func.count(UserResourceMonthStat.id)).where(UserResourceMonthStat.user_id.in_(user_ids)))
    month_rows = int(month_count.scalar_one() or 0)
    return {"users": len(user_ids), "month_rows": month_rows, "total_rows": total_rows}