from datetime import datetime
from pathlib import Path
import xml.etree.ElementTree as ET
import hashlib

from .util import format_timestamp

class IterationLogger:
    """Logs agent iterations to XML files"""
    
    def __init__(
            self,
            iterations_dir: Path,
            system_prompt: str,
            action_schema: str,
        ):
        """Initialize with directory for storing iteration files"""
        self.iterations_dir = iterations_dir
        self.iterations_dir.mkdir(parents=True, exist_ok=True)
        self._system_prompt_hash = hashlib.sha256(system_prompt.encode()).hexdigest()
        self._action_schema_hash = hashlib.sha256(action_schema.encode()).hexdigest()
        
    def log_iteration(
            self,
            timestamp: datetime,
            context: str,
            response: str,
        ):
        """
        Save an iteration to an XML file
        
        Args:
            context: The context as ElementTree
            response: Raw response from LLM
        """
        filename = f"iteration_{format_timestamp(timestamp)}.xml"
        filepath = self.iterations_dir / filename
        
        root = ET.Element("iteration")
        
        root.set("system_prompt_hash", self._system_prompt_hash)
        root.set("action_schema_hash", self._action_schema_hash)
        
        context_elem = ET.SubElement(root, "context")
        context_elem.text = context
        
        response_elem = ET.SubElement(root, "response")
        response_elem.text = response
        
        tree = ET.ElementTree(root)
        tree.write(filepath, encoding="utf-8", xml_declaration=True)