"""
Simple state tracking for appointment scheduling
"""

import json
import os
import logging
import datetime
from typing import Dict, Any, Optional, List

# Set up logger
logger = logging.getLogger("mango_ai_agent")

# Path to store appointment states
APPOINTMENT_STATES_FILE = "data/appointment_states.json"

# Ensure the data directory exists
os.makedirs(os.path.dirname(APPOINTMENT_STATES_FILE), exist_ok=True)

# In-memory storage for appointment states
appointment_states = {}

# Load existing appointment states if the file exists
if os.path.exists(APPOINTMENT_STATES_FILE):
    try:
        with open(APPOINTMENT_STATES_FILE, "r") as f:
            appointment_states = json.load(f)
    except Exception as e:
        logger.error(f"Error loading appointment states: {str(e)}")
        appointment_states = {}


def save_appointment_states():
    """Save appointment states to file"""
    try:
        with open(APPOINTMENT_STATES_FILE, "w") as f:
            json.dump(appointment_states, f, indent=2)
        return True
    except Exception as e:
        logger.error(f"Error saving appointment states: {str(e)}")
        return False


def get_appointment_state(conversation_id: str) -> Dict[str, Any]:
    """
    Get the appointment state for a conversation
    
    Args:
        conversation_id: The conversation ID
        
    Returns:
        Dict: The appointment state or an empty dict if not found
    """
    return appointment_states.get(conversation_id, {})


def update_appointment_state(conversation_id: str, state: Dict[str, Any]) -> bool:
    """
    Update the appointment state for a conversation
    
    Args:
        conversation_id: The conversation ID
        state: The appointment state
        
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        appointment_states[conversation_id] = state
        save_appointment_states()
        return True
    except Exception as e:
        logger.error(f"Error updating appointment state: {str(e)}")
        return False


def is_in_appointment_flow(conversation_id: str) -> bool:
    """
    Check if a conversation is in the appointment flow
    
    Args:
        conversation_id: The conversation ID
        
    Returns:
        bool: True if in appointment flow, False otherwise
    """
    state = get_appointment_state(conversation_id)
    return state.get("meeting_request_stage") == "collecting_info"


def get_missing_fields(conversation_id: str) -> List[str]:
    """
    Get missing fields for an appointment
    
    Args:
        conversation_id: The conversation ID
        
    Returns:
        List[str]: List of missing fields
    """
    state = get_appointment_state(conversation_id)
    return state.get("missing_fields", [])


def get_extracted_info(conversation_id: str) -> Dict[str, str]:
    """
    Get extracted information for an appointment
    
    Args:
        conversation_id: The conversation ID
        
    Returns:
        Dict[str, str]: Extracted information
    """
    state = get_appointment_state(conversation_id)
    return state.get("extracted_info", {})


def start_appointment_flow(
    conversation_id: str, 
    missing_fields: List[str], 
    extracted_info: Dict[str, str],
    suggested_topic: Optional[str] = None,
    confidence: float = 0.0
) -> bool:
    """
    Start the appointment flow for a conversation
    
    Args:
        conversation_id: The conversation ID
        missing_fields: List of missing fields
        extracted_info: Extracted information
        suggested_topic: Suggested topic for the appointment
        confidence: Confidence score for the meeting intent detection
        
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        # Merge any existing state to preserve information
        existing_state = get_appointment_state(conversation_id)
        existing_extracted_info = existing_state.get("extracted_info", {})
        
        # Combine existing and new extracted info, with new info taking precedence
        combined_info = {**existing_extracted_info, **extracted_info}
        
        # If we already have a topic but new info doesn't, preserve the existing one
        if "topic" in existing_extracted_info and "topic" not in extracted_info:
            combined_info["topic"] = existing_extracted_info["topic"]
        
        # If suggested topic is provided, use it
        if suggested_topic:
            combined_info["topic"] = suggested_topic
        
        state = {
            "meeting_request": True,
            "meeting_request_stage": "collecting_info",
            "missing_fields": missing_fields,
            "extracted_info": combined_info,
            "suggested_topic": suggested_topic or combined_info.get("topic"),
            "confidence": confidence,
            "last_updated": str(datetime.datetime.now())
        }
        update_appointment_state(conversation_id, state)
        return True
    except Exception as e:
        logger.error(f"Error starting appointment flow: {str(e)}")
        return False


def update_appointment_flow(
    conversation_id: str,
    missing_fields: List[str] = None,
    extracted_info: Dict[str, str] = None,
    confidence: float = None
) -> bool:
    """
    Update the appointment flow for a conversation
    
    Args:
        conversation_id: The conversation ID
        missing_fields: Updated list of missing fields (if None, keeps existing)
        extracted_info: Updated extracted information (if None, keeps existing)
        confidence: Updated confidence score (if None, keeps existing)
        
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        state = get_appointment_state(conversation_id)
        
        # If no appointment state exists, can't update
        if not state:
            logger.warning(f"No appointment state found for conversation {conversation_id}")
            return False
            
        # Update missing fields if provided
        if missing_fields is not None:
            state["missing_fields"] = missing_fields
        
        # Update extracted info if provided, merging with existing info
        if extracted_info is not None:
            existing_info = state.get("extracted_info", {})
            # Merge, with new info taking precedence
            state["extracted_info"] = {**existing_info, **extracted_info}
            
            # Automatically update missing fields based on what we've collected
            current_missing_fields = state.get("missing_fields", [])
            updated_missing_fields = []
            
            # Check which fields are still missing
            for field in current_missing_fields:
                # If the field is in extracted_info and has a non-empty value, it's no longer missing
                if field not in state["extracted_info"] or not state["extracted_info"][field]:
                    updated_missing_fields.append(field)
            
            # Update the missing fields
            state["missing_fields"] = updated_missing_fields
            logger.info(f"Updated missing fields for conversation {conversation_id}: {updated_missing_fields}")
        
        # Update confidence if provided
        if confidence is not None:
            state["confidence"] = confidence
        
        # Update last_updated timestamp
        state["last_updated"] = str(datetime.datetime.now())
        
        # If no more missing fields, mark as ready to schedule
        if not state["missing_fields"]:
            state["meeting_request_stage"] = "ready_to_schedule"
            logger.info(f"Conversation {conversation_id} is now ready to schedule")
        
        update_appointment_state(conversation_id, state)
        return True
    except Exception as e:
        logger.error(f"Error updating appointment flow: {str(e)}")
        return False


def complete_appointment_flow(conversation_id: str, appointment_id: str) -> bool:
    """
    Mark the appointment flow as completed
    
    Args:
        conversation_id: The conversation ID
        appointment_id: The appointment ID
        
    Returns:
        bool: True if successful, False otherwise
    """
    try:
        state = get_appointment_state(conversation_id)
        state["meeting_request_stage"] = "scheduled"
        state["appointment_id"] = appointment_id
        update_appointment_state(conversation_id, state)
        return True
    except Exception as e:
        logger.error(f"Error completing appointment flow: {str(e)}")
        return False