import asyncio
import json
import logging
import os
import pickle
import random
import re
import time
from functools import lru_cache, wraps
from io import BytesIO
from pathlib import Path
from pprint import pformat
from typing import Any, Dict, List, Optional, Tuple, Union

from .images import ImageCache
from .memory import MemoryManager
from .persistence import PersistentStore


def pp(*args, **kw):
    if "width" not in kw:
        kw["width"] = 300
    return pformat(*args, **kw)


@lru_cache(maxsize=300)
def parse_json(content: str) -> Dict:
    # Strict JSON only — model output is schema-enforced (ENV-18/19),
    # history entries are json.dumps products.
    return json.loads(content.strip())


def sanitize_external_text(text: str, max_len: int = 4000) -> str:
    """Neutralize attacker-influenced text before it enters a prompt (SAF-03)."""
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", text)
    text = text.replace("@everyone", "@​everyone").replace("@here", "@​here")
    return text[:max_len]


def exponential_backoff(base=2, max_delay=60, factor=1, jitter=0.1, max_attempts=None):
    """Generate sleep intervals for exponential backoff with jitter.

    Args:
        base: Base of the exponentiation operation
        max_delay: Maximum delay
        factor: Multiplication factor for each increase in backoff
        jitter: Additional randomness range to prevent thundering herd problem

    Yields:
        Delay for backoff as a floating point number.
    """
    attempt = 0
    while True:
        sleep = min(max_delay, factor * base**attempt)
        jitter_amount = jitter * sleep
        sleep += random.uniform(-jitter_amount, jitter_amount)
        yield sleep
        attempt += 1
        if max_attempts is not None and attempt > max_attempts:
            raise RuntimeError("Max attempts reached in exponential backoff.")


def async_cache_to_file(filename):
    cache_file = Path(filename)
    cache = None
    if cache_file.exists():
        try:
            with cache_file.open("rb") as fd:
                cache = pickle.load(fd)
        except Exception:
            cache = {}

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if cache is None:
                return await func(*args, **kwargs)
            key = json.dumps((func.__name__, list(args[1:]), kwargs), sort_keys=True)
            if key in cache:
                return cache[key]
            result = await func(*args, **kwargs)
            cache[key] = result
            with cache_file.open("wb") as fd:
                pickle.dump(cache, fd)
            return result

        return wrapper

    return decorator


class AIMessageBase(object):
    def __init__(self) -> None:
        self.vars: List[str] = []

    def __str__(self) -> str:
        return json.dumps({k: v for k, v in vars(self).items() if k in self.vars})


class AIMessage(AIMessageBase):
    def __init__(self, user: str, message: str, channel: str = "chat", direct: bool = False, historise_question: bool = True) -> None:
        self.user = user
        self.message = message
        self.urls: Optional[List[str]] = None
        self.channel = channel
        self.direct = direct
        self.historise_question = historise_question
        self.vars = ["user", "message", "channel", "direct", "historise_question"]


class AIResponse(AIMessageBase):
    def __init__(
        self,
        answer: Optional[str],
        answer_needed: bool,
        channel: Optional[str],
        staff: Optional[str],
        picture: Optional[str],
        picture_edit: bool,
        hack: bool,
    ) -> None:
        self.answer = answer
        self.answer_needed = answer_needed
        self.channel = channel
        self.staff = staff
        self.picture = picture
        self.picture_count = 1
        self.picture_edit = picture_edit
        self.hack = hack
        self.vars = ["answer", "answer_needed", "channel", "staff", "picture", "hack"]


class AIResponderBase(object):
    def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
        super().__init__()
        self.config = config
        self.channel = channel if channel is not None else "system"


class AIResponder(AIResponderBase):
    def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
        super().__init__(config, channel)
        self.history: List[Dict[str, Any]] = []
        self.memory: str = "I am an assistant."
        self.rate_limit_backoff = exponential_backoff()
        self.store: Optional[PersistentStore] = None
        if "history-directory" in self.config:
            directory = Path(self.config["history-directory"]).expanduser()
            self.store = PersistentStore(directory / "bot.db")
            # Legacy pickles import once, then live on as *.migrated (PER-03)
            self.store.migrate_pickles(self.channel, directory / f"{self.channel}.dat", directory / f"{self.channel}.memory")
            self.history = self.store.load_history(self.channel)
            stored_memory = self.store.load_memory(self.channel)
            if stored_memory is not None:
                self.memory = stored_memory
        self.memory_manager = MemoryManager(self.store, lambda: self.config, self.consolidate, self.channel)
        self.image_cache: Optional[ImageCache] = None
        if self.store is not None:
            self.image_cache = ImageCache(self.store, Path(self.config["history-directory"]).expanduser() / "images", lambda: self.config)
        logging.info(f"memmory:\n{self.memory}")

    # Dynamic values move to a context suffix so the persona prefix
    # stays byte-stable for the prompt cache (ENV-20)
    DYNAMIC_PLACEHOLDERS = ("{date}", "{time}", "{news}", "{memory}")

    def _context_lines(self, message: AIMessage) -> List[str]:
        context = [f"date: {time.strftime('%Y-%m-%d')} ({time.strftime('%A')})", f"time: {time.strftime('%H:%M:%S')}"]
        news_feed = self.config.get("news")
        if news_feed and os.path.exists(news_feed):
            with open(news_feed) as fd:
                context.append("news:\n" + sanitize_external_text(fd.read().strip()))
        participants = [message.user] + [entry_user for entry_user in self._history_users(20)]
        memory_block = self.memory_manager.memory_block(participants, self.memory)
        if memory_block:
            context.append("memory:\n" + memory_block)
        if self.image_cache is not None:
            recent_images = self.image_cache.recent(message.channel, 4)
            if recent_images:
                # the model cannot use picture_edit unless told images exist (IMG-16)
                context.append(
                    f"recent images in this channel: {len(recent_images)}. When the user asks to modify, reuse, combine or"
                    " include a previously shared image, you MUST set picture_edit=true — text-to-image cannot see earlier"
                    " images; only picture_edit passes them to the image model."
                )
        return context

    def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        messages = []
        persona = self.config.get(self.channel, self.config["system"])
        for placeholder in self.DYNAMIC_PLACEHOLDERS:
            persona = persona.replace(placeholder, "")
        system = persona.rstrip() + "\n\n## Context\n" + "\n".join(self._context_lines(message))
        messages.append({"role": "system", "content": system})
        if limit is not None:
            while len(self.history) > limit:
                self.shrink_history_by_one()
        for msg in self.history:
            messages.append(msg)
        if not message.urls:
            messages.append({"role": "user", "content": str(message)})
        else:
            content: List[Dict[str, Union[str, Dict[str, str]]]] = [{"type": "text", "text": str(message)}]
            for url in message.urls:
                content.append({"type": "image_url", "image_url": {"url": url}})
            messages.append({"role": "user", "content": content})
        return messages

    async def draw(self, description: str, count: int = 1) -> List[BytesIO]:
        if self.config.get("leonardo-token") is not None:
            return [await self.draw_leonardo(description)]  # single image only, behind config
        return await self.draw_openai(description, count)

    async def draw_leonardo(self, description: str) -> BytesIO:
        raise NotImplementedError()

    async def draw_openai(self, description: str, count: int = 1) -> List[BytesIO]:
        raise NotImplementedError()

    async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse:
        # Envelope arrives schema-validated (ENV-19); .get defaults keep old
        # history entries and hand-built test dicts working.
        answer = response.get("answer")
        answer_needed = bool(response.get("answer_needed", False))
        if answer is None:
            answer_needed = False
        else:
            answer = str(answer)
            answer = re.sub(r"@\[([^\]]*)\]\([^\)]*\)", r"\1", answer)
            answer = re.sub(r"\[[^\]]*\]\(([^\)]*)\)", r"\1", answer)
            if message.direct or message.user in message.message:
                answer_needed = True
        response_message = AIResponse(
            answer,
            answer_needed,
            response.get("channel"),
            response.get("staff"),
            response.get("picture"),
            bool(response.get("picture_edit", False)),
            bool(response.get("hack", False)),
        )
        try:
            response_message.picture_count = max(1, min(int(response.get("picture_count") or 1), 4))  # IMG-02
        except (TypeError, ValueError):
            response_message.picture_count = 1
        if response_message.staff is not None and response_message.answer is not None:
            response_message.answer_needed = True
        if response_message.channel is None:
            response_message.channel = message.channel
        return response_message

    def short_path(self, message: AIMessage, limit: int) -> bool:
        if message.direct or "short-path" not in self.config:
            return False
        for chan_re, user_re in self.config["short-path"]:
            chan_ma = re.match(chan_re, message.channel)
            user_ma = re.match(user_re, message.user)
            if chan_ma and user_ma:
                self.history.append({"role": "user", "content": str(message)})
                while len(self.history) > limit:
                    self.shrink_history_by_one()
                return True
        return False

    async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]:
        raise NotImplementedError()

    async def consolidate(self, observations: List[Dict[str, Any]], known_facts: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        raise NotImplementedError()

    async def classify(self, message: AIMessage, history_tail: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """Cheap reply/factual/emoji pre-pass (BEH-01); None = fail open."""
        raise NotImplementedError()

    @staticmethod
    def _entry_channel(item: Dict[str, Any]) -> Optional[str]:
        try:
            return parse_json(item["content"]).get("channel")
        except Exception:
            return None

    def _history_users(self, tail: int) -> List[str]:
        users = []
        for item in self.history[-tail:]:
            try:
                user = parse_json(item["content"]).get("user")
            except Exception:
                user = None
            if user:
                users.append(str(user))
        return users

    def shrink_history_by_one(self) -> None:
        if not self.history:
            return
        cap = self.config.get("history-per-channel", 3)
        counts: Dict[Optional[str], int] = {}
        for item in self.history:
            chan = self._entry_channel(item)
            counts[chan] = counts.get(chan, 0) + 1
        for index, item in enumerate(self.history):
            if counts[self._entry_channel(item)] > cap:
                del self.history[index]
                return
        del self.history[0]

    def update_history(self, question: Dict[str, Any], answer: Dict[str, Any], limit: int, historise_question: bool = True) -> None:
        if not isinstance(question["content"], str):
            question["content"] = question["content"][0]["text"]
        if historise_question:
            self.history.append(question)
        self.history.append(answer)
        while len(self.history) > limit:
            self.shrink_history_by_one()

    async def _persist_history(self) -> None:
        if self.store is not None:
            await asyncio.to_thread(self.store.save_history, self.channel, list(self.history))

    async def handle_picture(self, response: Dict) -> bool:
        # Prompt goes to the image API verbatim — no translate step (IMG-05)
        if not isinstance(response.get("picture"), (type(None), str)):
            logging.warning(f"picture key is wrong in response: {pp(response)}")
            return False
        return True

    def _parse_answer(self, answer: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        # Schema-enforced output should always parse; anything else is a
        # failed attempt — no repair model (ENV-18).
        try:
            return parse_json(answer["content"])
        except Exception as err:
            logging.error(f"failed to parse the answer: {pp(err)}\n{repr(answer['content'])}")
            return None

    async def observe_event(self, user: str, kind: str, content: str) -> None:
        """Feed a Discord event into the observation stream (MEM-01)."""
        await self.memory_manager.observe(user, kind, content)

    async def send(self, message: AIMessage) -> AIResponse:
        # Get the history limit from the configuration
        limit = self.config["history-limit"]

        # Check if a short path applies, return an empty AIResponse if it does
        if self.short_path(message, limit):
            await self._persist_history()
            return AIResponse(None, False, None, None, None, False, False)

        # Number of retries for sending the message; failed attempts are
        # spaced by exponential backoff (ENV-12 / D1)
        retries = 3
        backoff = exponential_backoff(max_delay=10)

        async def failed_attempt() -> None:
            nonlocal retries
            retries -= 1
            if retries > 0:
                await asyncio.sleep(next(backoff))

        while retries > 0:
            # Get the message queue
            messages = self.message(message, limit)
            logging.info(f"try to send this messages:\n{pp(messages)}")

            # Attempt to send the message to the AI
            answer, limit = await self.chat(messages, limit)

            if answer is None:
                await failed_attempt()
                continue

            # Attempt to parse the AI's response (strict — ENV-18)
            response = self._parse_answer(answer)
            if response is None or not await self.handle_picture(response):
                await failed_attempt()
                continue

            # Post-process the message and update the answer's content
            answer_message = await self.post_process(message, response)
            answer["content"] = str(answer_message)

            # Update message history; persistence runs off the loop (PER-05)
            self.update_history(messages[-1], answer, limit, message.historise_question)
            await self._persist_history()
            logging.info(f"got this answer:\n{str(answer_message)}")

            # Feed the observation stream — consolidation is batched (MEM-01/02)
            await self.observe_event(message.user, "message", message.message)
            if answer_message.answer is not None:
                await self.observe_event("assistant", "message", answer_message.answer)

            # Return the updated answer message
            return answer_message

        # Raise an error if the process failed after all retries
        raise RuntimeError("Failed to generate answer after multiple retries")