core/ai_engine.py

"""
core/ai_engine.py
All Groq API calls — task parsing, intelligent RAG-style scheduling, journaling.
"""

import json
import re
import os
from datetime import datetime, date, timedelta
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 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):
    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 Second Brain.

Take the user's raw task description and return a structured JSON object.

Dimensions:
1. title (string): Clean, action-oriented task title.
2. life_area (string): ONE of: Work, Health, Learning, Finance, Personal, Family, Other
3. urgency (string): ONE of: Habit | Urgent | Not Urgent
4. importance (string): ONE of: Move the Needle | Important | Not Important
5. state_of_mind (string): ONE of: Flow | Easy | Quick | Personal
6. time_estimate (integer): Realistic minutes to complete.
7. clarifications_needed (array): Short specific questions if uncertain about any dimension. Return [] if confident.

RETURN: Valid JSON only. No markdown, no prose. All 7 fields always present."""


def parse_task_with_groq(raw_text: str, user_context: dict = None,
                         user_goals: list = None, life_areas: list = None) -> dict:
    context_hint = ""
    if user_context and user_context.get("learned_patterns", {}).get("notes"):
        notes = user_context["learned_patterns"]["notes"]
        context_hint += f"\n\nUser patterns: {'; '.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)}"

    response = _groq().chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": TASK_CAPTURE_PROMPT + context_hint},
            {"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,
            "clarifications_needed": ["Could you give more details? Which area, how urgent, how long?"]
        }
    return result


# ── Module 2: Intelligent RAG-style Task Scheduler ────────────────────────────

SMART_SCHEDULER_PROMPT = """You are an intelligent task scheduler for Second Brain.

You assign UNSCHEDULED TASKS to specific future dates, like a smart personal assistant who understands the
user's rhythms, goals, energy patterns, and the current time.

REASONING PROCESS:
1. Read the user request carefully — honour it above all else
   - "clear my day" / "nothing today" / "free today" = assign NOTHING to today
   - "schedule for tomorrow" = assign to tomorrow
   - "this week" = spread across next 5 days
   - No explicit date = use next 1-7 days intelligently
2. Never assign to a date/time in the past (current datetime is given)
3. If current time is past 18:00, treat today as unavailable unless user explicitly asks
4. Prioritise by: deadline proximity > urgency > importance > goal alignment
5. Match tasks to days by state_of_mind:
   - Flow = peak days (Mon-Thu mornings if peak=Morning)
   - Quick/Easy = any day, fill gaps
   - Habit = today or tomorrow
6. Spread load — don't stack everything on one day
7. Tasks with deadlines must land BEFORE that deadline

RETURN FORMAT (JSON only, no markdown):
{
  "assignments": [
    {
      "task_id": 123,
      "title": "Task title",
      "assigned_date": "YYYY-MM-DD",
      "reasoning": "1 sentence why this date"
    }
  ],
  "skipped": [
    {
      "task_id": 456,
      "title": "Task title",
      "reason": "why not assigned"
    }
  ],
  "summary": "2-3 sentence plain-English explanation of what was scheduled and why"
}

NEVER assign to a past date. NEVER ignore an explicit user instruction about when to schedule."""


def smart_schedule_tasks(
    tasks: list,
    user_context: dict,
    user_goals: list,
    scheduling_prompt: str,
    current_dt: datetime = None,
) -> dict:
    """
    RAG-style scheduler: reads context + goals + patterns + current time + user request,
    assigns each unscheduled task to a specific future date.
    """
    if current_dt is None:
        current_dt = datetime.now()

    today    = current_dt.date()
    tomorrow = today + timedelta(days=1)
    next_7   = [(today + timedelta(days=i)).isoformat() for i in range(8)]

    prefs    = user_context.get("preferences", {})
    patterns = user_context.get("learned_patterns", {})
    feedback = user_context.get("scheduling_feedback", {})

    context_block = f"""CURRENT DATE/TIME: {current_dt.strftime('%Y-%m-%d %H:%M')} ({current_dt.strftime('%A')})
TODAY: {today.isoformat()}  |  TOMORROW: {tomorrow.isoformat()}
NEXT 7 DAYS: {', '.join(next_7)}

USER PREFERENCES:
- Wake: {prefs.get('wake_time', '08:00')} | Sleep: {prefs.get('sleep_time', '23:00')}
- Peak focus: {prefs.get('focus_peak', 'Morning')}
- Max flow block: {prefs.get('max_flow_block_minutes', 90)} min

LEARNED PATTERNS:
- Productive times: {patterns.get('productive_times', 'unknown')}
- Low energy times: {patterns.get('low_energy_times', 'unknown')}
- Avg task overrun: {patterns.get('avg_task_overrun_pct', 0)}%
- Flow batching: {patterns.get('flow_batch_capable', 'unknown')}
- Commonly skipped: {patterns.get('common_skipped_task_types', [])}
- Notes: {'; '.join(patterns.get('notes', [])[-3:])}

HISTORY:
- Days tracked: {feedback.get('total_days_scheduled', 0)}
- Avg completion: {round(feedback.get('avg_completion_rate', 0) * 100)}%

GOALS:
{chr(10).join(f'- {g}' for g in (user_goals or [])) or '(none set)'}"""

    tasks_block = json.dumps([{
        "task_id":       t.get("id", t.get("task_id")),
        "title":         t.get("title"),
        "life_area":     t.get("life_area"),
        "urgency":       t.get("urgency"),
        "importance":    t.get("importance"),
        "state_of_mind": t.get("state_of_mind"),
        "time_estimate": t.get("time_estimate"),
        "deadline_date": t.get("deadline_date") or "none",
    } for t in tasks], indent=2)

    user_message = f"""{context_block}

UNSCHEDULED TASKS ({len(tasks)} tasks):
{tasks_block}

USER REQUEST: "{scheduling_prompt}"

Assign each task to the best date. Follow the user request precisely."""

    try:
        response = _groq().chat.completions.create(
            model=GROQ_MODEL,
            messages=[
                {"role": "system", "content": SMART_SCHEDULER_PROMPT},
                {"role": "user",   "content": user_message}
            ],
            max_tokens=2048,
            temperature=0.15,
        )
        result = safe_json_parse(response.choices[0].message.content.strip())
    except Exception as e:
        result = None

    if result is None:
        return {
            "assignments": [],
            "skipped": [{"task_id": t.get("id", t.get("task_id")), "title": t.get("title"),
                         "reason": "AI scheduling failed"} for t in tasks],
            "summary": "Scheduling failed — please try again or rephrase your request."
        }

    # Safety pass: strip any assignments set in the past
    safe_assignments = []
    for a in result.get("assignments", []):
        try:
            assigned = date.fromisoformat(a["assigned_date"])
            if assigned >= today:
                safe_assignments.append(a)
            else:
                result.setdefault("skipped", []).append({
                    "task_id": a.get("task_id"),
                    "title":   a.get("title", ""),
                    "reason":  f"AI tried to assign to past date {a['assigned_date']} — blocked"
                })
        except (ValueError, KeyError):
            pass

    result["assignments"] = safe_assignments
    return result


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

JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for Second Brain.

Given user context, today's tasks, and the conversation so far — decide what to ask next.

FOCUS: completion reasons, energy patterns, time estimate accuracy, flow placement, overall balance.

RULES: ONE question at a time. Build on prior answers. After 5-7 exchanges, signal session_complete.
Warm, efficient tone — 2-minute check-in.

RETURN (JSON only):
{"question": "...", "question_focus": "...", "session_complete": false}
OR: {"question": null, "question_focus": null, "session_complete": true}"""


SYNTHESIS_PROMPT = """You are a pattern recognition engine for Second Brain.

Given a completed journaling conversation, return the UPDATED user context JSON.

UPDATE learned_patterns based on evidence: productive_times, low_energy_times, avg_task_overrun_pct,
flow_batch_capable, best_life_areas_morning, common_skipped_task_types.
Append 1-2 new insight notes (never remove existing).

ALWAYS UPDATE: scheduling_feedback (total_days_scheduled +1, rolling avg, last_7_rates),
history_summary (append today, keep last 14), last_updated (now), version (+1).

RETURN: Complete updated context JSON only. No markdown. Conservative — only update on clear evidence."""


def build_opening_question(context: dict, tasks_today: list) -> dict:
    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 = "No tasks were scheduled today — intentional, or did things go sideways?"
    elif completed == 0:
        q = f"None of today's {total} tasks got marked complete — derailed, or the plan didn't fit?"
    elif completed == total:
        q = f"You completed all {total} tasks — great day! Did it feel natural, or were you grinding through it?"
    elif len(incomplete) == 1:
        q = f'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 — 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 TASKS:
{json.dumps(tasks_today, indent=2)}

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

What should I ask next?"""

    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:
    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:
{json.dumps(context, indent=2)}

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

Completion rate: {completion_rate} ({completed}/{total})

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:
        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": []
        }
    }