import json
import zlib

from backend import schemas
from backend.database.base import db
from backend.database.models import Conversation, Message, MessageAgent
from backend.utils.logging_utils import get_logger
from backend.utils.tool_validation import add_validation_response

logger = get_logger(__name__)


def _sanitize_merged_ctx(message_data: dict) -> dict:
    merged_context = message_data.get("merged_context")
    if not isinstance(merged_context, dict):
        return message_data

    sanitized_context = {key: value for key, value in merged_context.items() if key != "dkuCallerTicket"}
    message_data["merged_context"] = sanitized_context
    return message_data


def _get_db_message(message_id: str, user_id: str, *, for_update: bool = False):
    """Return an active user-owned message, optionally locking it for the transaction."""
    if not user_id:
        return None
    query = (
        db.session.query(Message)
        .join(Conversation, Message.conversation_id == Conversation.conversation_id)
        .filter(Message.id == message_id, Conversation.user_id == user_id, Message.status == "active")
    )
    if for_update:
        query = query.with_for_update()
    return query.first()


def _get_message_attachments(message_id: str) -> schemas.MessageAttachment | None:
    """
    Helper to fetch attachments for a message by id, no user scoping.
    Matches store.py _get_message_attachments.
    """
    from backend.database.models import MessageAttachment

    result = db.session.query(MessageAttachment).filter(MessageAttachment.message_id == message_id).first()
    if result:
        # TODO why do we need to serialize if the column is jsonendoded already?
        return schemas.MessageAttachment(
            attachments=json.loads(result.attachments) if isinstance(result.attachments, str) else result.attachments,
            extraction_mode=result.extraction_mode.get("mode") if result.extraction_mode else None,
            quota_exceeded=result.extraction_mode.get("quotaExceeded", False) if result.extraction_mode else False,
        )
    return None


def build_message(msg: Message) -> None:
    # Decompress artifacts
    artifacts = {}
    if msg.artifacts:
        try:
            artifacts = json.loads(zlib.decompress(msg.artifacts).decode("utf-8"))
        except Exception as e:
            logger.exception("Error decompressing artifacts for message id: %s", e)
            artifacts = {}
    # Build the schema object, including computed fields
    message_data = {
        "id": msg.id,
        "role": msg.role,
        "content": msg.content,
        "event_log": msg.event_log,
        "has_event_log": bool(msg.event_log),
        "actions": msg.actions or {},
        "artifacts_metadata": artifacts,
        "merged_context": msg.merged_context or {},
        "selected_agent_ids": msg.selected_agent_ids,
        "used_agent_ids": msg.used_agent_ids,
        "llm_id": msg.llm_id,
        "agents_enabled": msg.agents_enabled,
        "status": msg.status,
        "created_at": msg.created_at,
        "tool_validation_requests": msg.tool_validation_requests,
        "tool_validation_responses": msg.tool_validation_responses,
        "memory_fragment": msg.memory_fragment,
    }

    # Attachments
    msg_attachements = _get_message_attachments(msg.id)
    if msg_attachements:
        message_data["attachments"] = msg_attachements.attachments
        message_data["extraction_mode"] = msg_attachements.extraction_mode
        message_data["quota_exceeded"] = msg_attachements.quota_exceeded
    # Feedback
    fb = {}
    if msg.feedback_rating is not None:
        fb = {
            "rating": int(msg.feedback_rating),
            "by": msg.feedback_by,
            "updatedAt": msg.feedback_updated_at,
        }
        if msg.feedback_text:
            fb["text"] = msg.feedback_text
    message_data["feedback"] = fb
    return schemas.MessageRead(**message_data)


def get_message(message_id: str, user_id: str) -> schemas.MessageRead | None:
    with db.session.begin():
        msg = _get_db_message(message_id, user_id)
        if not msg:
            return None
        return build_message(msg)


def create_message(message: schemas.MessageCreate) -> Message:
    with db.session.begin():
        message_dict = _sanitize_merged_ctx(message.model_dump(exclude_unset=True))

        # Handle artifacts
        artifacts = message_dict.pop("artifacts", None)
        if artifacts:
            compressed_artifacts = zlib.compress(json.dumps(artifacts).encode("utf-8"))
            message_dict["artifacts"] = compressed_artifacts
        # Handle trace
        if "trace" in message_dict:
            trace = message_dict.pop("trace", None)
            if trace:
                compressed_trace = zlib.compress(json.dumps(trace).encode("utf-8"))
                message_dict["trace"] = compressed_trace
        new_message = Message(**message_dict)

        db.session.add(new_message)
        # Update parent conversation to trigger auto-update of timestamps
        from backend.database.models import Conversation

        db.session.query(Conversation).filter(Conversation.conversation_id == message.conversation_id).update(
            {"conversation_id": message.conversation_id}
        )
        db.session.flush()
        return new_message


def update_message(message_id: str, user_id: str, message_update: schemas.MessageUpdate) -> Message:
    with db.session.begin():
        db_message = _get_db_message(message_id, user_id)
        if not db_message:
            return None

        update_data = _sanitize_merged_ctx(message_update.model_dump(exclude_unset=True))
        # Handle artifacts
        if "artifacts" in update_data:
            artifacts = update_data["artifacts"]
            if artifacts:
                compressed_artifacts = zlib.compress(json.dumps(artifacts).encode("utf-8"))
                update_data["artifacts"] = compressed_artifacts
            else:
                update_data["artifacts"] = None

        if "trace" in update_data:
            trace = update_data["trace"]
            if trace:
                compressed_trace = zlib.compress(json.dumps(trace).encode("utf-8"))
                update_data["trace"] = compressed_trace
            else:
                update_data["trace"] = None
        for key, value in update_data.items():
            setattr(db_message, key, value)

        db.session.flush()
        return db_message


def append_tool_validation_response(message_id: str, user_id: str, response: dict) -> list[dict] | None:
    """Atomically append a response unless its validation request was already answered."""
    with db.session.begin():
        message = _get_db_message(message_id, user_id, for_update=True)
        if not message:
            raise ValueError("Tool validation message not found")
        responses = add_validation_response(
            message.tool_validation_requests or [],
            message.tool_validation_responses or [],
            response,
        )
        if responses is None:
            return None
        message.tool_validation_responses = responses
        db.session.flush()
        return responses


def update_message_artifacts_batch(
    conversation_id: str,
    user_id: str,
    artifacts_by_message_id: dict[str, dict],
) -> None:
    """Replace artifact blobs for several user-owned messages in one transaction."""

    with db.session.begin():
        db_messages = _get_conversation_messages_for_artifact_updates(
            conversation_id,
            user_id,
            set(artifacts_by_message_id),
        )
        _apply_message_artifact_updates(db_messages, artifacts_by_message_id)


def delete_message(message_id: str, user_id: str):
    with db.session.begin():
        db_message = _get_db_message(message_id, user_id)
        if db_message:
            db.session.delete(db_message)
            return True
        return False


def update_message_feedback(message_id: str, user_id: str, rating: int | None, text: str | None) -> None:
    """
    Update feedback for a message, enforcing user ownership.
    """
    with db.session.begin():
        msg = _get_db_message(message_id=message_id, user_id=user_id)
        if not msg:
            raise PermissionError("Message not found or not owned by user")
        msg.feedback_rating = rating if rating in (0, 1) else None
        msg.feedback_text = text or None
        msg.feedback_by = user_id
        # feedback_updated_at will be auto-updated by SQLAlchemy


def clear_message_feedback(message_id: str, user_id: str) -> None:
    """
    Clear feedback fields for a message, enforcing user ownership.
    """
    with db.session.begin():
        msg = _get_db_message(message_id=message_id, user_id=user_id)
        if not msg:
            raise PermissionError("Message not found or not owned by user")
        msg.feedback_rating = None
        msg.feedback_text = None
        msg.feedback_by = None
        # feedback_updated_at will be auto-updated by SQLAlchemy


def _get_conversation_messages_for_artifact_updates(
    conversation_id: str,
    user_id: str,
    message_ids: set[str],
) -> dict[str, Message]:
    conversation = (
        db.session.query(Conversation)
        .filter(Conversation.conversation_id == conversation_id, Conversation.user_id == user_id)
        .first()
    )
    if conversation is None:
        raise PermissionError("Conversation not found or not owned by the user")

    if not message_ids:
        return {}

    db_messages = (
        db.session.query(Message)
        .filter(
            Message.conversation_id == conversation_id,
            Message.status == "active",
            Message.id.in_(message_ids),
        )
        .all()
    )
    messages_by_id = {message.id: message for message in db_messages}
    if set(messages_by_id) != message_ids:
        raise PermissionError("One or more artifact messages do not belong to the conversation")

    return messages_by_id


def _apply_message_artifact_updates(
    messages_by_id: dict[str, Message], artifacts_by_message_id: dict[str, dict]
) -> None:
    if not artifacts_by_message_id:
        return

    for message_id, artifacts in artifacts_by_message_id.items():
        db_message = messages_by_id.get(message_id)
        if db_message is None:
            raise PermissionError("One or more artifact messages do not belong to the conversation")
        db_message.artifacts = zlib.compress(json.dumps(artifacts).encode("utf-8")) if artifacts else None


def _append_messages_in_transaction(conv_id: str, messages: list[schemas.MessageCreate]) -> None:
    new_messages = []
    agent_links = []
    for msg in messages:
        msg_data = _sanitize_merged_ctx(msg.model_dump(exclude_unset=True))
        artifacts = msg_data.pop("artifacts", None)
        if artifacts is not None:
            msg_data["artifacts"] = zlib.compress(json.dumps(artifacts).encode("utf-8"))
        trace = msg_data.pop("trace", None)
        if trace is not None:
            msg_data["trace"] = zlib.compress(json.dumps(trace).encode("utf-8"))
        msg_data.pop("attachments", None)
        new_message = Message(**msg_data)
        new_messages.append(new_message)
        selected_ids = msg_data.get("selected_agent_ids") or []
        used_ids = msg_data.get("used_agent_ids") or []
        logger.info("Selected IDs: %s, Used IDs: %s", selected_ids, used_ids)
        seen = set(selected_ids) | set(used_ids)
        for aid in seen:
            agent_links.append(
                MessageAgent(
                    message_id=msg_data["id"],
                    agent_id=aid,
                    selected=1 if aid in set(selected_ids) else 0,
                    used=1 if aid in set(used_ids) else 0,
                )
            )

    db.session.add_all(new_messages)
    db.session.add_all(agent_links)
    db.session.query(Conversation).filter(Conversation.conversation_id == conv_id).update({"conversation_id": conv_id})


def append_messages(conv_id: str, messages: list[schemas.MessageCreate]) -> None:
    """Append multiple messages in one transaction."""
    with db.session.begin():
        _append_messages_in_transaction(conv_id, messages)


def append_messages_with_artifact_updates(
    conv_id: str,
    messages: list[schemas.MessageCreate],
    *,
    user_id: str,
    artifacts_by_message_id: dict[str, dict],
) -> None:
    """Atomically update existing artifact blobs and append new conversation messages."""

    with db.session.begin():
        db_messages = _get_conversation_messages_for_artifact_updates(
            conv_id,
            user_id,
            set(artifacts_by_message_id),
        )
        _apply_message_artifact_updates(db_messages, artifacts_by_message_id)
        _append_messages_in_transaction(conv_id, messages)


def get_message_events(message_id: str, user_id: str | None = None) -> list:
    """
    Return the event log for a message, enforcing user ownership if user_id is provided.
    """
    with db.session.begin():
        query = db.session.query(Message.event_log)
        if user_id:
            query = query.join(Conversation, Message.conversation_id == Conversation.conversation_id)
            query = query.filter(Message.id == message_id, Conversation.user_id == user_id, Message.status == "active")
            authorized = query.first()
            if not authorized:
                return []
            event_log = authorized[0]
        else:
            result = db.session.query(Message.event_log).filter(Message.id == message_id).first()
            if not result:
                return []
            event_log = result[0]
        return event_log if event_log else []


def get_message_trace(message_id: str, user_id: str | None = None) -> str:
    """
    Return the trace string for a message, enforcing user ownership if user_id is provided.
    """
    with db.session.begin():
        msg = _get_db_message(message_id, user_id)
        if not msg or not msg.trace:
            return ""
        try:
            trace = json.loads(zlib.decompress(msg.trace).decode("utf-8"))
        except Exception as e:
            logger.exception("Error decompressing trace for message id: %s", e)
            trace = ""
        return trace


# Document uploads
def insert_or_update_message_attachments(
    message_id: str,
    attachments: list[dict],
    extraction_mode: dict | None = None,
) -> None:
    """
    Insert or update message attachments metadata for a given message_id.
    """
    from backend.database.models import MessageAttachment

    with db.session.begin():
        # attachments = schemas.MessageAttachmentList.model_validate(attachments) if hasattr(schemas, 'MessageAttachmentList') else attachments
        obj = db.session.query(MessageAttachment).filter_by(message_id=message_id).first()
        if obj:
            obj.attachments = attachments
            # Always store as JSON (only mode and quota_exceeded present)
            obj.extraction_mode = extraction_mode
        else:
            obj = MessageAttachment(message_id=message_id, attachments=attachments)
            db.session.add(obj)
