from typing import List, Dict, Optional
from datetime import datetime
import json
from pathlib import Path


class ChatHistory:
	"""Manages multi-turn chat history"""
	
	def __init__(self, max_history: int = 10):
		self.max_history = max_history
		self.history: List[Dict[str, str]] = []
	
	def add_message(self, role: str, content: str, source_document: Optional[str] = None, chunks: Optional[List[str]] = None):
		"""Add a message to chat history
		
		Args:
			role: Message role ("user" or "assistant")
			content: Message content
			source_document: Optional document filename used for assistant messages
			chunks: Optional list of chunk texts used for assistant messages (in chunk mode)
		"""
		message = {
			"role": role,
			"content": content,
			"timestamp": datetime.now().isoformat()
		}
		
		# Store document source for assistant messages
		if role == "assistant" and source_document:
			message["source_document"] = source_document
		
		# Store chunks for assistant messages
		if role == "assistant" and chunks:
			message["chunks"] = chunks
		
		self.history.append(message)
		
		# Keep only last N messages
		if len(self.history) > self.max_history * 2:  # *2 because we have user + assistant pairs
			self.history = self.history[-self.max_history * 2:]
	
	def get_recent_history(self, n_turns: Optional[int] = None) -> List[Dict[str, str]]:
		"""Get recent chat history formatted for LLM"""
		if n_turns is None:
			n_turns = self.max_history
		
		# Get last N turns (each turn = user + assistant)
		recent = self.history[-n_turns * 2:] if len(self.history) > n_turns * 2 else self.history
		
		# Format for OpenAI API (remove timestamp)
		return [{"role": msg["role"], "content": msg["content"]} for msg in recent]
	
	def get_last_turn(self) -> List[Dict[str, str]]:
		"""Get only the last chat turn (last user + assistant messages)"""
		if len(self.history) < 2:
			return []
		
		# Get last 2 messages (user + assistant)
		last_two = self.history[-2:]
		
		# Format for OpenAI API (remove timestamp)
		return [{"role": msg["role"], "content": msg["content"]} for msg in last_two]
	
	def get_last_document(self) -> Optional[str]:
		"""Get the document filename used in the last assistant response
		
		Returns:
			Document filename if last message was assistant with a document, None otherwise
		"""
		if not self.history:
			return None
		
		# Check last message
		last_msg = self.history[-1]
		if last_msg.get("role") == "assistant":
			return last_msg.get("source_document")
		
		return None
	
	def get_last_turn_with_document(self) -> Optional[List[Dict[str, str]]]:
		"""Get the last chat turn that used a document (skipping general questions)
		
		Returns:
			List of messages from the last turn with a document, or None if no such turn exists
		"""
		# Search backwards through history to find last assistant message with a document
		for i in range(len(self.history) - 1, 0, -1):
			msg = self.history[i]
			if msg.get("role") == "assistant" and msg.get("source_document"):
				# Found an assistant message with a document
				# Get the turn (user + assistant pair)
				if i >= 1 and self.history[i-1].get("role") == "user":
					# Format for OpenAI API (remove timestamp, keep source_document in metadata)
					return [
						{"role": self.history[i-1]["role"], "content": self.history[i-1]["content"]},
						{"role": msg["role"], "content": msg["content"]}
					]
		
		return None
	
	def get_last_chunks(self) -> Optional[List[str]]:
		"""Get the chunks used in the last assistant response
		
		Returns:
			List of chunk texts if last message was assistant with chunks, None otherwise
		"""
		if not self.history:
			return None
		
		# Check last message
		last_msg = self.history[-1]
		if last_msg.get("role") == "assistant":
			return last_msg.get("chunks")
		
		return None
	
	def clear(self):
		"""Clear chat history"""
		self.history = []