eight-hourr/backend/services/planner.py

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,
        }