# TODO: Implement prefixes for Panels, similar to `get_prefixed_logger()`.

# archive_agent/core/CliManager.py
#  Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
#  This file is part of Archive Agent. See LICENSE for details.

import json
import logging
import time
from logging import Logger
import queue
import threading
from contextlib import contextmanager
from logging import Handler, LogRecord
from typing import Callable, List, Dict, cast, Optional, Any

import typer
from rich.console import Console, RenderableType, Group
from rich.live import Live
from rich.logging import RichHandler
from rich.panel import Panel
from rich.pretty import Pretty
from rich.table import Table

from qdrant_client.models import ScoredPoint

from archive_agent.ai.AiResult import AiResult

from archive_agent.ai.query.AiQuery import QuerySchema
from archive_agent.ai.vision.AiVisionSchema import VisionSchema
from archive_agent.util.PrefixedLogger import PrefixedLogger

from archive_agent.util.format import format_chunk_brief, get_point_reference_info
from archive_agent.db.QdrantSchema import parse_payload


class QueueHandler(Handler):
    """
    A logging handler that puts records into a queue.
    """

    def __init__(self, record_queue: queue.Queue):
        """
        Initialize the handler.
        :param record_queue: The queue to put records into.
        """
        super().__init__()
        self.record_queue = record_queue

    def emit(self, record: LogRecord):
        """
        Emit a record.
        """
        self.record_queue.put(record)


_LIVE_UPDATE_INTERVAL_S = 0.1


def _printer_thread_target(
        live: Live,
        q: queue.Queue,
        rich_handler: RichHandler,
        get_renderable: Callable[[], RenderableType]
):
    """
    The target function for the printer thread.
    Pulls items from a queue and prints them to the live console.
    Refreshes the live display at most every _LIVE_UPDATE_INTERVAL_S seconds
    to prevent rendering from becoming the bottleneck.
    """
    logger = logging.getLogger("ai.cli.printer")
    last_update = time.monotonic()
    last_backlog_warn = 0.0
    items_processed = 0
    while True:
        try:
            item = q.get(timeout=_LIVE_UPDATE_INTERVAL_S)
            if item is None:  # Sentinel value to signal thread to exit
                break

            if isinstance(item, LogRecord):
                rich_handler.handle(item)
            else:
                try:
                    live.console.print(item)
                except Exception as e:
                    logger.error(f"Failed to render output: {type(e).__name__}: {e}")
            q.task_done()
            items_processed += 1

            # Drain additional pending items without blocking (batch processing)
            while True:
                try:
                    item = q.get_nowait()
                    if item is None:  # Sentinel value to signal thread to exit
                        # Process remaining items, then update one final time and exit
                        try:
                            live.update(get_renderable())
                        except Exception as e:
                            logger.error(f"Live update error: {type(e).__name__}: {e}")
                        return
                    if isinstance(item, LogRecord):
                        rich_handler.handle(item)
                    else:
                        try:
                            live.console.print(item)
                        except Exception as e:
                            logger.error(f"Failed to render output: {type(e).__name__}: {e}")
                    q.task_done()
                    items_processed += 1
                except queue.Empty:
                    break

            # Periodic backlog warning (every 5 seconds)
            now = time.monotonic()
            qsize = q.qsize()
            if qsize > 50 and now - last_backlog_warn > 5.0:
                rich_handler.handle(logging.LogRecord(
                    name="ai.cli.printer", level=logging.WARNING,
                    pathname="", lineno=0, msg=f"Printer queue backlog: ({qsize}) items pending, ({items_processed}) processed",
                    args=None, exc_info=None
                ))
                last_backlog_warn = now
        except queue.Empty:
            pass  # Timeout occurred, just refresh the display
        except Exception as e:
            logger.error(f"Printer thread error: {type(e).__name__}: {e}")

        # Throttled live display update: at most every _LIVE_UPDATE_INTERVAL_S
        now = time.monotonic()
        if now - last_update >= _LIVE_UPDATE_INTERVAL_S:
            try:
                live.update(get_renderable())
            except Exception as e:
                logger.error(f"Live update error: {type(e).__name__}: {e}")
            last_update = now


class CliManager:
    """
    CLI manager.
    """

    VERBOSE_LOADER: bool = False  # enabled by --verbose flag
    VERBOSE_CHUNK: bool = False  # enabled by --verbose flag
    VERBOSE_RERANK: bool = False  # enabled by --verbose flag
    VERBOSE_EMBED: bool = False  # enabled by --verbose flag
    VERBOSE_QUERY: bool = False  # enabled by --verbose flag
    VERBOSE_VISION: bool = False  # enabled by --verbose flag
    VERBOSE_RETRIEVAL: bool = False  # enabled by --verbose flag
    VERBOSE_USAGE: bool = False  # enabled by --verbose flag

    def __init__(self, verbose: bool = False):
        """
        Initialize CLI manager.
        :param verbose: Verbosity switch.
        """
        CliManager.VERBOSE_LOADER = verbose
        CliManager.VERBOSE_CHUNK = verbose
        CliManager.VERBOSE_RERANK = verbose
        CliManager.VERBOSE_EMBED = verbose
        CliManager.VERBOSE_QUERY = verbose
        CliManager.VERBOSE_VISION = verbose
        CliManager.VERBOSE_RETRIEVAL = verbose
        CliManager.VERBOSE_USAGE = verbose

        self.logger = logging.getLogger()

        self.console = Console(markup=False)
        self._output_queue: Optional[queue.Queue] = None

        self.ai_usage_stats = {
            "chunk": 0,
            "embed": 0,
            "rerank": 0,
            "query": 0,
            "vision": 0
        }
        self.lock = threading.Lock()

    def get_prefixed_logger(self, prefix: str) -> Logger:
        """
        Get prefixed logger.
        :param prefix: Prefix.
        :return logger: Logger.
        """
        return PrefixedLogger(prefix, self.logger)

    @contextmanager
    def live_context(self, live: Live, get_renderable: Callable[[], RenderableType]) -> Any:
        """
        A context manager to correctly handle logging and printing during a Live display.
        It finds the existing RichHandler, redirects its output to the live console,
        and funnels all log messages through a queue to a dedicated printer thread.
        """
        q = queue.Queue()
        self._output_queue = q

        # Find the existing RichHandler from the configured logger.
        original_rich_handler = None
        for handler in self.logger.handlers:
            if isinstance(handler, RichHandler):
                original_rich_handler = handler
                break

        if not original_rich_handler:
            raise RuntimeError("RichHandler not found on the logger. Logging is not configured correctly.")

        # Store the original console and temporarily redirect the handler's output to the live console.
        original_console = original_rich_handler.console
        original_rich_handler.console = live.console

        # Start the printer thread, passing it the *original*, now-redirected handler.
        printer_thread = threading.Thread(
            target=_printer_thread_target,
            args=(live, q, original_rich_handler, get_renderable)
        )
        printer_thread.daemon = True
        printer_thread.start()

        # Temporarily replace the main logger's handlers with the QueueHandler to funnel all logs.
        queue_handler = QueueHandler(q)
        original_handlers = self.logger.handlers[:]
        self.logger.handlers = [queue_handler]

        try:
            yield
        finally:
            # Signal the printer thread to exit and wait for it to drain the queue.
            # Timeout prevents stall when Rich highlighting is slow on a backlog of messages.
            q.put(None)
            printer_thread.join(timeout=5.0)

            # Restore the original logger handlers.
            self.logger.handlers = original_handlers
            self._output_queue = None

            # IMPORTANT: Restore the original console to the handler.
            if original_rich_handler:
                original_rich_handler.console = original_console

    def _print(self, renderable: RenderableType):
        """
        Internal helper to print to the correct destination (queue or default console).
        """
        if self._output_queue:
            self._output_queue.put(renderable)
        else:
            self.console.print(renderable)

    def update_ai_usage(self, stats: Dict[str, int]):
        """
        Thread-safe method to update AI usage statistics.
        :param stats: A dictionary with token counts to add.
        """
        with self.lock:
            for category, value in stats.items():
                if category in self.ai_usage_stats:
                    self.ai_usage_stats[category] += value

    def get_ai_usage_renderable(self) -> Table:
        """
        Create a Rich Table to display AI usage statistics in a transposed layout.
        :return: A Rich Table object.
        """
        table = Table(
            title="AI Usage",
            show_header=False,
            show_edge=False,
            pad_edge=False,
            box=None,
            padding=(0, 2)
        )

        with self.lock:
            categories = [c.capitalize() for c in self.ai_usage_stats.keys()]
            values = [str(v) for v in self.ai_usage_stats.values()]

            # Define columns based on number of categories + 1 for the labels
            table.add_column(style="cyan", no_wrap=True)
            for _ in categories:
                table.add_column(justify="right", style="magenta")

            table.add_row("Category", *categories)
            table.add_row("Tokens", *values)

        return table

    @contextmanager
    def progress_context(self, progress_manager) -> Any:
        """
        A context manager for displaying hierarchical progress with live logging.
        Uses the provided ProgressManager from ContextManager.
        :param progress_manager: ProgressManager instance from ContextManager.
        """
        def get_renderable() -> Group:
            return Group(
                progress_manager.get_tree_renderable(),
                self.get_ai_usage_renderable()
            )

        with Live(get_renderable(), screen=False, redirect_stderr=False, transient=True) as live:
            with self.live_context(live, get_renderable):
                yield progress_manager, None  # yield ProgressManager directly

    def format_json(self, text: str) -> None:
        """
        Format text as JSON.
        :param text: Text.
        """
        try:
            data = json.loads(text)
            pretty = Pretty(data, expand_all=True)
            self._print(Panel(pretty, title="Structured output", style="blue", border_style="blue"))
        except json.JSONDecodeError:
            self._print(Panel(f"{text}", title="Raw output", style="red", border_style="red"))

    def prompt(self, message: str, is_cmd: bool, **kwargs) -> str:
        """
        Prompt user with message.
        :param message: Message.
        :param is_cmd: Enables "> " command style prompt.
        :param kwargs: Additional arguments for typer.prompt.
        :return: User input.
        """
        if is_cmd:
            self.logger.info(f"⚡ Archive Agent: {message}")
            return typer.prompt("", prompt_suffix="> ", **kwargs)
        else:
            self.logger.info(f"⚡ Archive Agent")
            return typer.prompt(message, prompt_suffix="", **kwargs)

    def format_ai_chunk(
            self,
            callback: Callable[[], AiResult],
            line_numbered_text: str,
    ) -> AiResult:
        """
        Format text to be chunked and AI result of chunk callback.
        :param callback: Chunk callback returning AI result.
        :param line_numbered_text: Text with line numbers.
        :return: AI result.
        """
        if CliManager.VERBOSE_CHUNK:
            self._print(Panel(f"{line_numbered_text}", title="Text", style="blue", border_style="blue"))
            self.logger.info("✨ Awaiting AI chunking…")

        result: AiResult = callback()

        self.update_ai_usage({"chunk": result.total_tokens})

        if CliManager.VERBOSE_USAGE:
            self.logger.info(f"Used ({result.total_tokens}) AI API token(s) for chunking")

        if CliManager.VERBOSE_CHUNK:
            self.format_json(result.output_text)

        return result

    def format_ai_rerank(
            self,
            callback: Callable[[], AiResult],
            indexed_chunks: Dict[int, str],
    ) -> AiResult:
        """
        Format chunks to be reranked and AI result of rerank callback.
        :param callback: Rerank callback returning AI result.
        :param indexed_chunks: Indexed chunks.
        :return: AI result.
        """
        if CliManager.VERBOSE_RERANK:
            indexed_chunks_str = "\n".join([
                f"{index:>3} : {format_chunk_brief(chunk=chunk)}"
                for index, chunk in indexed_chunks.items()
            ])
            self._print(Panel(f"{indexed_chunks_str}", title="Indexed Chunks", style="blue", border_style="blue"))
            self.logger.info("✨ Awaiting AI reranking…")

        result: AiResult = callback()

        self.update_ai_usage({"rerank": result.total_tokens})

        if CliManager.VERBOSE_USAGE:
            self.logger.info(f"Used ({result.total_tokens}) AI API token(s) for reranking")

        if CliManager.VERBOSE_RERANK:
            self.format_json(result.output_text)

        return result

    def format_ai_embed(
            self,
            callback: Callable[[], AiResult],
            text: str,
    ) -> AiResult:
        """
        Format text to be embedded.
        :param callback: Embed callback returning AI result.
        :param text: Text.
        :return: AI result.
        """
        if CliManager.VERBOSE_EMBED:
            self.format_chunk(text)
            self.logger.info("✨ Awaiting AI embedding…")

        result: AiResult = callback()

        self.update_ai_usage({"embed": result.total_tokens})

        if CliManager.VERBOSE_USAGE:
            self.logger.info(f"Used ({result.total_tokens}) AI API token(s) for embedding")

        return result

    def format_ai_query(
            self,
            callback: Callable[[], AiResult],
            prompt: str,
    ) -> AiResult:
        """
        Format prompt and AI result of query callback.
        :param callback: Query callback returning AI result.
        :param prompt: Prompt.
        :return: AI result.
        """
        if CliManager.VERBOSE_QUERY:
            self._print(Panel(f"{prompt}", title="Query", style="magenta", border_style="magenta"))
            self.logger.info("✨ Awaiting AI response…")

        result: AiResult = callback()

        self.update_ai_usage({"query": result.total_tokens})

        if CliManager.VERBOSE_USAGE:
            self.logger.info(f"Used ({result.total_tokens}) AI API token(s) for query")

        if CliManager.VERBOSE_QUERY:
            self.format_json(result.output_text)

        return result

    def format_ai_vision(
            self,
            callback: Callable[[], AiResult],
    ) -> AiResult:
        """
        Format AI result of vision callback.
        :param callback: Vision callback returning AI result.
        :return: AI result.
        """
        if CliManager.VERBOSE_VISION:
            self.logger.info("✨ Awaiting AI vision…")

        result: AiResult = callback()

        self.update_ai_usage({"vision": result.total_tokens})

        if CliManager.VERBOSE_VISION:
            vision_result = cast(VisionSchema, result.parsed_schema)
            if vision_result.is_rejected:
                self._print(
                    Panel(
                        f"{vision_result.rejection_reason}",
                        title="Vision rejected",
                        style="red",
                        border_style="red"
                    )
                )
            else:
                self.format_json(result.output_text)

        if CliManager.VERBOSE_USAGE:
            self.logger.info(f"Used ({result.total_tokens}) AI API token(s) for vision")

        return result

    def format_point(self, point: ScoredPoint) -> None:
        """
        Format point.
        :param point: Point.
        """
        self.logger.info(f"({point.score * 100:>6.2f} %) match: {get_point_reference_info(self.logger, point, verbose=True)}")

    def format_retrieved_points(self, points: List[ScoredPoint]) -> None:
        """
        Format chunks of retrieved points.
        :param points: Retrieved points.
        """
        if len(points) == 0:
            self.logger.info(f"⚠️ No retrieved results")
            return

        self.logger.info(f"✅ Retrieved ({len(points)}) chunk(s):")

        for point in points:

            assert point.payload is not None

            self.format_point(point)

            if CliManager.VERBOSE_RETRIEVAL:
                self.format_chunk(parse_payload(point.payload).chunk_text)

    def format_knee_cutoff(self, original_count: int, final_count: int, cutoff_index: int) -> None:
        """
        Format adaptive cutoff information.
        :param original_count: Original number of retrieved chunks.
        :param final_count: Chunks kept after cutoff.
        :param cutoff_index: Exclusive cutoff index.
        """
        self.logger.info(
            f"✅ Adaptive cutoff applied: {original_count} -> {final_count} chunk(s) "
            f"(cutoff index {cutoff_index})"
        )

    def format_knee_cutoff_settings(self, enabled: bool, sensitivity: float, min_chunks: int) -> None:
        """
        Format adaptive cutoff settings.
        :param enabled: Whether adaptive cutoff is enabled.
        :param sensitivity: Knee detection sensitivity.
        :param min_chunks: Minimum chunks to keep after cutoff.
        """
        self.logger.info(
            f"ℹ️ Adaptive cutoff settings: enabled={enabled}, "
            f"sensitivity={sensitivity}, min_chunks={min_chunks}"
        )

    def format_knee_cutoff_skipped(self, reason: str) -> None:
        """
        Format adaptive cutoff skip message.
        :param reason: Reason for skipping cutoff.
        """
        self.logger.info(f"ℹ️ Adaptive cutoff skipped: {reason}")

    def format_reranked_points(self, points: List[ScoredPoint]) -> None:
        """
        Format chunks of reranked points.
        :param points: Reranked points.
        """
        if len(points) == 0:
            self.logger.info(f"⚠️ No reranked results")
            return

        self.logger.info(f"✅ Reranked and limited down to ({len(points)}) chunk(s):")

        for point in points:

            assert point.payload is not None

            self.format_point(point)

            if CliManager.VERBOSE_RERANK:
                self.format_chunk(parse_payload(point.payload).chunk_text)

    def format_expanded_deduped_points(self, points: List[ScoredPoint]) -> None:
        """
        Format chunks of expanded and deduplicated points.
        :param points: Expanded and deduplicated points.
        """
        if len(points) == 0:
            self.logger.info(f"⚠️ No expanded and deduplicated results")
            return

        self.logger.info(f"✅ Expanded and deduplicated down to ({len(points)}) chunk(s):")

        for point in points:

            assert point.payload is not None

            self.format_point(point)

            if CliManager.VERBOSE_RERANK:
                self.format_chunk(parse_payload(point.payload).chunk_text)

    def format_chunk(self, chunk: str) -> None:
        """
        Format chunk.
        :param chunk: Chunk.
        """
        self._print(Panel(f"{chunk}", title="Chunk", style="orange3", border_style="orange3"))

    def format_question(self, question: str) -> None:
        """
        Format question.
        :param question: Question.
        """
        self._print(Panel(f"{question}", title="Question", style="magenta", border_style="magenta"))

    def format_query(self, query_result: QuerySchema, answer_text: str):
        """
        Format query.
        :param query_result: Query result.
        :param answer_text: Formatted answer.
        """
        if query_result.is_rejected:
            self._print(
                Panel(f"{query_result.rejection_reason}", title="Query rejected", style="red", border_style="red")
            )
        else:
            self._print(Panel(f"{answer_text}", title="Answer", style="green", border_style="green"))

    def usage(self):
        """
        Show AI token usage.
        """
        with self.lock:
            if any(self.ai_usage_stats.values()):
                self.logger.info(
                    f"Used AI API token(s): "
                    f"({self.ai_usage_stats['chunk']}) chunking, "
                    f"({self.ai_usage_stats['embed']}) embedding, "
                    f"({self.ai_usage_stats['rerank']}) reranking, "
                    f"({self.ai_usage_stats['query']}) query, "
                    f"({self.ai_usage_stats['vision']}) vision"
                )
            else:
                self.logger.info("No AI API tokens used")

        self.console.print(self.get_ai_usage_renderable())
