"""Redis-based chat persistence."""

from __future__ import annotations

import base64
import datetime as dt
import json
import logging
import os
from pathlib import Path
from typing import Any, cast

import redis

from rossum_agent.storage import ChatData, ChatMetadata, _build_chat_list_item, _extract_first_user_text

logger = logging.getLogger(__name__)


class RedisStorage:
    """Redis storage for chat conversations."""

    def __init__(self, host: str | None = None, port: int | None = None, ttl_days: int = 30) -> None:
        self.host = host or os.getenv("REDIS_HOST", "localhost")
        self.port = port if port is not None else int(os.getenv("REDIS_PORT", "6379"))
        self.ttl = dt.timedelta(days=ttl_days)
        self._client: redis.Redis | None = None

    @property
    def client(self) -> redis.Redis:
        if self._client is None:
            self._client = redis.Redis(
                host=self.host, port=self.port, decode_responses=False, socket_connect_timeout=5
            )
        return self._client

    def save_chat(
        self,
        user_id: str | None,
        chat_id: str,
        messages: list[dict[str, Any]],
        output_dir: str | Path | None = None,
        metadata: ChatMetadata | None = None,
    ) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            # Shallow-copy messages so we can attach feedback without mutating caller's list
            messages = [dict(msg) for msg in messages]
            try:
                existing_raw = self.client.get(key)
                if isinstance(existing_raw, bytes):
                    existing_data = json.loads(existing_raw.decode("utf-8"))
                    if isinstance(existing_data, dict):
                        # Preserve per-message feedback from existing messages
                        for i, old_msg in enumerate(existing_data.get("messages", [])):
                            if i < len(messages) and old_msg.get("feedback") is not None:
                                messages[i]["feedback"] = old_msg["feedback"]
                        # Migrate legacy top-level feedback dict
                        for k, v in (existing_data.get("feedback") or {}).items():
                            idx = int(k)
                            if idx < len(messages) and "feedback" not in messages[idx]:
                                messages[idx]["feedback"] = v
            except Exception:
                logger.warning(f"Failed to preserve feedback from existing chat {chat_id}")
            payload = {
                "messages": messages,
                "output_dir": str(output_dir) if output_dir else None,
                "metadata": metadata.to_dict() if metadata else ChatMetadata().to_dict(),
            }
            value = json.dumps(payload).encode("utf-8")
            self.client.setex(key, self.ttl, value)

            files_saved = 0
            if output_dir:
                output_path = Path(output_dir) if isinstance(output_dir, str) else output_dir
                files_saved = self.save_all_files(chat_id, output_path)

            logger.info(
                f"Saved chat {chat_id} to Redis "
                f"(messages={len(messages)}, user: {user_id or 'shared'}, files={files_saved})"
            )
            return True
        except Exception as e:
            logger.error(f"Failed to save chat {chat_id}: {e}", exc_info=True)
            return False

    def load_chat(self, user_id: str | None, chat_id: str, output_dir: Path | None = None) -> ChatData | None:
        try:
            key = self._get_chat_key(user_id, chat_id)
            value = self.client.get(key)
            if value is None:
                logger.info(f"Chat {chat_id} not found in Redis (user: {user_id or 'shared'})")
                return None

            data = json.loads(cast("bytes", value).decode("utf-8"))
            messages = data if isinstance(data, list) else data.get("messages", [])
            stored_output_dir = data.get("output_dir") if isinstance(data, dict) else None
            metadata_dict = data.get("metadata", {}) if isinstance(data, dict) else {}
            metadata = ChatMetadata.from_dict(metadata_dict)

            files_loaded = 0
            if output_dir:
                files_loaded = self.load_all_files(chat_id, output_dir)

            logger.info(
                f"Loaded chat {chat_id} from Redis "
                f"({len(messages)} messages, {files_loaded} files, user: {user_id or 'shared'})"
            )
            return ChatData(messages=messages, output_dir=stored_output_dir, metadata=metadata)
        except Exception as e:
            logger.error(f"Failed to load chat {chat_id}: {e}", exc_info=True)
            return None

    def delete_chat(self, user_id: str | None, chat_id: str) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            deleted = self.client.delete(key)
            logger.info(f"Deleted chat {chat_id} from Redis (deleted={deleted}, user: {user_id or 'shared'})")
            return bool(deleted)
        except Exception as e:
            logger.error(f"Failed to delete chat {chat_id}: {e}", exc_info=True)
            return False

    def chat_exists(self, user_id: str | None, chat_id: str) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            return bool(self.client.exists(key))
        except Exception as e:
            logger.error(f"Failed to check if chat {chat_id} exists: {e}", exc_info=True)
            return False

    def is_connected(self) -> bool:
        try:
            self.client.ping()
            return True
        except Exception:
            return False

    def _get_chat_key(self, user_id: str | None, chat_id: str) -> str:
        return f"user:{user_id}:chat:{chat_id}" if user_id else f"chat:{chat_id}"

    def _get_chat_pattern(self, user_id: str | None = None) -> str:
        return f"user:{user_id}:chat:*" if user_id else "chat:*"

    def list_all_chats(self, user_id: str | None = None) -> list[dict[str, Any]]:
        try:
            pattern = self._get_chat_pattern(user_id)
            keys = cast("list[bytes]", self.client.keys(pattern.encode("utf-8")))
            chats = []

            for key in keys:
                key_str = key.decode("utf-8")
                chat_id = key_str.replace(f"user:{user_id}:chat:", "") if user_id else key_str.replace("chat:", "")

                chat_data = self.load_chat(user_id, chat_id)

                if chat_data:
                    chats.append(
                        _build_chat_list_item(
                            chat_id,
                            len(chat_data.messages),
                            _extract_first_user_text(chat_data.messages),
                            chat_data.metadata,
                        )
                    )

            chats.sort(key=lambda x: x["timestamp"], reverse=True)
            logger.info(f"Found {len(chats)} chats in Redis (user: {user_id or 'shared'})")
            return chats
        except Exception as e:
            logger.error(f"Failed to list chats: {e}", exc_info=True)
            return []

    def save_file(self, chat_id: str, file_path: Path | str, content: bytes | None = None) -> bool:
        # If content is None, reads from file_path on disk
        try:
            if isinstance(file_path, str):
                file_path = Path(file_path)

            filename = file_path.name
            key = f"file:{chat_id}:{filename}"

            if content is None:
                if not file_path.exists():
                    logger.error(f"File not found: {file_path}")
                    return False
                content = file_path.read_bytes()

            # Store file with metadata
            metadata = {
                "filename": filename,
                "size": len(content),
                "timestamp": dt.datetime.now(dt.UTC).isoformat(),
                "content": base64.b64encode(content).decode("utf-8"),
            }

            value = json.dumps(metadata).encode("utf-8")
            self.client.setex(key, self.ttl, value)
            logger.info(f"Saved file {filename} for chat {chat_id} to Redis ({len(content)} bytes)")
            return True
        except Exception as e:
            logger.error(f"Failed to save file {filename} for chat {chat_id}: {e}", exc_info=True)
            return False

    def load_file(self, chat_id: str, filename: str) -> bytes | None:
        try:
            key = f"file:{chat_id}:{filename}"
            value = self.client.get(key)
            if value is None:
                logger.info(f"File {filename} not found for chat {chat_id}")
                return None

            metadata: dict[str, Any] = json.loads(cast("bytes", value).decode("utf-8"))
            content = base64.b64decode(metadata["content"])
            logger.info(f"Loaded file {filename} for chat {chat_id} ({len(content)} bytes)")
            return content
        except Exception as e:
            logger.error(f"Failed to load file {filename} for chat {chat_id}: {e}", exc_info=True)
            return None

    def list_files(self, chat_id: str) -> list[dict[str, Any]]:
        try:
            pattern = f"file:{chat_id}:*"
            keys = cast("list[bytes]", self.client.keys(pattern.encode("utf-8")))
            files = []

            for key in keys:
                key_str = key.decode("utf-8")
                filename = key_str.split(":")[-1]
                value = self.client.get(key)
                if value:
                    metadata: dict[str, Any] = json.loads(cast("bytes", value).decode("utf-8"))
                    files.append(
                        {
                            "filename": filename,
                            "size": metadata.get("size", 0),
                            "timestamp": metadata.get("timestamp", ""),
                        }
                    )

            logger.info(f"Found {len(files)} files for chat {chat_id}")
            return files
        except Exception as e:
            logger.error(f"Failed to list files for chat {chat_id}: {e}", exc_info=True)
            return []

    def delete_file(self, chat_id: str, filename: str) -> bool:
        try:
            key = f"file:{chat_id}:{filename}"
            deleted = self.client.delete(key)
            logger.info(f"Deleted file {filename} for chat {chat_id} (deleted={deleted})")
            return bool(deleted)
        except Exception as e:
            logger.error(f"Failed to delete file {filename} for chat {chat_id}: {e}", exc_info=True)
            return False

    def delete_all_files(self, chat_id: str) -> int:
        try:
            pattern = f"file:{chat_id}:*"
            keys = cast("list[bytes]", self.client.keys(pattern.encode("utf-8")))
            if not keys:
                logger.info(f"No files to delete for chat {chat_id}")
                return 0

            deleted = cast("int", self.client.delete(*keys))
            logger.info(f"Deleted {deleted} files for chat {chat_id}")
            return deleted
        except Exception as e:
            logger.error(f"Failed to delete files for chat {chat_id}: {e}", exc_info=True)
            return 0

    def save_all_files(self, chat_id: str, output_dir: Path) -> int:
        saved_count = 0
        try:
            if not output_dir.exists() or not output_dir.is_dir():
                logger.warning(f"Output directory does not exist: {output_dir}")
                return 0

            files = [f for f in output_dir.iterdir() if f.is_file()]
            for file_path in files:
                if self.save_file(chat_id, file_path):
                    saved_count += 1

            logger.info(f"Saved {saved_count}/{len(files)} files for chat {chat_id} to Redis")
            return saved_count
        except Exception as e:
            logger.error(f"Failed to save files for chat {chat_id}: {e}", exc_info=True)
            return saved_count

    def load_all_files(self, chat_id: str, output_dir: Path) -> int:
        loaded_count = 0
        try:
            output_dir.mkdir(parents=True, exist_ok=True)
            files = self.list_files(chat_id)

            for file_info in files:
                filename = file_info["filename"]
                content = self.load_file(chat_id, filename)
                if content:
                    file_path = output_dir / filename
                    file_path.write_bytes(content)
                    loaded_count += 1

            logger.info(f"Loaded {loaded_count}/{len(files)} files for chat {chat_id} from Redis")
            return loaded_count
        except Exception as e:
            logger.error(f"Failed to load files for chat {chat_id}: {e}", exc_info=True)
            return loaded_count

    def save_feedback(self, user_id: str | None, chat_id: str, turn_index: int, is_positive: bool) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            raw = self.client.get(key)
            if raw is None:
                return False
            data: dict[str, Any] = json.loads(cast("bytes", raw).decode("utf-8"))
            if isinstance(data, list):
                return False
            messages: list[dict[str, Any]] = data.get("messages", [])
            if turn_index >= len(messages):
                return False
            messages[turn_index]["feedback"] = is_positive
            # Remove from legacy top-level feedback dict if present
            legacy: dict[str, bool] = data.get("feedback") or {}
            legacy.pop(str(turn_index), None)
            if legacy:
                data["feedback"] = legacy
            else:
                data.pop("feedback", None)
            data["messages"] = messages
            self.client.setex(key, self.ttl, json.dumps(data).encode("utf-8"))
            logger.info(f"Saved feedback for chat {chat_id} turn {turn_index}: {is_positive}")
            return True
        except Exception as e:
            logger.error(f"Failed to save feedback for chat {chat_id}: {e}", exc_info=True)
            return False

    def get_feedback(self, user_id: str | None, chat_id: str) -> dict[int, bool]:
        try:
            key = self._get_chat_key(user_id, chat_id)
            raw = self.client.get(key)
            if raw is None:
                return {}
            data: dict[str, Any] = json.loads(cast("bytes", raw).decode("utf-8"))
            if isinstance(data, list):
                return {}
            result: dict[int, bool] = {}
            for i, msg in enumerate(data.get("messages", [])):
                fb = msg.get("feedback")
                if fb is not None:
                    result[i] = fb
            # Also include legacy top-level feedback (for old data not yet migrated)
            for k, v in (data.get("feedback") or {}).items():
                idx = int(k)
                if idx not in result:
                    result[idx] = v
            return result
        except Exception as e:
            logger.error(f"Failed to get feedback for chat {chat_id}: {e}", exc_info=True)
            return {}

    def delete_feedback(self, user_id: str | None, chat_id: str, turn_index: int) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            raw = self.client.get(key)
            if raw is None:
                return False
            data: dict[str, Any] = json.loads(cast("bytes", raw).decode("utf-8"))
            if isinstance(data, list):
                return False
            messages: list[dict[str, Any]] = data.get("messages", [])
            changed = turn_index < len(messages) and messages[turn_index].pop("feedback", None) is not None
            # Also remove from legacy top-level feedback dict if present
            legacy: dict[str, bool] = data.get("feedback") or {}
            if str(turn_index) in legacy:
                del legacy[str(turn_index)]
                data["feedback"] = legacy
                changed = True
            if not changed:
                return False
            data["messages"] = messages
            self.client.setex(key, self.ttl, json.dumps(data).encode("utf-8"))
            logger.info(f"Deleted feedback for chat {chat_id} turn {turn_index}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete feedback for chat {chat_id}: {e}", exc_info=True)
            return False

    def delete_all_feedback(self, user_id: str | None, chat_id: str) -> bool:
        try:
            key = self._get_chat_key(user_id, chat_id)
            raw = self.client.get(key)
            if raw is None:
                return False
            data: dict[str, Any] = json.loads(cast("bytes", raw).decode("utf-8"))
            if isinstance(data, list):
                return False
            for msg in data.get("messages", []):
                msg.pop("feedback", None)
            data.pop("feedback", None)  # Remove legacy top-level feedback dict
            self.client.setex(key, self.ttl, json.dumps(data).encode("utf-8"))
            logger.info(f"Deleted all feedback for chat {chat_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete all feedback for chat {chat_id}: {e}", exc_info=True)
            return False

    def close(self) -> None:
        if self._client is not None:
            self._client.close()
            self._client = None
            logger.info("Closed Redis connection")
