# #region update_language_stats [C:3] [TYPE Function] [SEMANTICS translate, language, stats, aggregation]
# @BRIEF Aggregate TranslationLanguage entries by language_code and update TranslationRunLanguageStats.
# @SIDE_EFFECT DB writes on language_stats objects.
# @RELATION DEPENDS_ON -> [TranslationRecord]
# @RELATION DEPENDS_ON -> [TranslationLanguage]
# @RELATION DEPENDS_ON -> [TranslationRunLanguageStats]

from collections import defaultdict

from sqlalchemy.orm import Session

from ...core.logger import logger
from ...models.translate import (
    TranslationLanguage,
    TranslationRecord,
    TranslationRunLanguageStats,
)
from .events import TranslationEventLog


def update_language_stats(
    db: Session,
    event_log: TranslationEventLog | None,
    run_id: str,
    language_stats_map: dict[str, TranslationRunLanguageStats],
) -> None:
    """Aggregate TranslationLanguage entries by language_code and update TranslationRunLanguageStats.

    Args:
        db: SQLAlchemy Session.
        event_log: Optional TranslationEventLog for logging events.
        run_id: The run ID to aggregate stats for.
        language_stats_map: Map of language_code -> TranslationRunLanguageStats to update.
    """
    records = (
        db.query(TranslationRecord)
        .filter(TranslationRecord.run_id == run_id)
        .all()
    )
    record_ids = [r.id for r in records]

    if not record_ids:
        logger.reason("No records for language stats aggregation", {"run_id": run_id})
        return

    lang_entries = (
        db.query(TranslationLanguage)
        .filter(TranslationLanguage.record_id.in_(record_ids))
        .all()
    )

    agg: dict[str, dict[str, int]] = defaultdict(
        lambda: {"total": 0, "translated": 0, "failed": 0, "skipped": 0}
    )
    for le in lang_entries:
        code = le.language_code
        agg[code]["total"] += 1
        if le.status in ("translated", "approved", "edited"):
            agg[code]["translated"] += 1
        elif le.status == "failed":
            agg[code]["failed"] += 1
        elif le.status == "skipped":
            agg[code]["skipped"] += 1

    total_chars = sum(len(le.translated_value or "") for le in lang_entries if le.translated_value)
    total_tokens = max(1, total_chars // 4)
    cost_per_token = 0.002 / 1000

    for lang_code, lang_stat in language_stats_map.items():
        data = agg.get(lang_code, {"total": 0, "translated": 0, "failed": 0, "skipped": 0})
        lang_stat.total_rows = data["total"]
        lang_stat.translated_rows = data["translated"]
        lang_stat.failed_rows = data["failed"]
        lang_stat.skipped_rows = data["skipped"]
        num_langs = len(language_stats_map)
        if num_langs > 0:
            lang_stat.token_count = total_tokens // num_langs
            lang_stat.estimated_cost = round((lang_stat.token_count / 1000) * cost_per_token, 6)

    db.flush()

    logger.reason("Language stats updated", {
        "run_id": run_id,
        "languages": list(language_stats_map.keys()),
        "total_tokens_est": total_tokens,
    })
# #endregion update_language_stats
