tools.py

import json
import os
import random
import logging
from datetime import datetime, timedelta, time
from typing import Type, Dict, List, Any, Optional
import pytz
import asyncio
from crewai.tools import BaseTool
from pydantic import BaseModel, Field

# Import the custom IntervalTree
from interval_tree import IntervalTree

# Setup standard logging
logger = logging.getLogger("scheduler")

# Google Authentication and API imports
GOOGLE_LIBS_AVAILABLE = False
try:
    from google.auth.transport.requests import Request
    from google.oauth2.credentials import Credentials
    from google_auth_oauthlib.flow import InstalledAppFlow
    from googleapiclient.discovery import build
    GOOGLE_LIBS_AVAILABLE = True
except ImportError:
    pass

# ==========================================
# Resilience: Exponential Backoff Retry Policy
# ==========================================

async def retry_with_backoff(
    coro_func, 
    *args, 
    max_retries: int = 3, 
    initial_delay: float = 1.0, 
    factor: float = 2.0, 
    jitter: bool = True,
    **kwargs
):
    """
    Executes an async function with exponential backoff and randomized jitter.
    """
    delay = initial_delay
    last_exception = None
    
    for attempt in range(max_retries + 1):
        try:
            return await coro_func(*args, **kwargs)
        except Exception as e:
            last_exception = e
            if attempt == max_retries:
                break
                
            sleep_time = delay * (factor ** attempt)
            if jitter:
                sleep_time += random.uniform(0, 1.0)
                
            logger.warning(
                f"Attempt {attempt + 1} failed: {str(e)}. Retrying in {sleep_time:.2f}s..."
            )
            await asyncio.sleep(sleep_time)
            
    raise last_exception

# ==========================================
# 1. GetParticipantTimezoneTool
# ==========================================

class GetParticipantTimezoneInput(BaseModel):
    name: str = Field(..., description="The name of the participant.")
    location: str = Field(..., description="The location/city of the participant.")

class GetParticipantTimezoneTool(BaseTool):
    name: str = "Get Participant Timezone Tool"
    description: str = "Maps a participant's location to their standard IANA timezone."
    args_schema: Type[BaseModel] = GetParticipantTimezoneInput

    def _run(self, name: str, location: str) -> str:
        location_map = {
            "new york": "America/New_York",
            "london": "Europe/London",
            "tokyo": "Asia/Tokyo",
            "sydney": "Australia/Sydney",
            "bangalore": "Asia/Kolkata",
            "bengaluru": "Asia/Kolkata",
            "mumbai": "Asia/Kolkata",
            "india": "Asia/Kolkata"
        }
        
        loc_lower = location.lower().strip()
        for key, tz in location_map.items():
            if key in loc_lower:
                return tz
                
        try:
            pytz.timezone(location)
            return location
        except Exception:
            name_lower = name.lower().strip()
            if "alice" in name_lower: return "America/New_York"
            elif "bob" in name_lower: return "Europe/London"
            elif "charlie" in name_lower: return "Asia/Tokyo"
            elif "david" in name_lower: return "Australia/Sydney"
            elif "eve" in name_lower or "rahul" in name_lower: return "Asia/Kolkata"
            return "UTC"

# ==========================================
# 2. GetMockCalendarEventsTool (Hybrid API + Mock)
# ==========================================

class GetMockCalendarEventsInput(BaseModel):
    participant_name: str = Field(..., description="The name of the participant to fetch calendar events for.")
    timezone_str: str = Field("UTC", description="The standard IANA timezone of the participant.")
    target_date: str = Field("2026-06-24", description="The target date to query calendar events for (YYYY-MM-DD).")

class GetMockCalendarEventsTool(BaseTool):
    name: str = "Get Mock Calendar Events Tool"
    description: str = (
        "Retrieves calendar events (working hours and busy slots) for a participant. "
        "Attempts to authenticate and query Google Calendar API if credentials are present; otherwise, falls back to local JSON."
    )
    args_schema: Type[BaseModel] = GetMockCalendarEventsInput

    _google_creds: Optional[Any] = None
    _primary_calendar_tz: Optional[str] = None

    def _get_google_credentials(self) -> Optional[Any]:
        if self._google_creds and self._google_creds.valid:
            return self._google_creds
            
        token_path = "token.json"
        creds_path = "credentials.json"
        
        creds = None
        if os.path.exists(token_path):
            try:
                creds = Credentials.from_authorized_user_file(token_path)
            except Exception:
                pass
                
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                try:
                    creds.refresh(Request())
                except Exception:
                    creds = None
            
            if not creds:
                if os.path.exists(creds_path) and GOOGLE_LIBS_AVAILABLE:
                    logger.info("Initiating Google Authentication Flow in browser...")
                    flow = InstalledAppFlow.from_client_secrets_file(
                        creds_path, 
                        scopes=['https://www.googleapis.com/auth/calendar.readonly']
                    )
                    creds = flow.run_local_server(port=0)
                    with open(token_path, 'w') as token_file:
                        token_file.write(creds.to_json())
                    logger.info("Authentication successful. Token saved to token.json.")
                else:
                    return None
                    
        self._google_creds = creds
        return creds

    async def _fetch_availability_async(self, participant_name: str, timezone_str: str, target_date: str) -> Dict[str, Any]:
        await asyncio.sleep(0.4)
        
        if participant_name.lower().strip() == "bob" and os.environ.get("SIMULATE_TRANSIENT_FAILURE") == "true":
            os.environ["SIMULATE_TRANSIENT_FAILURE"] = "false"
            raise ConnectionError("Simulated network timeout connecting to calendar server.")

        # Parse participant's timezone and target date
        try:
            participant_tz = pytz.timezone(timezone_str)
        except Exception:
            participant_tz = pytz.UTC
            
        try:
            target_dt = datetime.strptime(target_date, "%Y-%m-%d").date()
        except ValueError:
            target_dt = datetime.utcnow().date()

        creds = self._get_google_credentials()
        
        # Determine if this participant matches the logged-in Google Calendar user
        is_primary_user = False
        primary_user_env = os.environ.get("PRIMARY_USER")
        
        if primary_user_env:
            is_primary_user = (participant_name.lower().strip() == primary_user_env.lower().strip())
        else:
            if creds and GOOGLE_LIBS_AVAILABLE:
                if not hasattr(self, "_primary_calendar_tz") or self._primary_calendar_tz is None:
                    try:
                        logger.info("Fetching primary calendar metadata to auto-detect timezone...")
                        service = build('calendar', 'v3', credentials=creds)
                        calendar_info = service.calendars().get(calendarId='primary').execute()
                        self._primary_calendar_tz = calendar_info.get('timeZone')
                        logger.info(f"Auto-detected Google Calendar timezone: {self._primary_calendar_tz}")
                    except Exception as e:
                        logger.warning(f"Could not fetch primary calendar timezone: {str(e)}")
                        self._primary_calendar_tz = "UNKNOWN"
                
                if self._primary_calendar_tz and self._primary_calendar_tz != "UNKNOWN":
                    try:
                        p_tz = pytz.timezone(timezone_str).zone
                        c_tz = pytz.timezone(self._primary_calendar_tz).zone
                        is_primary_user = (p_tz == c_tz)
                    except Exception:
                        is_primary_user = (timezone_str.lower().strip() == self._primary_calendar_tz.lower().strip())
            
            # Fallback check: if the calendar timezone is not in the project participants' timezones,
            # or if credentials are not available, default to "Alice".
            valid_tzs = {"america/new_york", "europe/london", "asia/tokyo", "australia/sydney", "asia/kolkata"}
            tz_not_in_project = True
            if creds and GOOGLE_LIBS_AVAILABLE and hasattr(self, "_primary_calendar_tz") and self._primary_calendar_tz and self._primary_calendar_tz != "UNKNOWN":
                try:
                    c_tz_lower = pytz.timezone(self._primary_calendar_tz).zone.lower()
                except Exception:
                    c_tz_lower = self._primary_calendar_tz.lower().strip()
                if c_tz_lower in valid_tzs:
                    tz_not_in_project = False
            
            if tz_not_in_project:
                is_primary_user = (participant_name.lower().strip() == "alice")
        
        if creds and is_primary_user and GOOGLE_LIBS_AVAILABLE:
            try:
                logger.info(f"Querying live Google Calendar API for '{participant_name}' on date {target_date}...")
                service = build('calendar', 'v3', credentials=creds)
                
                # Fetch calendar events for the exact target local day of the participant
                start_dt = datetime.combine(target_dt, time(0, 0))
                end_dt = start_dt + timedelta(days=1)
                
                local_start = participant_tz.localize(start_dt)
                local_end = participant_tz.localize(end_dt)
                
                time_min = local_start.astimezone(pytz.UTC).isoformat().replace("+00:00", "Z")
                time_max = local_end.astimezone(pytz.UTC).isoformat().replace("+00:00", "Z")
                
                events_result = service.events().list(
                    calendarId='primary', 
                    timeMin=time_min,
                    timeMax=time_max, 
                    singleEvents=True,
                    orderBy='startTime'
                ).execute()
                
                events = events_result.get('items', [])
                busy_slots = []
                for event in events:
                    if event.get('transparency') == 'transparent':
                        continue
                    
                    start_str = event['start'].get('dateTime', event['start'].get('date'))
                    end_str = event['end'].get('dateTime', event['end'].get('date'))
                    
                    try:
                        # Standardize Zulu suffix to ISO-offset format
                        if start_str.endswith('Z'):
                            start_str = start_str[:-1] + '+00:00'
                        if end_str.endswith('Z'):
                            end_str = end_str[:-1] + '+00:00'
                            
                        # Parse timezone-aware datetime
                        dt_start_aware = datetime.fromisoformat(start_str)
                        dt_end_aware = datetime.fromisoformat(end_str)
                        
                        # Convert to the participant's local timezone
                        dt_start_local = dt_start_aware.astimezone(participant_tz)
                        dt_end_local = dt_end_aware.astimezone(participant_tz)
                        
                        # Strip timezone info to create naive local ISO strings matching schema
                        s_naive = dt_start_local.replace(tzinfo=None).isoformat()
                        e_naive = dt_end_local.replace(tzinfo=None).isoformat()
                        
                        busy_slots.append({
                            "start": s_naive,
                            "end": e_naive,
                            "priority": "high",
                            "description": event.get('summary', 'Busy Block')
                        })
                    except Exception as e:
                        logger.warning(f"Failed to parse event dates: {str(e)}")
                        continue
                
                return {
                    "participant": participant_name,
                    "data": {
                        "location": "New York, USA",
                        "timezone": timezone_str,
                        "working_hours": {"start": "09:00", "end": "17:00"},
                        "busy_slots": busy_slots
                    }
                }
            except Exception as e:
                logger.error(f"Google Calendar sync failed: {str(e)}. Falling back to local database.")
                
        # Mock database backup
        mock_file_path = "calendars_mock.json"
        if not os.path.exists(mock_file_path):
            mock_file_path = os.path.join(os.path.dirname(__file__), "calendars_mock.json")
            
        with open(mock_file_path, "r") as f:
            data = json.load(f)
            
        participants = data.get("participants", {})
        matched_key = None
        for key in participants.keys():
            if key.lower().strip() == participant_name.lower().strip():
                matched_key = key
                break
                
        if not matched_key:
            raise KeyError(f"Participant '{participant_name}' not found in mock database.")
            
        # Verify the mock JSON matches target date (since we mock, let's substitute target_date dynamically)
        p_data = participants[matched_key].copy()
        
        # Modify the date of mock busy slots to match the query target date
        for slot in p_data.get("busy_slots", []):
            try:
                # e.g., 2026-06-24T10:00:00 -> target_date + T10:00:00
                t_part = slot["start"].split("T")[1]
                slot["start"] = f"{target_date}T{t_part}"
                
                t_end_part = slot["end"].split("T")[1]
                slot["end"] = f"{target_date}T{t_end_part}"
            except Exception:
                continue
                
        return {
            "participant": matched_key,
            "data": p_data
        }

    def _run(self, participant_name: str, timezone_str: str = "UTC", target_date: str = "2026-06-24") -> str:
        try:
            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                
            result = loop.run_until_complete(
                retry_with_backoff(self._fetch_availability_async, participant_name, timezone_str, target_date)
            )
            return json.dumps(result, indent=2)
        except Exception as e:
            logger.critical(f"Calendar query failed for '{participant_name}': {str(e)}")
            return json.dumps({
                "participant": participant_name,
                "error": f"Calendar retrieval failed after retries: {str(e)}",
                "data": None
            })

# ==========================================
# 3. CalculateMeetingWindowTool (Sweep-Line Math)
# ==========================================

class CalculateMeetingWindowInput(BaseModel):
    participants_json: str = Field(
        ..., 
        description="A JSON string mapping each participant's name to their calendar details (timezone, working_hours, and busy_slots)."
    )
    meeting_duration_minutes: int = Field(
        45, 
        description="The duration of the meeting in minutes (default 45)."
    )
    target_date: str = Field(
        "2026-06-24", 
        description="The target date for the meeting in YYYY-MM-DD format."
    )

class CalculateMeetingWindowTool(BaseTool):
    name: str = "Calculate Meeting Window Tool"
    description: str = (
        "Calculates the optimal overlapping meeting window using a sweep-line interval search algorithm and IntervalTree."
        "Handles timezones, sleep ranges, and custom weights, resolving soft conflicts if necessary."
    )
    args_schema: Type[BaseModel] = CalculateMeetingWindowInput

    def _run(self, participants_json: str, meeting_duration_minutes: int = 45, target_date: str = "2026-06-24") -> str:
        try:
            participants_data = json.loads(participants_json)
        except Exception as e:
            return json.dumps({"error": f"Invalid JSON format for participants_json: {str(e)}"})
            
        try:
            target_dt = datetime.strptime(target_date, "%Y-%m-%d").date()
        except ValueError:
            return json.dumps({"error": "Invalid date format. Must be YYYY-MM-DD."})
            
        parsed_participants = {}
        failed_participants = []
        
        for name, details in participants_data.items():
            if isinstance(details, str):
                try:
                    details = json.loads(details)
                except:
                    pass
            
            data_dict = details.get("data", details) if isinstance(details, dict) else {}
            
            if details.get("error") or data_dict is None or not data_dict:
                failed_participants.append(name)
                continue
                
            timezone_str = data_dict.get("timezone", "UTC")
            working_hours = data_dict.get("working_hours", {"start": "09:00", "end": "17:00"})
            busy_slots = data_dict.get("busy_slots", [])
            
            try:
                tz = pytz.timezone(timezone_str)
            except Exception:
                tz = pytz.UTC
                timezone_str = "UTC"
                
            parsed_participants[name] = {
                "timezone": tz,
                "timezone_str": timezone_str,
                "working_hours": working_hours,
                "busy_slots": busy_slots
            }
            
        if not parsed_participants:
            return json.dumps({"error": "No valid participant calendars retrieved."})

        # Core Sweep-line setup
        start_utc_day = datetime.combine(target_dt, time(0, 0), tzinfo=pytz.UTC)
        end_utc_day = start_utc_day + timedelta(days=1)
        slot_duration = timedelta(minutes=meeting_duration_minutes)
        
        critical_boundaries = {start_utc_day, end_utc_day}
        
        # 1. Populate Interval Trees & Collect Time Boundaries
        for name, p in parsed_participants.items():
            tz = p["timezone"]
            
            busy_tree = IntervalTree()
            
            work_start_str = p["working_hours"].get("start", "09:00")
            work_end_str = p["working_hours"].get("end", "17:00")
            
            for boundary_str in [work_start_str, work_end_str, "07:00", "22:00"]:
                try:
                    b_time = datetime.strptime(boundary_str, "%H:%M").time()
                    local_dt = datetime.combine(target_dt, b_time)
                    local_dt_tz = tz.localize(local_dt)
                    critical_boundaries.add(local_dt_tz.astimezone(pytz.UTC))
                except Exception:
                    continue
                    
            for slot in p["busy_slots"]:
                try:
                    s_naive = datetime.fromisoformat(slot["start"])
                    e_naive = datetime.fromisoformat(slot["end"])
                    
                    busy_start = tz.localize(s_naive).astimezone(pytz.UTC)
                    busy_end = tz.localize(e_naive).astimezone(pytz.UTC)
                    
                    critical_boundaries.add(busy_start)
                    critical_boundaries.add(busy_end)
                    
                    busy_tree.insert(busy_start, busy_end, slot)
                except Exception:
                    continue
                    
            p["busy_tree"] = busy_tree
                    
        sorted_boundaries = sorted([
            t for t in critical_boundaries 
            if start_utc_day <= t <= end_utc_day
        ])
        
        # 2. Candidate evaluation list
        candidate_starts = set()
        for t in sorted_boundaries:
            if start_utc_day <= t <= end_utc_day - slot_duration:
                candidate_starts.add(t)
            t_offset = t - slot_duration
            if start_utc_day <= t_offset <= end_utc_day - slot_duration:
                candidate_starts.add(t_offset)
                
        if len(candidate_starts) < 10:
            curr = start_utc_day
            while curr + slot_duration <= end_utc_day:
                candidate_starts.add(curr)
                curr += timedelta(minutes=15)
                
        sorted_candidates = sorted(list(candidate_starts))
        
        best_slot_start = None
        best_slot_score = float('inf')
        best_slot_breakdown = {}
        
        # 3. Evaluate each candidate start time
        for slot_start in sorted_candidates:
            slot_end = slot_start + slot_duration
            
            total_penalty = 0
            participant_penalties = {}
            
            for name, p in parsed_participants.items():
                tz = p["timezone"]
                local_start = slot_start.astimezone(tz)
                local_end = slot_end.astimezone(tz)
                
                work_start_str = p["working_hours"].get("start", "09:00")
                work_end_str = p["working_hours"].get("end", "17:00")
                
                try:
                    work_start_time = datetime.strptime(work_start_str, "%H:%M").time()
                    work_end_time = datetime.strptime(work_end_str, "%H:%M").time()
                except ValueError:
                    work_start_time = time(9, 0)
                    work_end_time = time(17, 0)
                    
                local_start_time = local_start.time()
                local_end_time = local_end.time()
                
                # Calculate working hour penalty
                if (local_start.date() == local_end.date() and 
                    local_start_time >= work_start_time and 
                    local_end_time <= work_end_time):
                    work_penalty = 0
                else:
                    def time_penalty(t: time) -> int:
                        if t >= time(22, 0) or t < time(7, 0):
                            return 100
                        elif t < work_start_time or t > work_end_time:
                            return 15
                        return 0
                    work_penalty = max(time_penalty(local_start_time), time_penalty(local_end_time))
                    
                # Calculate busy slots penalty USING THE CUSTOM INTERVAL TREE
                busy_penalty = 0
                overlapping_events = []
                
                clashes = p["busy_tree"].overlap_search(slot_start, slot_end)
                for clash in clashes:
                    slot_data = clash["data"]
                    priority = slot_data.get("priority", "high").lower()
                    desc = slot_data.get("description", "Busy")
                    if priority == "high":
                        busy_penalty += 1000
                    else:
                        busy_penalty += 50
                    overlapping_events.append(f"{desc} ({priority})")
                        
                participant_total = work_penalty + busy_penalty
                total_penalty += participant_total
                
                participant_penalties[name] = {
                    "local_time_start": local_start.strftime("%I:%M %p"),
                    "local_time_end": local_end.strftime("%I:%M %p"),
                    "timezone": p["timezone_str"],
                    "work_penalty": work_penalty,
                    "busy_penalty": busy_penalty,
                    "overlapping_events": overlapping_events,
                    "status": "Working Hours & Free" if participant_total == 0 else (
                        "Busy (Low Priority)" if busy_penalty == 50 and work_penalty == 0 else (
                            "Outside Working Hours (Shoulder)" if work_penalty == 15 and busy_penalty == 0 else (
                                "Sleep Hours" if work_penalty == 100 else "High Priority Conflict / Hard Clash"
                            )
                        )
                    )
                }
                
            if total_penalty < best_slot_score:
                best_slot_score = total_penalty
                best_slot_start = slot_start
                best_slot_breakdown = participant_penalties
                
        if best_slot_start is None:
            return json.dumps({"error": "Could not calculate a valid interval slot."})
            
        best_slot_end = best_slot_start + slot_duration
        is_compromise = best_slot_score > 0
        
        result = {
            "target_date": target_date,
            "meeting_duration_minutes": meeting_duration_minutes,
            "best_slot_utc": {
                "start": best_slot_start.strftime("%Y-%m-%dT%H:%M:%SZ"),
                "end": best_slot_end.strftime("%Y-%m-%dT%H:%M:%SZ")
            },
            "disruption_score": best_slot_score,
            "is_compromise_solution": is_compromise,
            "participant_schedules": best_slot_breakdown,
            "failed_participants": failed_participants
        }
        
        return json.dumps(result, indent=2)