from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Dict, Generic, Optional

import chainlit as cl
from agentic_chat_ui.src.constants import TCChainlitCallBackRegistrar
from agentic_chat_ui.src.exceptions import StreamingError
from agentic_chat_ui.src.settings import (
    AuthChainlitCallBackRegistrarSettings,
    ChatChainlitCallBackRegistrarSettings,
)
from agentic_chat_ui.src.utils import hash_password

if TYPE_CHECKING:
    from agentic_chat_ui.src.view_runners.chainlit.runner import ChainlitViewRunner

logger = logging.getLogger(__name__)


class ChainlitCallBackRegistrar(ABC, Generic[TCChainlitCallBackRegistrar]):

    def __init__(self, config: TCChainlitCallBackRegistrar, chainlit_api: ChainlitViewRunner):
        self.config = config
        self.chainlit_api = chainlit_api

    @abstractmethod
    def register(self) -> None:
        raise NotImplementedError()


class ChatChainlitCallBackRegistrar(ChainlitCallBackRegistrar[ChatChainlitCallBackRegistrarSettings]):
    def __init__(self, config: ChatChainlitCallBackRegistrarSettings, chainlit_api: ChainlitViewRunner):
        super().__init__(config, chainlit_api)

    def register(self) -> None:
        controller = self.chainlit_api.controller  # local binding

        @cl.on_chat_start
        async def on_chat_start() -> None:
            await cl.Message(content="Hello! I'm connected to the Arxiv docs.").send()
            cl.user_session.set("history", [])

        @cl.on_message
        async def on_message(message: cl.Message) -> None:
            user_text = message.content

            # Retrieve conversation history
            history = cl.user_session.get("history") or []
            history.append({"role": "user", "content": user_text})

            # Create an empty assistant message
            msg = cl.Message(content="")
            await msg.send()

            try:
                # Stream tokens from controller → agent-api → LLM
                async for token in controller.astream_message(
                    user_text,
                    [h["content"] for h in history],
                ):
                    await msg.stream_token(token)

            except StreamingError:
                logger.exception("Streaming failed")
                await msg.stream_token("\n❌ Error while generating response")
            history.append({"role": "assistant", "content": msg.content})
            cl.user_session.set("history", history)

        logger.info("Chat callbacks registered")


class AuthChainlitCallBackRegistrar(ChainlitCallBackRegistrar[AuthChainlitCallBackRegistrarSettings]):
    def __init__(self, config: AuthChainlitCallBackRegistrarSettings, chainlit_api: ChainlitViewRunner):
        super().__init__(config, chainlit_api)
        self.users = self._load_users_from_config()

    def _load_users_from_config(self) -> Dict[str, Dict[str, str]]:
        """
        Load users from config.
        Returns:
            {
                "admin": {"password_hash": "...", "role": "user"}
            }
        """
        raw = self.config.users.get_secret_value()
        users: Dict[str, Dict[str, str]] = {}
        if not raw:
            logger.warning("CHAINLIT_USERS is not set")
            return users
        for entry in raw.split(","):
            try:
                username, password = entry.split(":", 1)
                users[username] = {
                    "password_hash": hash_password(password),
                    "role": "user",
                }
            except ValueError:
                logger.error("Invalid CHAINLIT_USERS entry: '%s'", entry)
        return users

    def register(self) -> None:
        @cl.password_auth_callback
        async def password_auth(
            username: str,
            password: str,
        ) -> Optional[cl.User]:
            user = self.users.get(username)
            if not user:
                return None
            if hash_password(password) != user["password_hash"]:
                return None
            return cl.User(
                identifier=username,
                metadata={
                    "role": user["role"],
                    "provider": "password",
                },
            )

        logger.info("Password auth callback registered")
