from __future__ import annotations import logging from datetime import date from typing import Any, Optional from config import settings from services.distributor import ( DayEntry, build_plan_for_dates, build_plan_for_day_list, entries_from_last_state, ) from services.gitlab_client import GitLabClient, GitLabTask from services.grok_client import GrokClient from services.holidays_util import iter_workdays from services.meeting_planner import merge_meetings_into_plan from services.state_store import get_last_entries logger = logging.getLogger(__name__) class Planner: def __init__(self) -> None: self.gitlab = GitLabClient() self.grok = GrokClient() async def build_plan( self, start: date, end: date, *, use_llm: Optional[bool] = None, use_last_if_empty: bool = True, ) -> dict[str, Any]: workdays = iter_workdays(start, end) return await self._build_plan_for_days( workdays, use_llm=use_llm, use_last_if_empty=use_last_if_empty, rule_builder=lambda tasks, _: build_plan_for_dates(tasks, start, end), ) async def build_plan_for_days( self, days: list[date], *, use_llm: Optional[bool] = None, use_last_if_empty: bool = True, ) -> dict[str, Any]: sorted_days = sorted(set(days)) return await self._build_plan_for_days( sorted_days, use_llm=use_llm, use_last_if_empty=use_last_if_empty, rule_builder=lambda tasks, workdays: build_plan_for_day_list(tasks, workdays), ) async def _build_plan_for_days( self, workdays: list[date], *, use_llm: Optional[bool], use_last_if_empty: bool, rule_builder, ) -> dict[str, Any]: gitlab_tasks = await self.gitlab.get_assigned_tasks() if gitlab_tasks: use_llm = settings.use_llm if use_llm is None else use_llm if use_llm and self.grok.enabled: try: plan, analysis = await self.grok.plan_pts_entries( gitlab_tasks, workdays, pts_project_name=settings.pts_project_name, pts_task_name=settings.pts_default_task_name, max_hours_per_day=settings.max_hours_per_day, hour_step=settings.hour_step, ) plan, analysis = await self._apply_variation( plan, gitlab_tasks, analysis, source="llm", ) return await self._finalize_plan( plan, source="llm", gitlab_tasks=gitlab_tasks, workdays=workdays, analysis=analysis, ) except Exception as exc: # noqa: BLE001 logger.warning("LLM planning failed, fallback to rules: %s", exc) plan = rule_builder(gitlab_tasks, workdays) plan, analysis = await self._apply_variation( plan, gitlab_tasks, {"fallback_reason": "rule_based_distribution"}, source="rule_based", ) return await self._finalize_plan( plan, source="rule_based", gitlab_tasks=gitlab_tasks, workdays=workdays, analysis=analysis, ) if use_last_if_empty: plan = entries_from_last_state(get_last_entries(), workdays) plan, analysis = await self._apply_variation( plan, [], {}, source="last_entries", ) return await self._finalize_plan( plan, source="last_entries", gitlab_tasks=[], workdays=workdays, analysis=analysis, ) return await self._finalize_plan( {d.isoformat(): [] for d in workdays}, source="empty", gitlab_tasks=[], workdays=workdays, analysis={}, ) def _has_duplicate_descriptions(self, plan: dict[str, list[DayEntry]]) -> bool: descriptions = [entry.description for entries in plan.values() for entry in entries] return len(descriptions) != len(set(descriptions)) async def _apply_variation( self, plan: dict[str, list[DayEntry]], tasks: list[GitLabTask], analysis: dict[str, Any], *, source: str, ) -> tuple[dict[str, list[DayEntry]], dict[str, Any]]: if not settings.use_description_variation or not self.grok.enabled: return plan, analysis entry_count = sum(len(entries) for entries in plan.values()) if entry_count <= 1: return plan, {**analysis, "variation": {"source": "variation_skipped", "reason": "single_entry"}} if source == "llm" and not self._has_duplicate_descriptions(plan): return plan, {**analysis, "variation": {"source": "variation_skipped", "reason": "llm_unique"}} try: plan, variation = await self.grok.diversify_descriptions(plan, tasks) return plan, {**analysis, "variation": variation} except Exception as exc: # noqa: BLE001 logger.warning("Description variation failed: %s", exc) return plan, {**analysis, "variation": {"source": "variation_failed", "error": str(exc)}} async def _finalize_plan( self, plan: dict[str, list[DayEntry]], *, source: str, gitlab_tasks: list[GitLabTask], workdays: list[date], analysis: dict[str, Any], ) -> dict[str, Any]: workday_keys = [d.isoformat() for d in workdays] plan, meeting_meta = await merge_meetings_into_plan(plan, workday_keys) if meeting_meta: analysis = {**analysis, "meetings": meeting_meta} return self._wrap_plan( plan, source=source, gitlab_tasks=gitlab_tasks, workdays=workdays, analysis=analysis, ) def _wrap_plan( self, plan: dict[str, list[DayEntry]], *, source: str, gitlab_tasks: list[GitLabTask], workdays: list[date], analysis: dict[str, Any], ) -> dict[str, Any]: return { "source": source, "gitlab_task_count": len(gitlab_tasks), "workdays": [d.isoformat() for d in workdays], "llm_enabled": self.grok.enabled, "analysis": analysis, "plan": { day: [ { "description": entry.description, "take_hours": entry.take_hours, "gitlab_task_id": entry.gitlab_task_id, "pts_task_name": entry.pts_task_name, } for entry in entries ] for day, entries in plan.items() }, "_plan_entries": plan, }