core/ai_engine.py

"""
core/ai_engine.py
All Groq API calls — prompts taken verbatim from the tested Colab notebook.
Covers: task parsing, scheduling, journaling Q&A, context synthesis.
"""

import json
import re
import os
from datetime import datetime, date
from copy import deepcopy

from groq import Groq

GROQ_MODEL = "llama-3.3-70b-versatile"
_client: Groq = None


def init_groq(api_key: str = None):
    global _client
    key = api_key or os.environ.get("GROQ_API_KEY", "")
    if not key:
        raise ValueError(
            "GROQ_API_KEY is not set. "
            "Add it in HuggingFace Space → Settings → Repository secrets."
        )
    _client = Groq(api_key=key)


def _groq() -> Groq:
    if _client is None:
        init_groq()
    return _client


# ── Shared util ───────────────────────────────────────────────────────────────

def safe_json_parse(text: str):
    """Parse JSON, stripping markdown fences if present. Returns None on failure."""
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        cleaned = re.sub(r'^```(?:json)?\s*|\s*```$', '', text, flags=re.MULTILINE).strip()
        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            m = re.search(r'\{[\s\S]*\}', cleaned)
            if m:
                try:
                    return json.loads(m.group())
                except Exception:
                    pass
    return None


# ── Module 1: Task Capture ────────────────────────────────────────────────────

TASK_CAPTURE_PROMPT = """You are a task classification assistant for a productivity app called The Second Brain.

Your job is to take a user's raw task description and return a structured JSON object.

Classify the task across these dimensions:

1. title (string): A clean, concise, action-oriented task title. Fix grammar.

2. life_area (string): Choose ONE from: Work, Health, Learning, Finance, Personal, Family, Other
   - Work: job, meetings, deadlines, clients, projects
   - Health: exercise, medical, diet, mental health
   - Learning: courses, books, skills, studying
   - Finance: bills, payments, budgeting, taxes, investments
   - Personal: hobbies, errands, home maintenance
   - Family: tasks involving family members
   - Other: does not fit any category

3. urgency (string): Choose ONE from:
   - Habit: recurring or routine task
   - Urgent: hard deadline or time pressure
   - Not Urgent: no specific deadline

4. importance (string): Choose ONE from:
   - Move the Needle: very high impact
   - Important: meaningful, should be done
   - Not Important: low real impact

5. state_of_mind (string): Choose ONE from:
   - Quick: 5-10 mins, very low focus
   - Easy: 10-20 mins, low focus
   - Flow: deep concentration needed
   - Personal: life admin, little goal impact

6. time_estimate (integer): Realistic minutes to complete.

7. deadline_date (string or null): If user says "by Friday", "due March 1", "before end of month", "deadline X" — extract as YYYY-MM-DD. Today is {TODAY}. Return null if no deadline mentioned.

8. clarifications_needed (array of strings):
   If NOT confident about a dimension, add a short specific question.
   If everything is clear, return []

STRICT RULES:
- Return ONLY valid JSON. No markdown, no explanation.
- Never guess if uncertain — ask a clarification question instead.
- Always return all 8 fields including deadline_date (null if none).

Example: {"title": "Finish project proposal", "life_area": "Work", "urgency": "Urgent", "importance": "Move the Needle", "state_of_mind": "Flow", "time_estimate": 90, "deadline_date": null, "clarifications_needed": []}"""


def parse_task_with_groq(raw_text: str, user_context: dict = None,
                         user_goals: list = None, life_areas: list = None) -> dict:
    """Parse raw task text into structured dimensions using Groq."""
    # Build context hint from AI memory + goals
    context_hint = ""
    if user_context and user_context.get("learned_patterns", {}).get("notes"):
        notes = user_context["learned_patterns"]["notes"]
        context_hint += f"\n\nUser context notes (use to inform classification): {'; '.join(notes[-3:])}"
    if user_goals:
        context_hint += f"\nUser goals: {'; '.join(user_goals[:5])}"
    if life_areas:
        context_hint += f"\nUser's life areas: {', '.join(life_areas)}"

    _prompt = TASK_CAPTURE_PROMPT.replace("{TODAY}", str(date.today())) + context_hint
    response = _groq().chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": _prompt},
            {"role": "user",   "content": f"Parse this task: {raw_text}"}
        ],
        max_tokens=512,
        temperature=0.1,
    )

    result = safe_json_parse(response.choices[0].message.content.strip())
    if result is None:
        result = {
            "title": raw_text,
            "life_area": None, "urgency": None, "importance": None,
            "state_of_mind": None, "time_estimate": None,
            "deadline_date": None,
            "clarifications_needed": [
                "Could you give more details about this task?",
                "Which area of your life does this belong to?",
                "Is this urgent or flexible?"
            ]
        }
    return result


# ── Module 2: Scheduling ──────────────────────────────────────────────────────

SCHEDULING_SYSTEM_PROMPT = """You are an intelligent daily scheduler for a productivity app called The Second Brain.

You receive a USER CONTEXT (preferences + learned patterns), a TASK LIST, and a SCHEDULING PROMPT.
Return a time-blocked schedule as a JSON object.

SCHEDULING RULES:
- Respect wake_time and sleep_time from context
- Place Flow tasks during the user's peak focus time
- If avg_task_overrun_pct > 0, add buffer proportionally to time estimates
- If flow_batch_capable is true, group Flow tasks; otherwise space them out
- Place Quick and Easy tasks around transitions and low-energy windows
- Place Personal/Habit tasks at day boundaries (start or end of day)
- Urgent tasks are scheduled before Not Urgent ones
- Move the Needle tasks get the best time slots
- Add 5-10 min breaks between tasks
- Respect any fixed commitments mentioned in the scheduling prompt
- Do NOT schedule past sleep_time
- If tasks won't realistically fit, put them in deferred_tasks
- If context is minimal (new user), use sensible defaults

RETURN FORMAT (JSON only, no markdown):
{
  "schedule_date": "YYYY-MM-DD",
  "scheduled_tasks": [
    {
      "task_id": "(id from input or index)",
      "title": "...",
      "life_area": "...",
      "start_time": "HH:MM",
      "end_time": "HH:MM",
      "duration_minutes": 60,
      "state_of_mind": "...",
      "scheduling_reason": "1-sentence explanation"
    }
  ],
  "deferred_tasks": [{"task_id": "...", "title": "...", "reason": "..."}],
  "day_summary": "2-3 sentences on day structure and reasoning",
  "warnings": ["any concerns e.g. day overloaded"]
}"""


def generate_schedule(context: dict, tasks: list, scheduling_prompt: str,
                      goals: list = None, schedule_date: str = None) -> dict:
    if not schedule_date:
        schedule_date = str(date.today())

    goals_section = ""
    if goals:
        goals_section = "\nUSER GOALS:\n" + "\n".join(f"- {g}" for g in goals)

    user_message = f"""Schedule Date: {schedule_date}

USER CONTEXT:
{json.dumps(context, indent=2)}
{goals_section}
TASKS TO SCHEDULE ({len(tasks)} tasks):
{json.dumps(tasks, indent=2)}

USER SCHEDULING PROMPT:
{scheduling_prompt}

Generate the optimal schedule."""

    response = _groq().chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": SCHEDULING_SYSTEM_PROMPT},
            {"role": "user",   "content": user_message}
        ],
        max_tokens=2048,
        temperature=0.2,
    )

    result = safe_json_parse(response.choices[0].message.content.strip())
    if result is None:
        result = {
            "error": "Could not parse schedule response.",
            "raw": response.choices[0].message.content
        }

    result["schedule_date"] = schedule_date
    result["generated_at"] = datetime.now().isoformat()
    return result


# ── Module 3: Journaling ──────────────────────────────────────────────────────

JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for a productivity app called The Second Brain.

You receive the user's context, today's schedule with completion status, and the conversation so far.
Your job: decide what targeted question to ask NEXT.

FOCUS AREAS (cover what's most relevant, don't ask all):
- Tasks not completed — why? wrong time? too tired? overestimated?
- Tasks that took much longer than estimated
- Energy levels — when were they sharp vs drained?
- Whether Flow tasks were placed well or hard to start
- Whether the day felt balanced or overloaded
- Patterns the user noticed about themselves

RULES:
- Ask ONE question at a time. Short and specific.
- Build on previous answers — don't repeat covered ground.
- After 4-6 good exchanges, signal completion.
- Keep tone warm and efficient — 2-minute check-in, not therapy.

RETURN FORMAT (JSON only):
{"question": "Your next question", "question_focus": "what aspect this targets", "session_complete": false}
OR when done:
{"question": null, "question_focus": null, "session_complete": true}"""


SYNTHESIS_PROMPT = """You are a pattern recognition engine for a productivity app called The Second Brain.

You have a completed journaling conversation. Extract learnings and return the UPDATED user context JSON.

UPDATE these fields in learned_patterns based on conversation evidence:
- productive_times: when user felt sharp/focused
- low_energy_times: when they felt drained or skipped tasks
- avg_task_overrun_pct: recalculate from actual vs estimated times mentioned
- flow_batch_capable: update if user gave clear evidence
- best_life_areas_morning: what they completed well before noon
- common_skipped_task_types: patterns in what gets consistently skipped
- notes: append 1-2 new insight notes (keep existing ones)

ALWAYS UPDATE:
- scheduling_feedback.total_days_scheduled: +1
- scheduling_feedback.avg_completion_rate: rolling average
- scheduling_feedback.last_7_day_completion_rates: append today, keep last 7
- history_summary: append brief today summary, keep last 14
- last_updated: now
- version: +1

RULES:
- Return ONLY the complete updated context JSON. Nothing else.
- Never remove existing patterns — only update or append.
- Be conservative — only update if there is clear evidence in the conversation."""


def build_opening_question(context: dict, tasks_today: list) -> dict:
    """Generate the first journal question based on task completion at a glance."""
    total     = len(tasks_today)
    completed = sum(1 for t in tasks_today if t.get("completed", False))
    incomplete = [t for t in tasks_today if not t.get("completed", False)]

    if total == 0:
        q = "It looks like you didn't have any tasks scheduled today — was that intentional or did things go sideways?"
    elif completed == 0:
        q = f"None of today's {total} tasks got marked complete — was the day unexpectedly derailed, or did the plan just not fit how your day went?"
    elif completed == total:
        q = f"You completed all {total} tasks today — great day! Did the schedule feel natural, or were you pushing through?"
    elif len(incomplete) == 1:
        q = f'You got almost everything done — the one task left was "{incomplete[0]["title"]}". What got in the way?'
    else:
        rate   = round(completed / total * 100)
        titles = ", ".join(f'"{t["title"]}"' for t in incomplete[:2])
        q = f"You completed {completed}/{total} tasks ({rate}%). Tasks like {titles} didn't get done — was that time, energy, or something else?"

    return {"question": q, "question_focus": "completion_overview", "session_complete": False}


def get_next_journal_question(context: dict, tasks_today: list,
                              conversation_history: list) -> dict:
    user_message = f"""USER CONTEXT:
{json.dumps(context, indent=2)}

TODAY'S SCHEDULE (with completion):
{json.dumps(tasks_today, indent=2)}

CONVERSATION SO FAR ({len(conversation_history)} exchanges):
{json.dumps(conversation_history, indent=2)}

What should I ask next? Return session_complete: true if enough has been covered."""

    response = _groq().chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": JOURNAL_QUESTION_PROMPT},
            {"role": "user",   "content": user_message}
        ],
        max_tokens=256,
        temperature=0.3,
    )

    result = safe_json_parse(response.choices[0].message.content.strip())
    if result is None:
        result = {
            "question": response.choices[0].message.content.strip(),
            "question_focus": "general",
            "session_complete": False
        }
    return result


def synthesize_journal(context: dict, tasks_today: list,
                       conversation_history: list) -> dict:
    """Synthesize conversation into updated context. Fallback to manual stats update if LLM fails."""
    total     = len(tasks_today)
    completed = sum(1 for t in tasks_today if t.get("completed", False))
    completion_rate = round(completed / total, 2) if total > 0 else 0.0

    user_message = f"""USER CONTEXT (current):
{json.dumps(context, indent=2)}

TODAY'S SCHEDULE + COMPLETION:
{json.dumps(tasks_today, indent=2)}

Today's completion rate: {completion_rate} ({completed}/{total})

FULL JOURNALING CONVERSATION:
{json.dumps(conversation_history, indent=2)}

Return the complete updated context JSON."""

    response = _groq().chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": SYNTHESIS_PROMPT},
            {"role": "user",   "content": user_message}
        ],
        max_tokens=2048,
        temperature=0.1,
    )

    updated = safe_json_parse(response.choices[0].message.content.strip())

    if updated is None:
        # Fallback: update stats manually if synthesis fails
        updated = deepcopy(context)
        updated["last_updated"] = datetime.now().isoformat()
        updated["version"]      = context.get("version", 1) + 1
        sf    = updated.setdefault("scheduling_feedback", {})
        sf["total_days_scheduled"] = sf.get("total_days_scheduled", 0) + 1
        rates = sf.get("last_7_day_completion_rates", [])
        rates.append(completion_rate)
        sf["last_7_day_completion_rates"] = rates[-7:]
        sf["avg_completion_rate"] = round(sum(rates) / len(rates), 2)

    return updated


# ── Context helpers ───────────────────────────────────────────────────────────

def create_blank_context(user_id, preferences: dict = None) -> dict:
    prefs = preferences or {}
    return {
        "user_id": user_id,
        "created_at":   datetime.now().isoformat(),
        "last_updated": datetime.now().isoformat(),
        "version": 1,
        "preferences": {
            "wake_time":              prefs.get("wake_time", "08:00"),
            "sleep_time":             prefs.get("sleep_time", "23:00"),
            "focus_peak":             prefs.get("focus_peak", "Morning"),
            "break_duration_minutes": 10,
            "max_flow_block_minutes": 90,
        },
        "learned_patterns": {
            "productive_times":          [],
            "low_energy_times":          [],
            "avg_task_overrun_pct":      0,
            "flow_batch_capable":        None,
            "best_life_areas_morning":   [],
            "habit_completion_rate":     {},
            "common_skipped_task_types": [],
            "notes":                     []
        },
        "history_summary": [],
        "scheduling_feedback": {
            "total_days_scheduled":        0,
            "avg_completion_rate":         0.0,
            "last_7_day_completion_rates": []
        }
    }