import copy

from sqlalchemy.orm.attributes import flag_modified

from backend import schemas
from backend.database.agenthub_models import DraftConversation
from backend.database.base import db
from backend.utils.tool_validation import add_validation_response


def get_draft_conversation(agent_id: str, user_id: str) -> schemas.DraftConversationRead | None:
    with db.session.begin():
        draft = (
            db.session.query(DraftConversation)
            .filter(DraftConversation.agent_id == agent_id, DraftConversation.user_id == user_id)
            .first()
        )
        if draft:
            return schemas.DraftConversationRead.model_validate(draft)
        return None


def create_draft_conversation(draft_conversation: schemas.DraftConversationCreate) -> schemas.DraftConversationRead:
    with db.session.begin():
        new_draft_conversation = DraftConversation(**draft_conversation.model_dump())
        db.session.add(new_draft_conversation)
        db.session.flush()
        return schemas.DraftConversationRead.model_validate(new_draft_conversation)


def upsert_draft_conversation(agent_id: str, user_id: str, convo: dict) -> None:
    """
    Insert or update a draft conversation for (agent_id, user_id).
    Sets created_at and last_modified timestamps as in store.py.
    """
    with db.session.begin():
        draft = (
            db.session.query(DraftConversation)
            .filter(DraftConversation.agent_id == agent_id, DraftConversation.user_id == user_id)
            .first()
        )
        if draft:
            draft.convo = dict(convo)
            flag_modified(draft, "convo")
        else:
            draft = DraftConversation(
                agent_id=agent_id,
                user_id=user_id,
                convo=convo,
            )
            db.session.add(draft)


def append_draft_tool_validation_response(agent_id: str, user_id: str, response: dict) -> dict | None:
    """Atomically append a validation response to the latest draft message."""
    with db.session.begin():
        draft = (
            db.session.query(DraftConversation)
            .filter(DraftConversation.agent_id == agent_id, DraftConversation.user_id == user_id)
            .with_for_update()
            .first()
        )
        if not draft:
            raise ValueError("Draft conversation not found")

        convo = copy.deepcopy(draft.convo or {})
        messages = convo.get("messages") or []
        if not messages:
            raise ValueError("Tool validation message not found")
        message = messages[-1]
        responses = add_validation_response(
            message.get("toolValidationRequests") or [],
            message.get("toolValidationResponses") or [],
            response,
        )
        if responses is None:
            return None

        message["toolValidationResponses"] = responses
        draft.convo = convo
        flag_modified(draft, "convo")
        db.session.flush()
        return convo


def delete_draft_conversation(agent_id: str, user_id: str) -> None:
    """
    Delete the draft conversation for the given agent_id and user_id.
    """
    with db.session.begin():
        db.session.query(DraftConversation).filter(
            DraftConversation.agent_id == agent_id, DraftConversation.user_id == user_id
        ).delete()
