eight-hourr/backend/services/grok_client.py

from __future__ import annotations

import json
import logging
import re
from datetime import date
from typing import Any

import httpx

from config import settings
from services.description_formatter import format_professional_description
from services.distributor import DayEntry
from services.gitlab_client import GitLabTask

logger = logging.getLogger(__name__)

DEFAULT_PTS_TASK_TYPES = [
    "Implement",
    "Meeting",
    "Test",
    "Survey",
    "Bug Fix",
    "Support",
    "Trouble Shooting",
    "Document",
    "Design",
    "Code Review",
    "Training",
    "Project Management",
    "Manager Task",
    "POC",
    "Interview",
    "Investigating",
    "Misc",
]


class GrokClient:
    def __init__(self) -> None:
        self.api_key = settings.xai_api_key
        self.base_url = settings.xai_base_url.rstrip("/")
        self.model = settings.grok_model

    @property
    def enabled(self) -> bool:
        return bool(self.api_key)

    async def plan_pts_entries(
        self,
        tasks: list[GitLabTask],
        workdays: list[date],
        *,
        pts_project_name: str,
        pts_task_name: str,
        max_hours_per_day: float,
        hour_step: float,
    ) -> tuple[dict[str, list[DayEntry]], dict[str, Any]]:
        if not self.enabled:
            raise ValueError("XAI_API_KEY is not configured")

        if not tasks:
            raise ValueError("No GitLab tasks to analyze")
        if not workdays:
            raise ValueError("No workdays in selected range")

        prompt = self._build_prompt(
            tasks,
            workdays,
            pts_project_name=pts_project_name,
            pts_task_name=pts_task_name,
            max_hours_per_day=max_hours_per_day,
            hour_step=hour_step,
        )

        raw = await self._chat(prompt)
        parsed = self._parse_json(raw)
        plan = self._to_plan(parsed, tasks, workdays, hour_step, max_hours_per_day)
        analysis = {
            "model": self.model,
            "summary": parsed.get("summary", ""),
            "notes": parsed.get("notes", ""),
            "days": parsed.get("days", []),
        }
        return plan, analysis

    def _build_prompt(
        self,
        tasks: list[GitLabTask],
        workdays: list[date],
        *,
        pts_project_name: str,
        pts_task_name: str,
        max_hours_per_day: float,
        hour_step: float,
    ) -> str:
        task_payload = [
            {
                "id": task.id,
                "reference": task.reference,
                "title": task.title,
                "source": task.source,
                "labels": task.labels or [],
            }
            for task in tasks
        ]
        days_payload = [d.isoformat() for d in workdays]

        return f"""You are a senior software engineer preparing professional PTS timesheet entries for Supermicro.

Plan daily work logs based on the GitLab tasks below.

## PTS constraints
- Project (fixed): {pts_project_name}
- Task type (fixed): {pts_task_name}
- Plan only: description + take_hours
- Max {max_hours_per_day} hours per day
- Hours must be multiples of {hour_step}
- 1 to 3 entries per day
- Prioritize board bugs (source=board_bug), then milestone items (source=milestone)

## Description style (English only — mandatory)
Write concise, professional engineering work logs suitable for management review.

Rules:
- English only. No Chinese or other languages.
- Use past tense or present-perfect tense (e.g., "Investigated", "Implemented", "Refactored", "Validated").
- Start with a strong action verb.
- Include the GitLab reference in parentheses, e.g. (super-cloud/scc-flex#1319).
- Do NOT paste the raw GitLab title verbatim — rewrite as what was accomplished that day.
- One sentence, 12–25 words.
- Sound credible and specific, not generic.
- When the same gitlab_task_id appears on multiple days, each description MUST be a unique paraphrase:
  vary verbs and daily focus (investigated, reproduced, patched, validated, refactored, reviewed, etc.).
- Never reuse identical description text across entries.

Good examples:
- "Investigated Job Report snapshot regression after template deletion (super-cloud/scc-flex#1319)."
- "Refactored Jobs Schedule facade error responses for API consistency (super-cloud/scc-flex#1294)."
- "Validated firmware baseline history progress calculation under multi-target runs (super-cloud/scc-flex#1296)."

Bad examples:
- "Worked on bug" (too vague)
- Copy-paste of full GitLab title without rewrite
- Any non-English text

## Workdays
{json.dumps(days_payload)}

## GitLab tasks
{json.dumps(task_payload, indent=2)}

Return JSON only:
{{
  "summary": "Brief English summary of the overall plan",
  "notes": "English explanation of allocation logic",
  "days": [
    {{
      "date": "YYYY-MM-DD",
      "entries": [
        {{
          "gitlab_task_id": 123,
          "description": "Investigated Job Report snapshot regression after template deletion (super-cloud/scc-flex#1319).",
          "take_hours": 2.5,
          "reasoning": "English rationale for task selection and hours"
        }}
      ]
    }}
  ]
}}
"""

    async def _chat(self, prompt: str, *, temperature: float = 0.2) -> str:
        payload = {
            "model": self.model,
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "You are a professional technical writing assistant. "
                        "Output valid JSON only. All text fields must be in English."
                    ),
                },
                {"role": "user", "content": prompt},
            ],
            "temperature": temperature,
            "max_tokens": 4096,
        }

        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(
                f"{self.base_url}/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
            response.raise_for_status()
            data = response.json()

        content = data["choices"][0]["message"]["content"]
        if not content:
            raise ValueError("Grok returned empty content")
        return content

    def _parse_json(self, text: str) -> dict[str, Any]:
        cleaned = text.strip()
        if cleaned.startswith("```"):
            cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned)
            cleaned = re.sub(r"\s*```$", "", cleaned)

        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            match = re.search(r"\{.*\}", cleaned, re.DOTALL)
            if not match:
                raise ValueError("Grok response is not valid JSON") from None
            return json.loads(match.group())

    def _to_plan(
        self,
        parsed: dict[str, Any],
        tasks: list[GitLabTask],
        workdays: list[date],
        hour_step: float,
        max_hours_per_day: float,
    ) -> dict[str, list[DayEntry]]:
        valid_ids = {task.id for task in tasks}
        task_by_id = {task.id: task for task in tasks}
        allowed_days = {d.isoformat() for d in workdays}
        plan: dict[str, list[DayEntry]] = {d.isoformat(): [] for d in workdays}

        for day_block in parsed.get("days", []):
            day = str(day_block.get("date", ""))
            if day not in allowed_days:
                continue

            entries: list[DayEntry] = []
            total_hours = 0.0
            for item in day_block.get("entries", []):
                task_id = item.get("gitlab_task_id")
                if task_id not in valid_ids:
                    continue

                hours = float(item.get("take_hours", 0))
                hours = round(round(hours / hour_step) * hour_step, 1)
                if hours <= 0:
                    continue
                if total_hours + hours > max_hours_per_day:
                    hours = max(0.0, max_hours_per_day - total_hours)
                    hours = round(round(hours / hour_step) * hour_step, 1)
                if hours <= 0:
                    continue

                description = str(item.get("description", "")).strip()
                if not description or self._contains_cjk(description):
                    description = format_professional_description(task_by_id[task_id])

                entries.append(
                    DayEntry(
                        date=day,
                        description=description,
                        take_hours=hours,
                        gitlab_task_id=task_id,
                    )
                )
                total_hours += hours

            if entries:
                plan[day] = entries

        if not any(plan.values()):
            raise ValueError("Grok plan contained no valid entries")

        return plan

    def _contains_cjk(self, text: str) -> bool:
        return bool(re.search(r"[\u4e00-\u9fff\u3400-\u4dbf]", text))

    def _match_task_name(self, name: str, allowed: list[str], fallback: str) -> str:
        cleaned = str(name or "").strip()
        if cleaned in allowed:
            return cleaned
        lowered = cleaned.lower()
        for option in allowed:
            if lowered == option.lower():
                return option
        for option in allowed:
            if lowered in option.lower() or option.lower() in lowered:
                return option
        return fallback if fallback in allowed else allowed[0]

    async def classify_calendar_meetings(
        self,
        meetings: list[dict[str, Any]],
        *,
        pts_task_options: list[str],
        default_task_name: str,
    ) -> tuple[dict[str, dict[str, str]], dict[str, Any]]:
        if not meetings:
            return {}, {"source": "skipped", "reason": "no_meetings"}

        allowed = pts_task_options or DEFAULT_PTS_TASK_TYPES
        fallback = self._match_task_name(default_task_name, allowed, "Meeting")

        prompt = f"""You classify Microsoft Teams / Outlook calendar events into PTS timesheet task types for Supermicro engineers.

For EACH meeting below, pick the best PTS task type and write a professional English work-log description.

## Allowed PTS task types (pick exactly one per meeting)
{json.dumps(allowed, ensure_ascii=False)}

## Classification guide
- Standups, syncs, team meetings, 1:1s, general discussions -> Meeting
- Sprint planning, roadmap, status review, project sync -> Project Management
- Design review, architecture review, UI/UX review -> Design
- Code review, PR review -> Code Review
- Training, workshop, onboarding session -> Training
- Customer call, support escalation, troubleshooting call -> Support or Trouble Shooting
- Interview (candidate) -> Interview
- Demo, POC review -> POC
- Requirements gathering, user research -> Survey
- Incident/postmortem investigation -> Investigating
- Documentation walkthrough -> Document

## Description rules
- English only, one sentence, 12–25 words, past tense
- Mention Teams/meeting context naturally
- Do NOT change take_hours

## Meetings
{json.dumps(meetings, ensure_ascii=False, indent=2)}

Return JSON only:
{{
  "summary": "brief note on classification approach",
  "meetings": [
    {{
      "key": "YYYY-MM-DD#0",
      "pts_task_name": "Meeting",
      "description": "Attended sprint planning for SCC v3.9 release scope.",
      "reasoning": "sprint planning keyword"
    }}
  ]
}}
"""

        raw = await self._chat(prompt, temperature=0.3)
        data = self._parse_json(raw)
        classified: dict[str, dict[str, str]] = {}
        for item in data.get("meetings", []):
            key = str(item.get("key", "")).strip()
            if not key:
                continue
            description = str(item.get("description", "")).strip()
            if not description or self._contains_cjk(description):
                continue
            task_name = self._match_task_name(
                str(item.get("pts_task_name", "")),
                allowed,
                fallback,
            )
            classified[key] = {
                "pts_task_name": task_name,
                "description": description,
                "reasoning": str(item.get("reasoning", "")).strip(),
            }

        return classified, {
            "source": "grok_meeting_classify",
            "summary": data.get("summary", ""),
            "classified_count": len(classified),
        }

    async def diversify_descriptions(
        self,
        plan: dict[str, list[DayEntry]],
        tasks: list[GitLabTask],
    ) -> tuple[dict[str, list[DayEntry]], dict[str, Any]]:
        flat: list[dict[str, Any]] = []
        for day_key in sorted(plan.keys()):
            for idx, entry in enumerate(plan[day_key]):
                flat.append(
                    {
                        "key": f"{day_key}#{idx}",
                        "date": day_key,
                        "gitlab_task_id": entry.gitlab_task_id,
                        "current_description": entry.description,
                        "take_hours": entry.take_hours,
                    }
                )

        if len(flat) <= 1:
            return plan, {"source": "variation_skipped", "reason": "single_entry"}

        task_by_id = {t.id: t for t in tasks}
        task_context = [
            {
                "id": t.id,
                "reference": t.reference,
                "title": t.title,
                "source": t.source,
            }
            for t in tasks
        ]

        prompt = f"""Rewrite PTS timesheet descriptions for Supermicro engineering work logs.

For EACH entry below, produce a UNIQUE professional English description that:
- Keeps the same GitLab issue reference and gitlab_task_id when present
- Stays faithful to the task title / current description — describe plausible daily engineering work
- Uses different wording, verbs, and focus than every other entry in this batch
- Sounds like a real one-day work log (one sentence, 12–25 words), not copy-paste
- English only. Past tense or present-perfect. Start with a strong action verb.
- Do NOT change dates or take_hours

GitLab tasks (context):
{json.dumps(task_context, ensure_ascii=False, indent=2)}

Entries to rewrite:
{json.dumps(flat, ensure_ascii=False, indent=2)}

Return JSON only:
{{
  "summary": "short English note on variation approach",
  "entries": [
    {{"key": "YYYY-MM-DD#0", "description": "unique English work log"}}
  ]
}}
"""

        raw = await self._chat(prompt, temperature=0.65)
        data = self._parse_json(raw)
        rewritten = {
            item["key"]: str(item.get("description", "")).strip()
            for item in data.get("entries", [])
            if item.get("key")
        }
        if not rewritten:
            raise ValueError("Grok variation returned no entries")

        new_plan: dict[str, list[DayEntry]] = {}
        for day_key in sorted(plan.keys()):
            new_plan[day_key] = []
            for idx, entry in enumerate(plan[day_key]):
                key = f"{day_key}#{idx}"
                description = rewritten.get(key, entry.description)
                if (not description or self._contains_cjk(description)) and entry.gitlab_task_id in task_by_id:
                    description = format_professional_description(task_by_id[entry.gitlab_task_id])
                new_plan[day_key].append(
                    DayEntry(
                        date=day_key,
                        description=description,
                        take_hours=entry.take_hours,
                        gitlab_task_id=entry.gitlab_task_id,
                    )
                )

        return new_plan, {
            "source": "grok_variation",
            "summary": data.get("summary", ""),
            "rewritten_count": len(rewritten),
        }