from datetime import datetime
from typing import Optional
from sqlmodel import SQLModel, Field, Column, select
from pydantic import BaseModel
from enum import Enum

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