# -*- coding: utf-8 -*-
"""
DWS_COACH_AREA_HOURS — 助教区域课时汇总表

按 (site_id, stat_month, assistant_id, area_code) 粒度聚合助教课时，
为财务看板助教分析板块的区域过滤提供数据支撑。

数据来源：dwd_assistant_service_log + dwd_assistant_service_log_ex + dim_table
更新策略：delete-before-insert（按 site_id + stat_month 删除后插入）

CHANGE 2026-03-29 | Prompt: 助教分析按区域细化 | 新建 ETL 任务
"""
from __future__ import annotations

import logging
from datetime import date
from decimal import Decimal
from typing import Any, Dict, List

from neozqyy_shared.area_mapping import (
    ALL_AREA_CODES,
    SPECIFIC_AREA_CODES,
    resolve_area_code,
)

from tasks.dws.base_dws_task import BaseDwsTask, TaskContext

logger = logging.getLogger(__name__)

_ZERO = Decimal("0")

# 课程类型映射（skill_name → 字段前缀）
_SKILL_MAP = {
    "基础课": "base",
    "附加课": "bonus",
    "包厢课": "room",
}


class CoachAreaHoursTask(BaseDwsTask):
    """助教区域课时汇总 ETL 任务"""

    def get_task_code(self) -> str:
        return "DWS_COACH_AREA_HOURS"

    def get_target_table(self) -> str:
        return "dws.dws_coach_area_hours"

    def get_primary_keys(self) -> List[str]:
        return ["site_id", "stat_month", "assistant_id", "area_code"]

    def extract(self, context: TaskContext) -> Dict[str, Any]:
        end_date = (
            context.window_end.date()
            if hasattr(context.window_end, "date")
            else context.window_end
        )
        site_id = context.store_id
        # 取当月第一天
        stat_month = end_date.replace(day=1)

        sql = """
            SELECT
                s.site_assistant_id AS assistant_id,
                dt.site_table_area_name AS area_name,
                s.skill_name,
                s.income_seconds,
                COALESCE(ex.is_trash, 0) AS is_trash
            FROM dwd.dwd_assistant_service_log s
            LEFT JOIN dwd.dwd_assistant_service_log_ex ex
                ON s.assistant_service_id = ex.assistant_service_id
            LEFT JOIN dwd.dim_table dt
                ON s.site_table_id = dt.table_id
                AND dt.scd2_is_current = 1
            WHERE s.site_id = %s
              AND s.is_delete = 0
              AND DATE(s.start_use_time) >= %s
              AND DATE(s.start_use_time) < %s
        """
        # 下月第一天
        if stat_month.month == 12:
            next_month = stat_month.replace(year=stat_month.year + 1, month=1)
        else:
            next_month = stat_month.replace(month=stat_month.month + 1)

        rows = self.db.query(sql, (site_id, stat_month, next_month))
        records = [dict(r) for r in rows] if rows else []

        self.logger.info(
            "%s: 提取 %d 条服务记录，月份 %s",
            self.get_task_code(), len(records), stat_month,
        )
        return {
            "records": records,
            "stat_month": stat_month,
            "site_id": site_id,
            "tenant_id": self.config.get("app.tenant_id", site_id),
        }

    def transform(
        self, extracted: Dict[str, Any], context: TaskContext
    ) -> List[Dict[str, Any]]:
        return transform_coach_area_hours(
            records=extracted["records"],
            stat_month=extracted["stat_month"],
            site_id=extracted["site_id"],
            tenant_id=extracted["tenant_id"],
        )

    def load(
        self, transformed: List[Dict[str, Any]], context: TaskContext
    ) -> dict:
        if not transformed:
            return {"inserted": 0}

        site_id = transformed[0]["site_id"]
        stat_month = transformed[0]["stat_month"]

        # delete-before-insert
        self.db.execute(
            "DELETE FROM dws.dws_coach_area_hours WHERE site_id = %s AND stat_month = %s",
            (site_id, stat_month),
        )

        cols = [
            "site_id", "tenant_id", "stat_month", "assistant_id", "area_code",
            "base_hours", "bonus_hours", "room_hours",
            "effective_hours", "trashed_hours",
            "base_service_count", "bonus_service_count", "room_service_count",
        ]
        placeholders = ", ".join(["%s"] * len(cols))
        insert_sql = f"INSERT INTO dws.dws_coach_area_hours ({', '.join(cols)}) VALUES ({placeholders})"

        for row in transformed:
            self.db.execute(insert_sql, tuple(row[c] for c in cols))

        self.logger.info(
            "%s: 写入 %d 行（site_id=%s, month=%s）",
            self.get_task_code(), len(transformed), site_id, stat_month,
        )
        return {"inserted": len(transformed)}


# ── 纯函数（可被回填脚本和属性测试直接调用） ─────────────────────────


def transform_coach_area_hours(
    records: List[Dict[str, Any]],
    stat_month: date,
    site_id: int,
    tenant_id: int,
) -> List[Dict[str, Any]]:
    """
    将 DWD 服务记录按 (assistant_id, area_code) 聚合为区域课时。

    返回列表，每个元素对应一行 dws_coach_area_hours。
    包含 9 个 area_code（hallA~ktv + hall + all）× N 个助教。
    """
    # 按 (assistant_id, area_code) 聚合
    # key: (assistant_id, area_code) → {base_hours, bonus_hours, ...}
    agg: Dict[tuple, Dict[str, Any]] = {}

    for rec in records:
        assistant_id = rec["assistant_id"]
        area_name = rec.get("area_name")
        skill_name = rec.get("skill_name")
        income_seconds = rec.get("income_seconds") or 0
        is_trash = rec.get("is_trash", 0)

        hours = Decimal(str(income_seconds)) / Decimal("3600")
        skill_prefix = _SKILL_MAP.get(skill_name)  # base/bonus/room/None

        # 解析区域
        area_code = resolve_area_code(area_name) if area_name else None
        if area_code is None:
            # 未知区域或 NULL table_id → 不计入具体区域，只计入 all
            target_areas = []
        else:
            target_areas = [area_code]

        # 更新具体区域
        for ac in target_areas:
            _accumulate(agg, assistant_id, ac, skill_prefix, hours, is_trash)

        # CHANGE 2026-03-29 | hall = 台球大厅（hallA+hallB+hallC），不含 vip/snooker/mahjong/ktv
        if area_code in ("hallA", "hallB", "hallC"):
            _accumulate(agg, assistant_id, "hall", skill_prefix, hours, is_trash)
        # all 始终累加
        _accumulate(agg, assistant_id, "all", skill_prefix, hours, is_trash)

    # 构建输出行
    result = []
    for (aid, ac), data in sorted(agg.items()):
        total_hours = data["base_hours"] + data["bonus_hours"] + data["room_hours"]
        effective = total_hours - data["trashed_hours"]
        result.append({
            "site_id": site_id,
            "tenant_id": tenant_id,
            "stat_month": stat_month,
            "assistant_id": aid,
            "area_code": ac,
            "base_hours": data["base_hours"],
            "bonus_hours": data["bonus_hours"],
            "room_hours": data["room_hours"],
            "effective_hours": max(effective, _ZERO),
            "trashed_hours": data["trashed_hours"],
            "base_service_count": data["base_service_count"],
            "bonus_service_count": data["bonus_service_count"],
            "room_service_count": data["room_service_count"],
        })

    return result


def _accumulate(
    agg: Dict[tuple, Dict[str, Any]],
    assistant_id: int,
    area_code: str,
    skill_prefix: str | None,
    hours: Decimal,
    is_trash: int,
) -> None:
    """累加单条服务记录到聚合字典"""
    key = (assistant_id, area_code)
    if key not in agg:
        agg[key] = {
            "base_hours": _ZERO,
            "bonus_hours": _ZERO,
            "room_hours": _ZERO,
            "trashed_hours": _ZERO,
            "base_service_count": 0,
            "bonus_service_count": 0,
            "room_service_count": 0,
        }
    bucket = agg[key]

    if is_trash == 1:
        bucket["trashed_hours"] += hours
        return

    if skill_prefix:
        bucket[f"{skill_prefix}_hours"] += hours
        bucket[f"{skill_prefix}_service_count"] += 1
    # skill_name 为 NULL 的记录：不计入任何课程类型的 hours，但已计入 hall/all
