import uuid

from backend.database import crud
from backend.database.crud.agent import get_agent  # _get_full_agent_details
from backend.database.user_store_protocol import IUserStore
from backend.schemas import schemas
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


class UserStore(IUserStore):
    """
    Wraps any Store implementation and scopes all reads/writes
    to the current user (by user_id).
    """

    def __init__(self, user_id: str, groups: list[str]):
        # # TODO SQLAlchemy/Pydantic Should we throw an exception if user is not defined
        # -> No because in backend init we create a UserStore without user id
        # if(not user_id):
        #     raise ValueError("User store requires a valid user_id")
        self._user = user_id
        self._groups = groups or []

    def _build_published_agent(self, agent: schemas.AgentRead | dict) -> schemas.AgentRead | None:
        """Project an accessible agent to the published snapshot exposed to non-owners."""

        # Convert dict to AgentRead if needed
        if isinstance(agent, dict):
            agent = schemas.AgentRead.model_validate(agent)

        pv = get_agent(agent.id).published_version

        if not pv:
            return None

        return agent.model_copy(update={**pv.model_dump(), "published_version": pv})

    def _build_published_agent_summary(self, agent: schemas.AgentRead | dict) -> schemas.AgentRead | None:
        """Project a light agent summary to its published snapshot for shared-access flows."""
        if isinstance(agent, dict):
            agent = schemas.AgentRead.model_validate(agent)

        pv = agent.published_version
        if not pv:
            return None

        return agent.model_copy(update={**pv.model_dump(), "published_version": pv})

    def _build_published_agent_from_loaded(self, agent: schemas.AgentRead) -> schemas.AgentRead | None:
        """Project an already-loaded full agent to the published payload exposed to non-owners."""
        pv = agent.published_version
        if not pv:
            return None

        return agent.model_copy(update={**pv.model_dump(), "published_version": pv})

    def _user_groups(self) -> list[str]:
        return self._groups

    def get_agents_by_owner(self, owner: str) -> list[schemas.AgentRead]:
        # Only allow querying your own agents
        if owner != self._user:
            return []
        return crud.get_agents_by_owner(self._user)

    # --- User-scoped Agent methods ---
    def get_agent(self, agent_id: str) -> schemas.AgentRead | None:
        """Get agent if user owns it or it's shared with them"""
        agent = crud.get_agent(agent_id)
        if not agent:
            logger.info("Agent %s not found", agent_id)
            return None
        logger.info("Found agent %s owned by %s", agent_id, agent.owner)
        # Owner can see everything
        if agent.owner == self._user:
            return agent

        # Check if shared with user
        shares = crud.get_agent_shares(agent_id)
        if not shares:
            logger.info("Agent %s is not shared with anyone", agent_id)
            return None
        shared_with_user = any(
            s.principal_type == schemas.PrincipalTypeEnum.USER and s.principal == self._user for s in shares
        )
        shared_with_group = any(
            s.principal_type == schemas.PrincipalTypeEnum.GROUP and s.principal in self._groups for s in shares
        )
        if shared_with_user or shared_with_group:
            return self._build_published_agent(agent)

        # Not authorized
        return None

    def get_all_agents(self, all_users: bool = False) -> list[schemas.AgentRead]:
        """
        Union of:
            • agents I own
            • agents shared with me (filtering out drafts)
        """
        if all_users:
            owned = crud.get_all_agents()
        else:
            owned = crud.get_agents_by_owner(self._user)

        shared: list[dict] = crud.get_agents_shared_with(self._user, self._user_groups())

        # For shared agents, show only the published version data
        processed_shared = []
        for agent in shared:
            pa = self._build_published_agent(agent)
            if pa is not None:
                processed_shared.append(pa)

        return owned + processed_shared

    def get_user_agent_summaries_by_ids(self, agent_ids: list[str]) -> dict[str, schemas.AgentRead]:
        """Return accessible user-agent summaries enriched with draft/published identity fields for routing."""
        if not agent_ids:
            return {}

        summaries: dict[str, schemas.AgentRead] = {}
        requested_ids = list(dict.fromkeys(agent_ids))
        owned_ids = {agent.id for agent in crud.get_agent_summaries_by_ids(requested_ids) if agent.owner == self._user}
        shared_ids = {
            agent["id"] for agent in crud.get_agents_shared_with_by_ids(self._user, self._user_groups(), requested_ids)
        }
        accessible_ids = [agent_id for agent_id in requested_ids if agent_id in owned_ids or agent_id in shared_ids]
        enriched_agents_by_id = {
            agent.id: agent for agent in crud.get_agents_identity_and_publication_details_by_ids(accessible_ids)
        }

        for agent_id in accessible_ids:
            agent = enriched_agents_by_id.get(agent_id)
            if agent is None:
                continue
            if agent.owner == self._user:
                summaries[agent_id] = agent
                continue
            summary = self._build_published_agent_summary(agent)
            if summary is not None:
                summaries[summary.id] = summary

        return summaries

    def get_user_agents_by_ids(self, agent_ids: list[str]) -> dict[str, schemas.AgentRead]:
        """Return full user-agent payloads for the specific ids that will actually execute."""
        if not agent_ids:
            return {}

        requested_ids = list(dict.fromkeys(agent_ids))
        resolved: dict[str, schemas.AgentRead] = {}
        loaded_by_id = {agent.id: agent for agent in crud.get_agents_by_ids(requested_ids)}

        for agent_id, agent in loaded_by_id.items():
            if agent.owner == self._user:
                resolved[agent_id] = agent

        shared_ids = {
            agent["id"] for agent in crud.get_agents_shared_with_by_ids(self._user, self._user_groups(), requested_ids)
        }
        for agent_id in requested_ids:
            if agent_id in resolved or agent_id not in shared_ids:
                continue
            agent = loaded_by_id.get(agent_id)
            if agent is None:
                continue
            published_agent = self._build_published_agent_from_loaded(agent)
            if published_agent is not None:
                resolved[agent_id] = published_agent

        return resolved

    def create_agent(self, agent_dict) -> schemas.AgentRead:
        agent_data = dict(agent_dict)
        agent_data["owner"] = self._user
        if not agent_data.get("id"):
            agent_data["id"] = str(uuid.uuid4())
        agent_create = schemas.AgentCreate(**agent_data)
        return crud.create_agent(agent_create)

    def update_agent(self, agent_id, update_data, bypass_user=False) -> schemas.AgentRead | None:
        # Check if user owns the agent unless bypass_user is True
        # (bypass_user is used by indexing monitor)
        agent = crud.get_agent(agent_id)
        if not agent or (not bypass_user and agent.owner != self._user):
            return None
        # Make sure id and owner are not updated
        agent_data = agent.model_dump(exclude={"id", "owner", "created_at", "last_modified"})
        agent_data.update(update_data)
        agent_update = schemas.AgentUpdate(**agent_data)
        return crud.update_agent(agent_id, agent_update)

    def delete_agent(self, agent_id) -> bool:
        # Check if user owns the agent
        agent = crud.get_agent(agent_id)
        if not agent or agent.owner != self._user:
            return False
        return crud.delete_agent(agent_id)

    # ---------------- sharing ---------------------------------------
    def replace_agent_shares(self, agent_id: str, shares: list[dict]) -> None:
        # Only the owner may modify sharing
        agent = crud.get_agent(agent_id)
        if not agent or agent.owner != self._user:
            raise PermissionError("Only the owner can share this agent")
        # Validate shares payload
        for s in shares:
            if "type" not in s or "principal" not in s:
                raise ValueError("Invalid share entry; expected keys: 'type', 'principal'")
            if s["type"] not in ("user", "group"):
                raise ValueError(f"Invalid share type: {s['type']}")
        crud.replace_agent_shares(
            agent_id,
            [
                schemas.AgentShareBase(
                    principal=s["principal"],
                    principal_type=(
                        schemas.PrincipalTypeEnum.USER if s["type"] == "user" else schemas.PrincipalTypeEnum.GROUP
                    ),
                )
                for s in shares
            ],
        )

    def get_agent_shares(self, agent_id: str) -> list[schemas.AgentShareRead]:
        return crud.get_agent_shares(agent_id) or []

    def get_agents_shared_with(self) -> list[schemas.AgentRead]:
        """
        Forward to the base-store,
        """
        return crud.get_agents_shared_with(self._user, self._user_groups())

    def get_share_counts(self, agent_ids: list[str]) -> dict[str, int]:
        # notice: count shares for any agent, not just those I own
        counts = crud.get_share_counts(agent_ids)
        return {aid: counts.get(aid, 0) for aid in agent_ids}

    # --- User-scoped Conversation methods ---

    def get_conversations(self) -> list[schemas.ConversationRead]:
        return crud.get_conversations_by_user(self._user)

    def get_conversations_ids(self) -> list[str]:
        return crud.get_conversations_ids_by_user(self._user)

    def get_conversations_metadata(self) -> list[schemas.ConversationMetadata]:
        return crud.get_conversations_metadata(self._user)

    def update_conversation(self, conversation_id, conversation_obj) -> schemas.ConversationRead | None:
        conversation_obj = dict(conversation_obj)
        conversation_obj["user_id"] = self._user
        conversation_update = schemas.ConversationUpdate(**conversation_obj)
        return crud.update_conversation(conversation_id, conversation_update)

    def delete_conversation(self, conversation_id, permanent_delete: bool = False) -> bool:
        return crud.delete_conversation(conversation_id, self._user, permanent_delete)

    def delete_all_conversations(self, permanent_delete: bool = False) -> None:
        return crud.delete_all_conversations(self._user, permanent_delete)

    def get_conversation(self, conversation_id: str) -> schemas.ConversationRead | None:
        return crud.get_user_conversation(conversation_id, self._user)

    def create_conversation(self, conv_id: str, new_conv: schemas.ConversationCreate) -> schemas.ConversationRead:
        new_conv.conversation_id = conv_id
        new_conv.user_id = self._user
        # conversation_create = schemas.ConversationCreate(**conversation_obj)
        return crud.create_conversation(new_conv)

    # -------- incremental helpers -------------------------------------
    def ensure_conversation_exists(
        self,
        conv_id: str,
        agent_ids: list[str],
        mode_agents: bool = True,
        llm_id: str | None = None,
        doc_gen_enabled: bool = True,
    ) -> schemas.FullConversationRead:
        conversation = crud.get_user_conversation(conversation_id=conv_id, user_id=self._user)
        if not conversation:
            new_convo = schemas.ConversationCreate(
                conversation_id=conv_id,
                user_id=self._user,
                agent_ids=agent_ids,
                agents_enabled=mode_agents,
                llm_id=llm_id,
                doc_gen_enabled=doc_gen_enabled,
                title="",
            )
            return crud.create_conversation(conversation=new_convo)
        return conversation

    def append_message(self, conversation_id: str, message: dict) -> None:
        new_msg = schemas.MessageCreate(**message, conversation_id=conversation_id)
        crud.create_message(new_msg)

    def get_message(self, message_id) -> schemas.MessageRead | None:
        msg = crud.get_message(message_id, self._user)
        return msg

    def get_message_artifacts_meta(self, message_id) -> schemas.ArtifactsMetadata | None:
        msg = crud.get_message(message_id, self._user)
        return msg.artifacts_metadata if msg else None

    def update_message(self, message_id: str, updates: dict) -> None:
        # SQLAlchemy/Pydantic I added user_id HERE
        crud.update_message(message_id=message_id, user_id=self._user, message_update=schemas.MessageUpdate(**updates))

    def append_tool_validation_response(self, message_id: str, response: dict) -> list[dict] | None:
        return crud.append_tool_validation_response(message_id, self._user, response)

    def append_chart_plan(self, message_id: str, chart_plan_entry: dict) -> dict | None:
        return crud.append_chart_plan(message_id=message_id, user_id=self._user, chart_plan_entry=chart_plan_entry)

    def update_chart_plan(self, message_id: str, *, chart_plan_index: int, chart_data: dict) -> dict:
        return crud.update_chart_plan(
            message_id=message_id,
            user_id=self._user,
            chart_plan_index=chart_plan_index,
            chart_data=chart_data,
        )

    def apply_chart_plan_decision(self, message_id: str, **decision) -> dict:
        return crud.apply_chart_plan_decision(message_id=message_id, user_id=self._user, **decision)

    def update_message_artifacts_batch(
        self,
        conversation_id: str,
        artifacts_by_message_id: dict[str, dict],
    ) -> None:
        crud.update_message_artifacts_batch(
            conversation_id,
            self._user,
            artifacts_by_message_id,
        )

    def update_message_feedback(self, message_id: str, rating: int | None, text: str | None) -> None:
        crud.update_message_feedback(message_id, user_id=self._user, rating=rating, text=text)

    def clear_message_feedback(self, message_id: str) -> None:
        # SQLAlchemy/Pydantic I added user_id HERE
        crud.clear_message_feedback(message_id, user_id=self._user)

    def append_messages(self, conversation_id: str, messages: list[dict]) -> None:
        crud.append_messages(conversation_id, messages)

    def append_messages_with_artifact_updates(
        self,
        conversation_id: str,
        messages: list[dict],
        artifacts_by_message_id: dict[str, dict],
    ) -> None:
        crud.append_messages_with_artifact_updates(
            conversation_id,
            messages,
            user_id=self._user,
            artifacts_by_message_id=artifacts_by_message_id,
        )

    def update_conversation_meta(
        self,
        conversation_id: str,
        *,
        title=None,
        agent_ids=None,
        agents_enabled: bool | None = None,
        llm_id: str | None = None,
        doc_gen_enabled: bool | None = None,
    ) -> None:
        update_fields = {}
        if title is not None:
            update_fields["title"] = title
        if agent_ids is not None:
            update_fields["agent_ids"] = agent_ids
        if agents_enabled is not None:
            update_fields["agents_enabled"] = agents_enabled
        if llm_id is not None:
            update_fields["llm_id"] = llm_id
        if doc_gen_enabled is not None:
            update_fields["doc_gen_enabled"] = doc_gen_enabled
        crud.update_conversation_meta(conversation_id, schemas.ConversationMetadataUpdate(**update_fields))

    # --- Shared across users: message events ---

    def get_message_events(self, message_id: str) -> list[dict]:
        return crud.get_message_events(message_id, self._user)

    # --- Preferences (scoped to me) ---

    def get_preferences(self) -> schemas.PreferenceRead | None:
        return crud.get_preferences(self._user)

    def update_preferences(self, prefs: dict) -> schemas.PreferenceRead | None:
        prefs_update = schemas.PreferenceUpdate(**prefs)
        return crud.update_preferences(self._user, prefs_update)

    # --- Draft-chat helpers ---

    def get_draft_conversation(self, agent_id: str) -> schemas.DraftConversationRead | None:
        return crud.get_draft_conversation(agent_id, self._user)

    def upsert_draft_conversation(self, agent_id: str, convo: dict) -> None:
        return crud.upsert_draft_conversation(agent_id=agent_id, user_id=self._user, convo=convo)

    def append_draft_tool_validation_response(self, agent_id: str, response: dict) -> dict | None:
        return crud.append_draft_tool_validation_response(agent_id, self._user, response)

    def delete_draft_conversation(self, agent_id: str) -> None:
        return crud.delete_draft_conversation(agent_id, self._user)

    # --- Helper method to check existence without permission ---
    def agent_exists(self, agent_id: str) -> bool:
        """Check if agent exists in database (no permission check)"""
        return bool(self.get_agent(agent_id))

    def delete_message(self, message_id: str) -> bool:
        """Delete a message (with user permission check)"""
        # First check if user has permission to delete this message
        message = self.get_message(message_id)
        if not message:
            return False

        # User can only delete messages from their own conversations
        return crud.delete_message(message_id, self._user)

    def list_user_agent_owners(self) -> list[dict]:
        return crud.list_user_agent_owners()

    def get_message_trace(self, message_id: str) -> str | None:
        """
        Retrieve only the trace data for the given message using the provided user id.
        Returns:
            The trace data if available and the message belongs to the user, None otherwise.
        """
        return crud.get_message_trace(message_id, self._user)

    # def get_message_attachments(self, message_id: str) -> list[dict]:
    #     return crud.get_message_attachments(message_id, self._user)

    def insert_or_update_message_attachments(
        self,
        message_id: str,
        attachments_json: str,
        extraction_mode: str | None = None,
        quota_exceeded: bool | None = None,
    ) -> None:
        """Insert or update message attachments."""
        extr_mode = (
            {
                "mode": extraction_mode,
                "quotaExceeded": quota_exceeded if quota_exceeded is not None else False,
            }
            if extraction_mode is not None
            else None
        )
        crud.insert_or_update_message_attachments(message_id, attachments_json, extr_mode)

    def get_derived_documents(self, conv_id: str) -> list[schemas.DerivedDocument]:
        return crud.get_derived_documents(conv_id, self._user)

    def get_derived_document_by_path(self, document_path: str) -> schemas.DerivedDocument | None:
        return crud.get_derived_document_by_path(document_path, self._user)

    def upsert_derived_document(
        self, conv_id: str, document_name: str, document_path: str, metadata: dict | None = None
    ) -> schemas.DerivedDocument:
        return crud.upsert_derived_document(conv_id, self._user, document_name, document_path, metadata)

    def delete_derived_document_by_source_path(self, document_path: str) -> bool:
        return crud.delete_derived_document_by_source_path(document_path, self._user)
