task_manager_agent/llm.py

"""
llm.py — Task Manager Agent
=============================
All Groq (llama-3.3-70b-versatile) calls:
  - extract tasks from emails
  - extract tasks from meeting notes
  - prioritize + enrich tasks
"""

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

from groq import Groq
from dotenv import load_dotenv
load_dotenv()

logger = logging.getLogger("TaskManagerAgent.LLM")

client = Groq(api_key=os.getenv("GROQ_API_KEY"))
MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
TODAY = datetime.now().strftime("%A, %B %d, %Y")


# ── helpers ───────────────────────────────────────────────────────────────────

def _chat(system: str, user: str, temperature: float = 0.3) -> str:
    resp = client.chat.completions.create(
        model=MODEL,
        temperature=temperature,
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": user},
        ],
    )
    return resp.choices[0].message.content.strip()


def _safe_json(text: str) -> Any:
    """Extract and parse first JSON block from LLM output."""
    # Strip markdown fences
    text = re.sub(r"```(?:json)?", "", text).strip().rstrip("`")
    # Find outermost array or object
    for start_char, end_char in [("[", "]"), ("{", "}")]:
        start = text.find(start_char)
        end = text.rfind(end_char)
        if start != -1 and end != -1:
            try:
                return json.loads(text[start : end + 1])
            except json.JSONDecodeError:
                pass
    logger.warning("Could not parse JSON from LLM response")
    return []


# ── email task extraction ─────────────────────────────────────────────────────

EXTRACT_EMAIL_SYSTEM = f"""You are an expert executive assistant AI.
Today is {TODAY}.

Your job is to read emails and extract ACTIONABLE tasks — things the user must DO.
Focus on:
- Direct requests or assignments ("Please send me…", "Can you prepare…", "I need you to…")
- Commitments the user made ("I'll get back to you", "Will share by Friday")
- Deadlines or follow-ups buried in email threads
- Approvals, decisions, or reviews required

Return ONLY a JSON array of task objects. Each task object must have:
{{
  "title":       "short imperative action (verb + object)",
  "description": "1-2 sentence context from the email",
  "source":      "email",
  "source_ref":  "<email subject>",
  "sender":      "<from address>",
  "deadline":    "YYYY-MM-DD or null",
  "deadline_confidence": "high|medium|low",
  "estimated_minutes": <number or null>,
  "tags":        ["tag1", "tag2"],
  "priority_raw": "urgent|high|medium|low"
}}

Return [] if no actionable tasks are found.
Do NOT include tasks that are just FYI, news, or promotions.
Do NOT duplicate tasks already in the existing list."""


def extract_tasks_from_emails(
    emails: list[dict],
    existing_tasks: list[dict],
) -> list[dict]:
    if not emails:
        return []

    existing_titles = [t.get("title", "") for t in existing_tasks]
    existing_context = "\n".join(f"- {t}" for t in existing_titles[:30]) if existing_titles else "None"

    email_blocks = []
    for i, e in enumerate(emails, 1):
        email_blocks.append(
            f"=== EMAIL {i} ===\n"
            f"From: {e['sender']}\n"
            f"Subject: {e['subject']}\n"
            f"Date: {e['date']}\n"
            f"Body:\n{e['body']}\n"
        )

    user_prompt = (
        f"EXISTING TASKS (do not duplicate):\n{existing_context}\n\n"
        f"EMAILS TO ANALYSE:\n{''.join(email_blocks)}"
    )

    try:
        raw = _chat(EXTRACT_EMAIL_SYSTEM, user_prompt)
        tasks = _safe_json(raw)
        if not isinstance(tasks, list):
            tasks = []
        logger.info(f"Email extraction: {len(tasks)} tasks found")
        return tasks
    except Exception as e:
        logger.error(f"Email task extraction failed: {e}", exc_info=True)
        return []


# ── meeting task extraction ───────────────────────────────────────────────────

EXTRACT_MEETING_SYSTEM = f"""You are an expert executive assistant AI.
Today is {TODAY}.

Your job is to extract ACTIONABLE tasks from meeting descriptions and calendar events.
Look for:
- Action items mentioned in descriptions ("Action: …", "TODO:", "Next steps:")
- Implied follow-ups (prep materials, send recap, schedule next meeting)
- Deadlines or deliverables tied to the meeting

Return ONLY a JSON array of task objects. Each must have:
{{
  "title":       "short imperative action",
  "description": "context from the meeting",
  "source":      "meeting",
  "source_ref":  "<meeting title>",
  "sender":      null,
  "deadline":    "YYYY-MM-DD or null",
  "deadline_confidence": "high|medium|low",
  "estimated_minutes": <number or null>,
  "tags":        ["tag1", "tag2"],
  "priority_raw": "urgent|high|medium|low"
}}

Return [] if nothing actionable found.
Do NOT duplicate tasks already in the existing list."""


def extract_tasks_from_meeting_notes(
    meetings: list[dict],
    existing_tasks: list[dict],
) -> list[dict]:
    if not meetings:
        return []

    existing_titles = [t.get("title", "") for t in existing_tasks]
    existing_context = "\n".join(f"- {t}" for t in existing_titles[:30]) if existing_titles else "None"

    meeting_blocks = []
    for i, m in enumerate(meetings, 1):
        attendees = ", ".join(m.get("attendees", [])[:5]) or "N/A"
        meeting_blocks.append(
            f"=== MEETING {i} ===\n"
            f"Title: {m['title']}\n"
            f"When: {m['start']} → {m['end']}\n"
            f"Attendees: {attendees}\n"
            f"Description:\n{m.get('description', '(none)')}\n"
        )

    user_prompt = (
        f"EXISTING TASKS (do not duplicate):\n{existing_context}\n\n"
        f"MEETINGS TO ANALYSE:\n{''.join(meeting_blocks)}"
    )

    try:
        raw = _chat(EXTRACT_MEETING_SYSTEM, user_prompt)
        tasks = _safe_json(raw)
        if not isinstance(tasks, list):
            tasks = []
        logger.info(f"Meeting extraction: {len(tasks)} tasks found")
        return tasks
    except Exception as e:
        logger.error(f"Meeting task extraction failed: {e}", exc_info=True)
        return []


# ── prioritization ────────────────────────────────────────────────────────────

PRIORITIZE_SYSTEM = f"""You are a world-class productivity coach and AI assistant.
Today is {TODAY}.

You will receive a list of raw tasks. Your job is to:
1. Score each task's PRIORITY on a 1-10 scale using these criteria:
   - Urgency (deadline proximity)
   - Impact (business / personal significance)
   - Effort required (lower effort = slightly higher priority, all else equal)
   - Dependencies (if others are blocked, bump up)

2. Assign a CATEGORY from: [work, personal, admin, communication, research, finance, health, other]

3. Suggest a DUE DATE if none exists (or confirm/adjust if one does).

4. Write a SHORT "why_urgent" note (1 sentence) for tasks scored 7+.

Return ONLY a JSON array. Each object must include ALL original fields PLUS:
{{
  ...original fields...,
  "priority_score": <1-10>,
  "priority_label": "critical|high|medium|low",
  "category": "<category>",
  "suggested_due_date": "YYYY-MM-DD or null",
  "why_urgent": "one sentence or null",
  "order": <integer starting from 1, lowest = highest priority>
}}

Sort the array by priority_score descending (order 1 = most important)."""


def prioritize_tasks(
    new_tasks: list[dict],
    all_context_tasks: list[dict],
) -> list[dict]:
    if not new_tasks:
        return []

    # Give the LLM context about the full task landscape
    context_summary = []
    for t in all_context_tasks[:20]:
        context_summary.append(
            f"- [{t.get('priority_raw', '?')}] {t.get('title', '?')} "
            f"(deadline: {t.get('deadline', 'none')})"
        )
    context_str = "\n".join(context_summary) if context_summary else "No existing tasks"

    user_prompt = (
        f"FULL TASK CONTEXT (existing + new):\n{context_str}\n\n"
        f"TASKS TO PRIORITIZE:\n{json.dumps(new_tasks, indent=2)}"
    )

    try:
        raw = _chat(PRIORITIZE_SYSTEM, user_prompt, temperature=0.2)
        tasks = _safe_json(raw)
        if not isinstance(tasks, list) or not tasks:
            # Fallback: return original tasks with default priority
            logger.warning("Prioritization returned empty — using defaults")
            return _apply_default_priority(new_tasks)
        logger.info(f"Prioritized {len(tasks)} tasks")
        return tasks
    except Exception as e:
        logger.error(f"Prioritization failed: {e}", exc_info=True)
        return _apply_default_priority(new_tasks)


def _apply_default_priority(tasks: list[dict]) -> list[dict]:
    priority_map = {"urgent": 9, "high": 7, "medium": 5, "low": 3}
    for i, t in enumerate(tasks):
        score = priority_map.get(t.get("priority_raw", "medium"), 5)
        t["priority_score"] = score
        t["priority_label"] = t.get("priority_raw", "medium")
        t["category"] = "work"
        t["suggested_due_date"] = t.get("deadline")
        t["why_urgent"] = None
        t["order"] = i + 1
    return sorted(tasks, key=lambda x: x["priority_score"], reverse=True)


# ── daily digest summary ──────────────────────────────────────────────────────

DIGEST_SYSTEM = f"""You are a sharp, concise executive assistant.
Today is {TODAY}.

Write a crisp task digest. Structure:
1. One-line headline ("You have N high-priority tasks today")
2. Top 3 critical tasks with a one-line action each
3. Quick summary of remaining tasks grouped by category
4. One motivational closing line

Keep the total under 250 words. Use bullet points sparingly — prefer clean paragraphs.
Do NOT use emoji."""


def generate_digest_summary(tasks: list[dict]) -> str:
    if not tasks:
        return "No new tasks extracted today. Your slate is clean."
    try:
        task_list = json.dumps(
            [{"title": t.get("title"), "priority_score": t.get("priority_score"),
              "deadline": t.get("deadline") or t.get("suggested_due_date"),
              "category": t.get("category"), "why_urgent": t.get("why_urgent")}
             for t in tasks[:15]],
            indent=2
        )
        return _chat(DIGEST_SYSTEM, f"TASKS:\n{task_list}", temperature=0.5)
    except Exception as e:
        logger.error(f"Digest generation failed: {e}")
        return f"{len(tasks)} new tasks extracted and prioritized."