import os
import mimetypes
from langchain_core.messages import HumanMessage, SystemMessage
from pathlib import Path
from fetch_question import get_files_by_task_id

# Function to handle attachments based on file type
def handle_attachment(task_id, file_name):
    """
    Handles file attachments for a task.
    
    Args:
        task_id: The ID of the task
        file_name: The name of the attachment file
    
    Returns:
        Dictionary with attachment details, including how it should be handled
    """
    if not file_name or file_name.strip() == "":
        return {"status": "no_attachment"}
    
    # Get file content using existing function
    file_data = get_files_by_task_id(task_id)
    
    if file_data["status"] == "error":
        print(f"Error fetching attachment: {file_data.get('error')}")
        return {"status": "error", "message": file_data.get("error")}
    
    content_type = file_data["content_type"]
    raw_content = file_data["content"]
    
    # Determine if this file type can be directly processed by Claude
    # Images, spreadsheets, and PDFs can be processed directly
    direct_process = (
        content_type.startswith("image/") or 
        "spreadsheet" in content_type or 
        "excel" in content_type or
        content_type == "application/pdf" or
        "csv" in content_type
    )
    
    if direct_process:
        # For files that Claude can handle directly, return content for direct inclusion
        print(f"Attachment {file_name} (type: {content_type}) will be passed directly to Claude")
        return {
            "status": "success",
            "handling": "direct",
            "content_type": content_type,
            "raw_content": raw_content,
            "file_name": file_name
        }
    else:
        # For other files, save to disk and return path
        save_dir = os.path.join(os.getcwd(), "task_files")
        os.makedirs(save_dir, exist_ok=True)
        
        # Determine file extension
        file_extension = Path(file_name).suffix
        if not file_extension:
            extension = mimetypes.guess_extension(content_type) or ".bin"
            file_path = os.path.join(save_dir, f"{task_id}{extension}")
        else:
            file_path = os.path.join(save_dir, file_name)
            
        # Save the file
        with open(file_path, 'wb') as f:
            f.write(raw_content)
            
        print(f"Attachment {file_name} (type: {content_type}) saved to {file_path}")
        return {
            "status": "success",
            "handling": "tool",
            "content_type": content_type,
            "file_path": file_path,
            "file_name": file_name
        }