import re
from datetime import datetime
from typing import Dict, List, Any

# Global conversation storage dictionary
SYSTEM_CHART_CONVERSATIONS: Dict[str, Dict[str, Any]] = {}

def save_conversation(extraction_id: str, paper_title: str, conversation_history: List, paper_content: str, paper_type: str) -> str:
    """
    Save a conversation to the global conversation storage.
    
    Args:
        extraction_id: Unique identifier for the extraction task
        paper_title: Title of the paper
        conversation_history: List of conversation messages
        paper_content: Full content of the paper
        paper_type: Type of paper (Experimental, Computational, Mixed)
    
    Returns:
        str: Unique conversation key
    """
    # Initialize extraction_id entry if it doesn't exist
    if extraction_id not in SYSTEM_CHART_CONVERSATIONS:
        SYSTEM_CHART_CONVERSATIONS[extraction_id] = {}
    
    # Create a safe key from paper title (matching original behavior)
    safe_title = re.sub(r'[^\w\s-]', '', paper_title)[:100]
    safe_title = re.sub(r'[-\s]+', '_', safe_title)
    
    conversation_key = f"{safe_title}_{len(SYSTEM_CHART_CONVERSATIONS[extraction_id])}"
    
    # Save the conversation
    SYSTEM_CHART_CONVERSATIONS[extraction_id][conversation_key] = {
        "paper_title": paper_title,
        "conversation_history": conversation_history.copy(),  # Make a copy to avoid modification
        "paper_content": paper_content,
        "paper_type": paper_type,
        "created_at": datetime.now().isoformat(),
        "last_updated": datetime.now().isoformat()
    }
    
    print(f"💾 Saved conversation for paper: {paper_title} (key: {conversation_key})")
    return conversation_key

def get_conversations(extraction_id: str) -> Dict[str, Dict[str, Any]]:
    """
    Get all conversations for a specific extraction ID.
    
    Args:
        extraction_id: Unique identifier for the extraction task
    
    Returns:
        Dict of conversation_key -> conversation_data
    """
    return SYSTEM_CHART_CONVERSATIONS.get(extraction_id, {})

def get_all_conversations() -> Dict[str, Dict[str, Any]]:
    """
    Get all conversations across every extraction ID.

    Returns:
        Dict of extraction_id -> {conversation_key -> conversation_data}
    """
    return SYSTEM_CHART_CONVERSATIONS

def update_conversation(extraction_id: str, conversation_key: str, new_message, response):
    """
    Update an existing conversation with a new message and response.
    
    Args:
        extraction_id: Unique identifier for the extraction task
        conversation_key: Unique conversation key
        new_message: New message to add
        response: Response to the message
    
    Returns:
        bool: True if update was successful, False otherwise
    """
    if (extraction_id in SYSTEM_CHART_CONVERSATIONS and 
        conversation_key in SYSTEM_CHART_CONVERSATIONS[extraction_id]):
        
        conversation = SYSTEM_CHART_CONVERSATIONS[extraction_id][conversation_key]
        conversation["conversation_history"].append(new_message)
        conversation["conversation_history"].append(response)
        conversation["last_updated"] = datetime.now().isoformat()
        return True
    return False