# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========

from enum import Enum

from pydantic import BaseModel
from sqlmodel import Field, select

from app.model.abstract.model import AbstractModel, DefaultTimes


class UserStatActionEnum(str, Enum):
    download_count = "download_count"
    register_count = "register_count"
    task_complete_count = "task_complete_count"
    task_failed_count = "task_failed_count"
    file_download_count = "file_download_count"
    file_generate_count = "file_generate_count"
    paid_amount_on_avg_task = "paid_amount_on_avg_task"


class UserStatActionIn(BaseModel):
    user_id: int | None = None
    action: UserStatActionEnum
    value: int = 1
    model_type: str | None = None


class UserStat(AbstractModel, DefaultTimes, table=True):
    id: int = Field(default=None, primary_key=True)
    user_id: int = Field(foreign_key="user.id", index=True, description="User ID")
    # Model usage type: 'cloud' or 'local'
    model_type: str = Field(default="unused", description="Model usage type: 'cloud' or 'local'")
    # Product page statistics
    download_count: int = Field(default=0, description="Number of downloads by the user")
    register_count: int = Field(default=0, description="Number of registrations (for product page)")
    task_complete_count: int = Field(default=0, description="Number of tasks completed by the user")
    task_failed_count: int = Field(default=0, description="Number of tasks failed by the user")
    file_download_count: int = Field(default=0, description="Number of files downloaded by the user")
    file_generate_count: int = Field(default=0, description="Number of files generated by the user")
    # Payment statistics
    paid_amount_on_avg_task: int = Field(default=0, description="Total paid amount on average task completion")

    @classmethod
    def record_action(cls, session, action_in: UserStatActionIn):
        """
        Record or update user operation statistics using a Pydantic model.
        If no record exists for the user, create one. Otherwise, update the corresponding field.
        Supported actions: download_count, register_count, task_complete_count, task_failed_count, file_download_count, file_generate_count, paid_amount_on_avg_task.
        If model_type is provided, update it as well.
        """
        stat = session.exec(select(cls).where(cls.user_id == action_in.user_id)).first()
        if not stat:
            stat = cls(user_id=action_in.user_id)
            session.add(stat)
        if action_in.action in [
            UserStatActionEnum.download_count,
            UserStatActionEnum.register_count,
            UserStatActionEnum.task_complete_count,
            UserStatActionEnum.task_failed_count,
            UserStatActionEnum.file_download_count,
            UserStatActionEnum.file_generate_count,
        ]:
            setattr(stat, action_in.action.value, getattr(stat, action_in.action.value, 0) + action_in.value)
        elif action_in.action == UserStatActionEnum.paid_amount_on_avg_task:
            stat.paid_amount_on_avg_task += action_in.value
        else:
            raise ValueError(f"Unsupported action: {action_in.action}")
        if action_in.model_type is not None:
            stat.model_type = action_in.model_type
        session.add(stat)
        session.commit()
        session.refresh(stat)
        return stat


class UserStatOut(BaseModel):
    model_type: str | None = None
    download_count: int = 0
    register_count: int = 0
    task_complete_count: int = 0
    task_failed_count: int = 0
    file_download_count: int = 0
    file_generate_count: int = 0
    paid_amount_on_avg_task: int = 0
    # cusotmer
    task_queries: int = 0
    mcp_install_count: int = 0
    storage_used: float = 0