app.py

# app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os, re, json
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
from groq import Groq

# Initialize Groq client using environment variable
GROQ_KEY = "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"
if not GROQ_KEY:
    raise RuntimeError("Missing GROQ_API_KEY. Set it as an environment variable or in Hugging Face Secrets.")

client = Groq(api_key=GROQ_KEY)

app = FastAPI(title="Task Parsing API (Groq-powered)")

# -------------------- PROMPTS --------------------
llm_task_create_agent_system_prompt = f"""
You are a task creation agent. Extract the following from user input:

- Task Name (Text)
- Task Priority (High, Medium, Low)
- Task Deadline (STRICTLY 'YYYY-MM-DD HH:MM:SS')
- Task Status (Pending, Completed)
- Task Type (Work, Health, Personal)
- Date Created (If not given, use current time)

⚠️ ALL datetime values must follow 'YYYY-MM-DD HH:MM:SS'. NO natural language like 'tomorrow' or 'evening'.

Output JUST a JSON object. Example:
{{"Task Name": "Submit report", "Task Priority": "High", "Task Deadline": "2025-07-17 18:00:00", "Task Status": "Pending", "Task Type": "Work", "Date Created": "2025-07-01 12:00:00"}}
"""

# -------------------- HELPERS --------------------
def extract_json_from_response(text: str):
    match = re.search(r"\{[\s\S]*\}", text)
    if not match:
        return None
    try:
        return json.loads(match.group())
    except json.JSONDecodeError:
        return None

DATETIME_REGEX = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}"

def validate_datetime_format(dt_str: str) -> bool:
    if re.fullmatch(DATETIME_REGEX, dt_str):
        try:
            datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
            return True
        except:
            return False
    return False

# -------------------- GROQ CALLER --------------------
def call_groq_system(system_prompt: str, user_message: str, model: str = "llama-3.3-70b-versatile"):
    """Generic wrapper to call Groq chat API."""
    chat_completion = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    )
    return chat_completion.choices[0].message.content.strip()

# -------------------- AGENTS --------------------
def run_task_creation_agent(nl_task: str):
    reply = call_groq_system(llm_task_create_agent_system_prompt, nl_task)
    parsed = extract_json_from_response(reply)
    if not parsed:
        raise ValueError(f"LLM did not return valid JSON. Raw:\n{reply}")

    if not validate_datetime_format(parsed.get("Task Deadline", "")):
        raise ValueError(f"Invalid 'Task Deadline': {parsed.get('Task Deadline')}")

    if not validate_datetime_format(parsed.get("Date Created", "")):
        parsed["Date Created"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    parsed.setdefault("Task Status", "Pending")
    parsed.setdefault("Task Priority", "Medium")
    parsed.setdefault("Task Type", "Work")

    return parsed


def scheduler_agent(task_json: Dict[str, Any]) -> str:
    """Simplified scheduler — returns the deadline as the scheduled time."""
    deadline = task_json.get("Task Deadline")
    if not validate_datetime_format(deadline):
        raise ValueError("Scheduler requires a valid Task Deadline.")
    return deadline


def reminder_agent(task_json: Dict[str, Any], scheduled_time_str: str):
    scheduled = datetime.strptime(scheduled_time_str, "%Y-%m-%d %H:%M:%S")
    priority = (task_json.get("Task Priority") or "Medium").lower()

    if priority == "high":
        offsets = [timedelta(hours=24), timedelta(hours=1)]
    elif priority == "low":
        offsets = [timedelta(days=1), timedelta(days=3)]
    else:
        offsets = [timedelta(hours=12), timedelta(hours=1)]

    reminders = []
    for off in offsets:
        r = scheduled - off
        if r > datetime.now():
            reminders.append(r.strftime("%Y-%m-%d %H:%M:%S"))
    return reminders

# -------------------- MODELS --------------------
class ParseRequest(BaseModel):
    text: str

class ParseResponse(BaseModel):
    task_json: Dict[str, Any]
    scheduled_time: str
    reminders: list

# -------------------- ENDPOINTS --------------------
@app.post("/parse-task", response_model=ParseResponse)
def parse_task(req: ParseRequest):
    try:
        task_json = run_task_creation_agent(req.text)
        scheduled_time = scheduler_agent(task_json)
        reminders = reminder_agent(task_json, scheduled_time)
        return ParseResponse(task_json=task_json, scheduled_time=scheduled_time, reminders=reminders)
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/health")
def health():
    return {"status": "ok"}