"""World model warmup helpers for startup."""

from __future__ import annotations

import os
from datetime import datetime, timedelta

import sqlalchemy as sa
from sqlalchemy.exc import OperationalError, ProgrammingError

from .core.logging import get_logger
from .database import AsyncSessionLocal
from .main_db_guard import is_missing_table_error
from .models import Project, WorkItem
from .services.ai_pmo_kernel_service import AIPMOKernelService
from .utils.timezone import now_cn_aware

logger = get_logger(__name__)
ACTIVE_PROJECT_STATUSES = ("active", "planning", "todo", "doing")


async def get_recent_active_projects(
    *,
    days: int,
    limit_per_tenant: int,
    max_projects: int,
) -> list[tuple[int, int]]:
    since = (now_cn_aware() - timedelta(days=max(1, int(days)))).replace(tzinfo=None)
    per_tenant_limit = max(1, int(limit_per_tenant))
    total_limit = max(per_tenant_limit, int(max_projects))

    async with AsyncSessionLocal() as db:
        try:
            rows = (
                await db.execute(
                    sa.select(
                        WorkItem.tenant_id.label("tenant_id"),
                        WorkItem.project_id.label("project_id"),
                        sa.func.max(WorkItem.created_at).label("last_seen_at"),
                    )
                    .join(
                        Project,
                        sa.and_(
                            Project.tenant_id == WorkItem.tenant_id,
                            Project.id == WorkItem.project_id,
                        ),
                    )
                    .where(
                        sa.and_(
                            WorkItem.project_id.is_not(None),
                            WorkItem.created_at >= since,
                            Project.status.in_(ACTIVE_PROJECT_STATUSES),
                        )
                    )
                    .group_by(WorkItem.tenant_id, WorkItem.project_id)
                )
            ).all()
        except (OperationalError, ProgrammingError) as exc:
            if not is_missing_table_error(exc):
                raise
            logger.warning("[WM_WARMUP] 活跃项目查询跳过", reason="missing_table", error=str(exc))
            return []

    bucket: dict[int, list[tuple[datetime, int]]] = {}
    for row in rows:
        tenant_id = int(getattr(row, "tenant_id"))
        project_id = int(getattr(row, "project_id"))
        last_seen = getattr(row, "last_seen_at")
        last_seen_at = last_seen if isinstance(last_seen, datetime) else datetime.min
        bucket.setdefault(tenant_id, []).append((last_seen_at, project_id))

    picked: list[tuple[datetime, int, int]] = []
    for tenant_id, entries in bucket.items():
        entries.sort(key=lambda item: item[0], reverse=True)
        for last_seen_at, project_id in entries[:per_tenant_limit]:
            picked.append((last_seen_at, tenant_id, project_id))

    picked.sort(key=lambda item: item[0], reverse=True)
    return [(tenant_id, project_id) for _, tenant_id, project_id in picked[:total_limit]]


async def warmup_world_models() -> None:
    days = max(1, int(os.getenv("WORLD_MODEL_WARMUP_DAYS", "7")))
    limit_per_tenant = max(1, int(os.getenv("WORLD_MODEL_WARMUP_LIMIT_PER_TENANT", "5")))
    max_projects = max(limit_per_tenant, int(os.getenv("WORLD_MODEL_WARMUP_MAX_PROJECTS", "50")))
    try:
        active_projects = await get_recent_active_projects(
            days=days,
            limit_per_tenant=limit_per_tenant,
            max_projects=max_projects,
        )
        warmed = 0
        for tenant_id, project_id in active_projects:
            try:
                async with AsyncSessionLocal() as db:
                    kernel = AIPMOKernelService(db, tenant_id=int(tenant_id))
                    await kernel.build_world_model(
                        tenant_id=int(tenant_id),
                        project_id=int(project_id),
                        sender="wm_warmup",
                    )
                warmed += 1
            except Exception as exc:
                logger.warning(
                    "[WM_WARMUP] 项目预热失败",
                    tenant_id=tenant_id,
                    project_id=project_id,
                    error=str(exc),
                )
        logger.info(
            "[WM_WARMUP] 预热完成",
            count=warmed,
            total=len(active_projects),
            days=days,
            limit_per_tenant=limit_per_tenant,
            max_projects=max_projects,
        )
    except Exception as exc:
        logger.warning("[WM_WARMUP] 预热失败", error=str(exc))