ha-text-ai/custom_components/ha_text_ai/history.py

"""
History management for HA Text AI integration.

@license: MIT (https://opensource.org/licenses/MIT)
@author: SMKRV
@github: https://github.com/smkrv/ha-text-ai
@source: https://github.com/smkrv/ha-text-ai
"""
from __future__ import annotations

import json
import logging
import os
import shutil
import traceback
from datetime import datetime
from typing import Any

import aiofiles

from homeassistant.core import HomeAssistant
from homeassistant.util import dt as dt_util

from .const import (
    ABSOLUTE_MAX_HISTORY_SIZE,
    MAX_ATTRIBUTE_SIZE,
    MAX_HISTORY_FILE_SIZE,
    TRUNCATION_INDICATOR,
)

# Per-entry storage cap (32KB per field) to prevent disk exhaustion
MAX_STORED_FIELD_SIZE = 32 * 1024
MAX_ARCHIVE_FILES = 3

_LOGGER = logging.getLogger(__name__)

class AsyncFileHandler:
    """Async context manager for file operations."""

    def __init__(self, file_path: str, mode: str = "a"):
        self.file_path = file_path
        self.mode = mode

    async def __aenter__(self):
        self.file = await aiofiles.open(self.file_path, self.mode)
        return self.file

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.file.close()

def _assert_not_symlink(path: str) -> None:
    """Refuse to operate on a path that resolves to a symlink.

    Why: another component or an attacker with filesystem access could
    replace our history file with a symlink pointing at arbitrary disk
    locations. Then os.remove or shutil.move would hit the target
    instead of our managed file. Check before destructive ops.
    """
    if os.path.islink(path):
        raise OSError(f"Refusing to operate on symlink: {path}")

class HistoryManager:
    """Manages conversation history for an instance."""

    def __init__(
        self,
        hass: HomeAssistant,
        instance_name: str,
        normalized_name: str,
        history_dir: str,
        max_history_size: int,
    ) -> None:
        self.hass = hass
        self.instance_name = instance_name
        self.normalized_name = normalized_name
        self._history_dir = history_dir
        self.max_history_size = min(
            max(1, max_history_size), ABSOLUTE_MAX_HISTORY_SIZE
        )
        self._history_file = os.path.join(
            history_dir, f"{normalized_name}_history.json"
        )
        self._max_history_file_size = MAX_HISTORY_FILE_SIZE
        self._conversation_history: list[dict[str, Any]] = []

    @property
    def conversation_history(self) -> list[dict[str, Any]]:
        return self._conversation_history

    @property
    def history_size(self) -> int:
        return len(self._conversation_history)

    async def async_initialize(self) -> None:
        """Initialize history: directories, file, migration."""
        await self._create_history_dir()
        await self._check_history_directory()
        await self._initialize_history_file()
        await self._migrate_history_from_txt_to_json()

    async def _file_exists(self, path: str) -> bool:
        try:
            return await self.hass.async_add_executor_job(os.path.exists, path)
        except Exception as e:
            _LOGGER.error("Error checking file existence for %s: %s", path, e)
            return False

    async def _create_history_dir(self) -> None:
        try:
            await self.hass.async_add_executor_job(
                os.makedirs, self._history_dir, 0o755, True
            )
        except PermissionError:
            _LOGGER.error("Permission denied creating history directory: %s", self._history_dir)
            raise
        except OSError as e:
            _LOGGER.error("Error creating history directory %s: %s", self._history_dir, e)
            raise

    async def _check_history_directory(self) -> None:
        """Check history directory permissions and writability."""
        try:
            test_file_path = os.path.join(self._history_dir, ".write_test")
            await self.hass.async_add_executor_job(
                self._sync_test_directory_write, test_file_path
            )
        except PermissionError:
            _LOGGER.error("No write permissions for history directory: %s", self._history_dir)
        except Exception as e:
            _LOGGER.error("Error checking history directory: %s", e)

    @staticmethod
    def _sync_test_directory_write(test_file_path: str) -> None:
        try:
            os.makedirs(os.path.dirname(test_file_path), mode=0o755, exist_ok=True)
            with open(test_file_path, "w") as f:
                f.write("Permission test")
            os.remove(test_file_path)
        except Exception as e:
            _LOGGER.error("Directory write test failed: %s", e)

    async def _initialize_history_file(self) -> None:
        """Initialize history file and load existing history."""
        try:
            if await self._file_exists(self._history_file):
                async with AsyncFileHandler(self._history_file, "r") as f:
                    content = await f.read()
                    if content:
                        history = json.loads(content)
                        if isinstance(history, list):
                            self._conversation_history = history[
                                -self.max_history_size :
                            ]
                            _LOGGER.debug(
                                "Loaded %d history entries for %s",
                                len(self._conversation_history),
                                self.instance_name,
                            )
            else:
                async with AsyncFileHandler(self._history_file, "w") as f:
                    await f.write(json.dumps([]))

            await self._check_history_size()
        except Exception as e:
            _LOGGER.error("Could not initialize history file: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def update_history(self, question: str, response: dict) -> None:
        """Update conversation history.

        In-memory history stores full text for context retrieval.
        On-disk storage caps per-field size to prevent disk exhaustion.
        Display truncation is handled by get_limited_history().
        """
        try:
            content = response.get("content", "")
            history_entry = {
                "timestamp": dt_util.utcnow().isoformat(),
                "question": question[:MAX_STORED_FIELD_SIZE],
                "response": content[:MAX_STORED_FIELD_SIZE],
            }

            self._conversation_history.append(history_entry)

            while len(self._conversation_history) > self.max_history_size:
                self._conversation_history.pop(0)

            await self._save_history_to_file()
        except Exception as e:
            _LOGGER.error("Error updating history: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def _save_history_to_file(self) -> None:
        """Serialize in-memory history to file with rotation if needed."""
        try:
            data = json.dumps(self._conversation_history, indent=2)
            data_size = len(data.encode("utf-8"))

            if data_size > MAX_HISTORY_FILE_SIZE:
                await self._rotate_history()

            async with AsyncFileHandler(self._history_file, "w") as f:
                await f.write(data)
        except Exception as e:
            _LOGGER.error("Error writing history file: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def _check_history_size(self) -> None:
        if len(self._conversation_history) > self.max_history_size:
            _LOGGER.warning(
                "History size (%d) exceeds maximum (%d). Trimming...",
                len(self._conversation_history), self.max_history_size,
            )
            self._conversation_history = self._conversation_history[
                -self.max_history_size :
            ]

    async def _check_file_size(self, file_path: str) -> int:
        try:
            if await self._file_exists(file_path):
                return await self.hass.async_add_executor_job(
                    os.path.getsize, file_path
                )
            return 0
        except Exception as e:
            _LOGGER.error("Error checking file size for %s: %s", file_path, e)
            return 0

    async def _rotate_history(self) -> None:
        try:
            _LOGGER.debug("Starting history rotation for %s", self._history_file)
            await self._rotate_history_files()
        except Exception as e:
            _LOGGER.error("Error rotating history: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def _rotate_history_files(self) -> None:
        """Rotate history files with size validation."""
        try:
            if await self._file_exists(self._history_file):
                current_size = await self._check_file_size(self._history_file)

                if current_size > MAX_HISTORY_FILE_SIZE:
                    _LOGGER.info(
                        "Rotating history file. Current size: %d, Max: %d",
                        current_size, MAX_HISTORY_FILE_SIZE,
                    )

                    archive_file = os.path.join(
                        self._history_dir,
                        f"{self.normalized_name}_history_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}.json",
                    )

                    await self.hass.async_add_executor_job(
                        _assert_not_symlink, self._history_file
                    )
                    await self.hass.async_add_executor_job(
                        shutil.move, self._history_file, archive_file
                    )

                    async with AsyncFileHandler(self._history_file, "w") as f:
                        await f.write(
                            json.dumps(
                                self._conversation_history[
                                    -self.max_history_size :
                                ],
                                indent=2,
                            )
                        )

                    _LOGGER.info("History file rotated to: %s", archive_file)

                    # Clean up old archive files, keep only MAX_ARCHIVE_FILES
                    await self._cleanup_archives()
        except Exception as e:
            _LOGGER.error("History rotation failed: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def _cleanup_archives(self) -> None:
        """Remove old archive files beyond MAX_ARCHIVE_FILES."""
        try:
            prefix = f"{self.normalized_name}_history_"

            def find_archives():
                archives = []
                for f in os.listdir(self._history_dir):
                    if f.startswith(prefix) and f.endswith(".json") and f != os.path.basename(self._history_file):
                        archives.append(os.path.join(self._history_dir, f))
                archives.sort()
                return archives

            archives = await self.hass.async_add_executor_job(find_archives)
            if len(archives) > MAX_ARCHIVE_FILES:
                for old_file in archives[:-MAX_ARCHIVE_FILES]:
                    await self.hass.async_add_executor_job(
                        _assert_not_symlink, old_file
                    )
                    await self.hass.async_add_executor_job(os.remove, old_file)
                    _LOGGER.debug("Removed old archive: %s", old_file)
        except Exception as e:
            _LOGGER.warning("Archive cleanup error: %s", e)

    async def _migrate_history_from_txt_to_json(self) -> None:
        """Migrate old .txt history to .json format."""
        try:
            old_history_file = os.path.join(
                self._history_dir, f"{self.normalized_name}_history.txt"
            )

            if not await self._file_exists(old_history_file):
                return

            # Skip migration if JSON history already has entries
            if self._conversation_history:
                _LOGGER.debug(
                    "JSON history already has %d entries for %s, skipping txt migration",
                    len(self._conversation_history), self.instance_name,
                )
                return

            _LOGGER.info(
                "Found old history file for %s, migrating to JSON", self.instance_name
            )

            history_entries = []
            async with AsyncFileHandler(old_history_file, "r") as f:
                content = await f.read()

            for line in content.split("\n"):
                if not line or line.startswith("History initialized at:"):
                    continue
                try:
                    parts = line.split(": ", 1)
                    if len(parts) != 2:
                        continue
                    timestamp = parts[0]
                    content_parts = parts[1].split(" - ")
                    if len(content_parts) != 2:
                        continue
                    question = content_parts[0].replace("Question: ", "")
                    response = content_parts[1].replace("Response: ", "")
                    history_entries.append(
                        {
                            "timestamp": timestamp,
                            "question": question,
                            "response": response,
                        }
                    )
                except Exception as e:
                    _LOGGER.warning("Error parsing history line: %s. Error: %s", line, e)
                    continue

            if history_entries:
                async with AsyncFileHandler(self._history_file, "w") as f:
                    await f.write(json.dumps(history_entries, indent=2))

                backup_file = old_history_file + ".backup"
                await self.hass.async_add_executor_job(
                    shutil.move, old_history_file, backup_file
                )

                _LOGGER.info(
                    "Migrated %d entries from txt to JSON for %s. Old file: %s",
                    len(history_entries), self.instance_name, backup_file,
                )

                self._conversation_history = history_entries
        except Exception as e:
            _LOGGER.error("Error during history migration for %s: %s", self.instance_name, e)
            _LOGGER.debug(traceback.format_exc())

    async def async_clear_history(self) -> None:
        """Clear conversation history."""
        try:
            self._conversation_history = []
            if await self._file_exists(self._history_file):
                await self.hass.async_add_executor_job(
                    _assert_not_symlink, self._history_file
                )
                await self.hass.async_add_executor_job(os.remove, self._history_file)
            _LOGGER.info("History for %s cleared", self.instance_name)
        except Exception as e:
            _LOGGER.error("Error clearing history: %s", e)
            _LOGGER.debug(traceback.format_exc())

    async def async_get_history(
        self,
        limit: int | None = None,
        filter_model: str | None = None,
        start_date: str | None = None,
        include_metadata: bool = False,
        sort_order: str = "newest",
        default_model: str = "",
    ) -> list[dict[str, Any]]:
        """Get conversation history with optional filtering and sorting."""
        try:
            history = self._conversation_history.copy()

            if filter_model:
                history = [
                    entry for entry in history if entry.get("model") == filter_model
                ]

            if start_date:
                try:
                    start_dt = datetime.fromisoformat(
                        start_date.replace("Z", "+00:00")
                    )
                    history = [
                        entry
                        for entry in history
                        if datetime.fromisoformat(
                            entry["timestamp"].replace("Z", "+00:00")
                        )
                        >= start_dt
                    ]
                except (ValueError, KeyError) as e:
                    _LOGGER.warning("Invalid start_date format: %s. Error: %s", start_date, e)

            if sort_order == "oldest":
                history.sort(key=lambda x: x.get("timestamp", ""))
            else:
                history.sort(key=lambda x: x.get("timestamp", ""), reverse=True)

            # Clamp limit to ABSOLUTE_MAX_HISTORY_SIZE to prevent pathological
            # caller requests from producing multi-MB service payloads.
            if limit and limit > 0:
                effective_limit = min(int(limit), ABSOLUTE_MAX_HISTORY_SIZE)
                history = history[:effective_limit]

            if include_metadata:
                enriched = []
                for entry in history:
                    enriched_entry = dict(entry)
                    enriched_entry["metadata"] = {
                        "entry_size": len(str(entry)),
                        "question_length": len(entry.get("question", "")),
                        "response_length": len(entry.get("response", "")),
                        "model_used": entry.get("model", default_model),
                        "instance": self.instance_name,
                    }
                    enriched.append(enriched_entry)
                return enriched

            return history
        except Exception as e:
            _LOGGER.error("Error getting history: %s", e)
            return []

    def get_limited_history(self, max_display: int = 5) -> dict[str, Any]:
        """Get limited conversation history for sensor attributes.

        Returns last `max_display` entries with truncated text for HA state.
        """
        recent = self._conversation_history[-max_display:]
        limited_history = [
            {
                "timestamp": entry["timestamp"],
                "question": self._truncate_text(entry["question"], 4096),
                "response": self._truncate_text(entry["response"], 4096),
            }
            for entry in recent
        ]

        return {
            "entries": limited_history,
            "info": {
                "total_entries": len(self._conversation_history),
                "displayed_entries": len(limited_history),
            },
        }

    @staticmethod
    def _truncate_text(text: str, max_length: int = MAX_ATTRIBUTE_SIZE) -> str:
        """Safely truncate text to maximum length with indicator."""
        if not text:
            return ""
        if len(text) <= max_length:
            return text
        return text[:max_length] + TRUNCATION_INDICATOR