""" 简单任务记录服务

这是一个轻量级的同步任务管理器，仅供 Agent 工具（agents/tools/*.py）用于记录任务状态。它不负责任务的实际执行。

对于主要的生成接口（generations.py），请使用 task_manager，
它提供完整的异步任务执行和工作池调度功能。

用法说明：
- Agent 工具：用于简单的任务增删改查（CRUD）操作
- API 端点：请从 task_manager.py 使用 task_manager
""" 
import logging
import json
import threading
from typing import Optional, List, Dict, Any
from datetime import datetime

from sqlmodel import Session, select, desc

from src.config.settings import DB_PATH
from src.models.schemas import Task
from src.config.database import engine, init_db
from src.models.entities import TaskDB

logger = logging.getLogger(__name__)

class TaskManager:
    def __init__(self):
        self._lock = threading.Lock()
        init_db()

    def create_task(self, type: str, model: str, params: Dict[str, Any], user_id: Optional[str] = None) -> Task:
        # 创建 TaskDB object
        # 注意: id, created_at, updated_at have defaults in TaskDB
        task_db = TaskDB(
            type=type,
            model=model,
            params=params,
            status="pending",
            user_id=user_id
        )
        
        with Session(engine) as session:
            session.add(task_db)
            session.commit()
            session.refresh(task_db)
            
            return self._db_to_task(task_db)

    def get_task(self, task_id: str) -> Optional[Task]:
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                return None
            return self._db_to_task(task_db)
            
    def list_tasks(self, type: Optional[str] = None, limit: int = 50, offset: int = 0) -> List[Task]:
        with Session(engine) as session:
            statement = select(TaskDB).order_by(TaskDB.created_at.desc())
            if type:
                statement = statement.where(TaskDB.type == type)
            statement = statement.offset(offset).limit(limit)
            
            results = session.exec(statement).all()
            return [self._db_to_task(task) for task in results]

    def update_task(self, task_id: str, **updates) -> Optional[Task]:
        with self._lock: # Keep lock for thread safety if needed
            with Session(engine) as session:
                task_db = session.get(TaskDB, task_id)
                if not task_db:
                    return None
                
                updates['updated_at'] = datetime.now().timestamp()
                
                for k, v in updates.items():
                    if hasattr(task_db, k):
                        setattr(task_db, k, v)
                
                session.add(task_db)
                session.commit()
                session.refresh(task_db)
                
                return self._db_to_task(task_db)

    def _db_to_task(self, task_db: TaskDB) -> Task:
        return Task(
            id=task_db.id,
            type=task_db.type, # type: ignore
            status=task_db.status,
            created_at=task_db.created_at,
            updated_at=task_db.updated_at,
            model=task_db.model or "", # Handle optional model if Pydantic requires str
            provider=task_db.provider,
            params=task_db.params,
            provider_task_id=task_db.provider_task_id,
            result=task_db.result,
            error=task_db.error,
            retry_count=task_db.retry_count,
            max_retries=task_db.max_retries,
            started_at=task_db.started_at,
            completed_at=task_db.completed_at,
            user_id=task_db.user_id,
            project_id=task_db.project_id,
            deleted_at=task_db.deleted_at
        )

task_manager = TaskManager()
