import datetime
import json
import os
from typing import Dict, List, Optional
from langchain.schema import HumanMessage, AIMessage
from langchain.memory import ConversationBufferMemory

class CaptionHistory:
    """
    Manages caption generation history using Langchain
    """
    def __init__(self):
        self.memory = ConversationBufferMemory(
            return_messages=True,
            memory_key="chat_history"
        )
        self.history_file = "caption_history.json"
        self.load_history()  # Load existing history on initialization

    def add_interaction(self, image_name: str, model: str,
    caption: str, timestamp: str|None = None):
        if not timestamp:
            timestamp = datetime.datetime.now().isoformat()

        interaction = {
            "timestamp": timestamp,
            "image_name": image_name,
            "model": model,
            "caption": caption
        }

        # Add to langchain memory
        human_msg = HumanMessage(
            content=f"Generate caption for {image_name} using {model}"
        )
        ai_msg = AIMessage(content=caption)

        self.memory.chat_memory.add_user_message(human_msg.content)
        self.memory.chat_memory.add_ai_message(ai_msg.content)

        # Save the file 
        self.save_interaction(interaction)

    def get_history(self) -> List[Optional[Dict[str, str]]]:
        try:
            with open(self.history_file, mode="r") as f:
                return json.load(f)
        except FileNotFoundError:
            return []

    def save_interaction(self, interaction: Dict[str, str]) -> None:
        history = self.get_history()
        history.append(interaction)
        with open(self.history_file, mode="w") as f:
            json.dump(history, f, indent=2)

    def load_history(self):
        """Fixed: Proper string formatting in f-strings"""
        history = self.get_history()
        for item in history:
            human_msg = HumanMessage(
                content=f"Generate caption for {item['image_name']} using {item['model']}"  # Fixed: proper quotes
            )
            ai_msg = AIMessage(content=item["caption"]) 
            self.memory.chat_memory.add_user_message(human_msg.content)
            self.memory.chat_memory.add_ai_message(ai_msg.content)

    def clear_history(self):
        self.memory.clear()
        if os.path.exists(self.history_file):
            os.remove(self.history_file)