agentic_implementation/tools.py

from schemas import (
    FetchEmailsParams,
    ShowEmailParams,
    AnalyzeEmailsParams,
    DraftReplyParams,
    SendReplyParams,
)
from typing import Any, Dict
from email_scraper import scrape_emails_from_sender, _load_email_db, _save_email_db, _is_date_in_range
from datetime import datetime
from typing import List
from openai import OpenAI
import json
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()

# Initialize OpenAI client
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)


def extract_date_range(query: str) -> Dict[str, str]:
    """
    Use an LLM to extract a date range from a user query.
    Returns {"start_date":"DD-MMM-YYYY","end_date":"DD-MMM-YYYY"}.
    """
    today_str = datetime.today().strftime("%d-%b-%Y")
    system_prompt = f"""
You are a date‐range extractor. Today is {today_str}.

Given a user query (in natural language), return _only_ valid JSON with:
  {{
    "start_date": "DD-MMM-YYYY",
    "end_date":   "DD-MMM-YYYY"
  }}

Interpret relative dates as:
- "today"       → {today_str} to {today_str}
- "yesterday"   → 1 day ago to 1 day ago
- "last week"   → 7 days ago to {today_str}
- "last month"  → 30 days ago to {today_str}
- "last N days" → N days ago to {today_str}

Examples:
- "emails from dev agarwal last week"  
  → {{ "start_date": "01-Jun-2025", "end_date": "{today_str}" }}
- "show me emails yesterday"  
  → {{ "start_date": "06-Jun-2025", "end_date": "06-Jun-2025" }}

Return _only_ the JSON object—no extra text.
"""

    messages = [
        {"role": "system",  "content": system_prompt},
        {"role": "user",    "content": query}
    ]
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0.0,
        messages=messages
    )
    content = resp.choices[0].message.content.strip()

    # Try direct parse; if the model added fluff, strip to the JSON block.
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        start = content.find("{")
        end = content.rfind("}") + 1
        return json.loads(content[start:end])


def fetch_emails(email: str, query: str) -> Dict:
    """
    Fetch emails from a sender within a date range extracted from the query.
    Now returns both date info and emails.
    
    Args:
        email: The sender's email address
        query: The original user query (for date extraction)
    
    Returns:
        Dict with date_info and emails
    """
    # Extract date range from query
    date_info = extract_date_range(query)
    start_date = date_info.get("start_date")
    end_date = date_info.get("end_date")
    
    # Fetch emails using the existing scraper
    emails = scrape_emails_from_sender(email, start_date, end_date)
    
    # Return both date info and emails
    return {
        "date_info": date_info,
        "emails": emails,
        "email_count": len(emails)
    }


def show_email(message_id: str) -> Dict:
    """
    Retrieve the full email record (date, time, subject, content, etc.)
    from the local cache by message_id.
    """
    db = _load_email_db()  # returns { sender_email: { "emails": [...], "last_scraped": ... }, ... }

    # Search each sender's email list
    for sender_data in db.values():
        for email in sender_data.get("emails", []):
            if email.get("message_id") == message_id:
                return email

    # If we didn't find it, raise or return an error structure
    raise ValueError(f"No email found with message_id '{message_id}'")


def draft_reply(email: Dict, tone: str) -> str:
    # call LLM to generate reply
    # return a dummy reply for now
    print(f"Drafting reply for email {email['id']} with tone: {tone}")
    return f"Drafted reply for email {email['id']} with tone {tone}."
    ...


def send_reply(message_id: str, reply_body: str) -> Dict:
    # SMTP / Gmail API send
    print(f"Sending reply to message {message_id} with body: {reply_body}")
    ...


def analyze_emails(emails: List[Dict]) -> Dict:
    """
    Summarize and extract insights from a list of emails.
    Returns a dict with this schema:
      {
        "summary": str,        # a concise overview of all emails
        "insights": [str, ...] # list of key observations or stats
      }
    """
    # 1) Prepare the email payload
    emails_payload = json.dumps(emails, ensure_ascii=False)

    # 2) Build the LLM prompt
    system_prompt = """
You are an expert email analyst. You will be given a JSON array of email objects,
each with keys: date, time, subject, content, message_id.

Your job is to produce _only_ valid JSON with two fields:
1. summary: a 1–2 sentence high-level overview of these emails.
2. insights: a list of 3–5 bullet-style observations or statistics 
   (e.g. "2 job offers found", "overall positive tone", "next action: reply").

Output exactly:

{
  "summary": "...",
  "insights": ["...", "...", ...]
}
"""
    messages = [
        {"role": "system",  "content": system_prompt},
        {"role": "user",    "content": f"Here are the emails:\n{emails_payload}"}
    ]

    # 3) Call the LLM
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0.0,
        messages=messages
    )

    # 4) Parse and return
    content = response.choices[0].message.content.strip()
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        # In case the model outputs extra text, extract the JSON block
        start = content.find('{')
        end = content.rfind('}') + 1
        return json.loads(content[start:end])


TOOL_MAPPING = {
    "fetch_emails": fetch_emails,
    "show_email": show_email,
    "analyze_emails": analyze_emails,
    "draft_reply": draft_reply,
    "send_reply": send_reply,
}