import logging
from typing import (
    Optional,
    Union,
)
from uuid import uuid4

from pydantic import (
    UUID4,
    ValidationError,
)
from sqlalchemy import select

from galaxy.config import GalaxyAppConfiguration
from galaxy.config.url_headers import UrlHeadersConfigFactory
from galaxy.exceptions import (
    InsufficientPermissionsException,
    ItemAlreadyClaimedException,
    ItemMustBeClaimed,
    ObjectNotFound,
    RequestParameterInvalidException,
    RequestParameterMissingException,
)
from galaxy.managers.workflows import WorkflowContentsManager
from galaxy.model import (
    ToolLandingRequest as ToolLandingRequestModel,
    WorkflowLandingRequest as WorkflowLandingRequestModel,
)
from galaxy.model.dataset_collections.types.sample_sheet_util import (
    validate_column_definitions,
    validate_row,
)
from galaxy.model.scoped_session import galaxy_scoped_session
from galaxy.schema.schema import (
    ClaimLandingPayload,
    CreateToolLandingRequestPayload,
    CreateWorkflowLandingRequestPayload,
    LandingRequestState,
    ToolLandingRequest,
    WorkflowLandingRequest,
)
from galaxy.security.idencoding import IdEncodingHelper
from galaxy.security.vault import (
    InvalidVaultConfigException,
    Vault,
)
from galaxy.structured_app import (
    MinimalManagerApp,
    StructuredApp,
)
from galaxy.tool_util.parameters import (
    landing_decode,
    LandingRequestInternalToolState,
    LandingRequestToolState,
)
from galaxy.tool_util_models.parameters import (
    DataOrCollectionRequestAdapter,
    DataRequestCollectionUri,
    ToolParameterBundleModel,
)
from galaxy.util import safe_str_cmp
from .context import ProvidesUserContext
from .headers_encryption import (
    decrypt_headers_in_data,
    encrypt_headers_in_data,
    has_sensitive_headers,
)
from .tools import (
    get_tool_from_toolbox,
    ToolRunReference,
)

LandingRequestModel = Union[ToolLandingRequestModel, WorkflowLandingRequestModel]

FETCH_TOOL_ID = "__DATA_FETCH__"

log = logging.getLogger(__name__)


class LandingRequestManager:

    def __init__(
        self,
        sa_session: galaxy_scoped_session,
        security: IdEncodingHelper,
        workflow_contents_manager: WorkflowContentsManager,
        app: MinimalManagerApp,
        config: GalaxyAppConfiguration,
        vault: Optional[Vault] = None,
    ):
        self.sa_session = sa_session
        self.security = security
        self.workflow_contents_manager = workflow_contents_manager
        self.app = app
        self.vault = vault
        self.url_headers_config = UrlHeadersConfigFactory.from_app_config(config)

    def create_tool_landing_request(self, payload: CreateToolLandingRequestPayload, user_id=None) -> ToolLandingRequest:
        tool_id = payload.tool_id
        tool_version = payload.tool_version
        request_state = payload.request_state

        ref = ToolRunReference(tool_id=tool_id, tool_version=tool_version, tool_uuid=None)
        tool = get_tool_from_toolbox(self.app.toolbox, ref, user=None)
        landing_request_state = LandingRequestToolState(request_state or {})
        # Okay this is a hack until tool request API commit is merged, tools don't yet have a parameter
        # schema - so we can't do this properly.
        if tool.parameters is not None:
            parameter_bundle = ToolParameterBundleModel(parameters=tool.parameters)
            internal_landing_request_state = landing_decode(
                landing_request_state, parameter_bundle, self.security.decode_id
            )
        else:
            assert tool.id == FETCH_TOOL_ID, f"tool '{tool.id}' has no parameter schema but is not {FETCH_TOOL_ID}"
            # we have validated the payload as part of the API request
            # nothing else to decode ideally so just swap to internal model state object
            internal_landing_request_state = LandingRequestInternalToolState(
                input_state=landing_request_state.input_state
            )

        # Validate sample sheet metadata in request_state for __DATA_FETCH__ tool
        if tool.id == "__DATA_FETCH__" and request_state:

            # Check each item in request_state for sample sheet metadata
            for item in landing_request_state.input_state.get("request_state", []):
                # Try to parse as DataRequestCollectionUri to access sample sheet fields
                if isinstance(item, dict) and item.get("class") == "Collection":
                    column_definitions = item.get("column_definitions")
                    rows = item.get("rows")
                    collection_type = item.get("collection_type", "")

                    if column_definitions is not None or rows is not None:
                        # Validate that sample sheet metadata is only used with sample_sheet collection types
                        if not collection_type.startswith("sample_sheet"):
                            raise RequestParameterInvalidException(
                                f"Sample sheet metadata (column_definitions, rows) can only be used with collection_type 'sample_sheet' or 'sample_sheet:<type>', not '{collection_type}'"
                            )

                        # Validate column definitions structure
                        if column_definitions is not None:
                            validate_column_definitions(column_definitions)

                        # Validate rows against column definitions and element identifiers
                        if rows:
                            element_identifiers = [elem.get("identifier") for elem in item.get("elements", [])]
                            for identifier, row in rows.items():
                                if identifier not in element_identifiers:
                                    raise RequestParameterInvalidException(
                                        f"Row identifier '{identifier}' not found in collection elements"
                                    )
                                validate_row(row, column_definitions, element_identifiers)

        model = ToolLandingRequestModel()
        model.tool_id = tool_id
        model.tool_version = tool_version
        model.uuid = uuid4()
        model.client_secret = payload.client_secret
        model.public = payload.public
        model.origin = str(payload.origin) if payload.origin else None
        if user_id:
            model.user_id = user_id

        request_state = self._encrypt_headers_in_request_state(
            internal_landing_request_state.input_state, str(model.uuid)
        )
        model.request_state = request_state

        self._save(model)
        return self._tool_response(model)

    def create_workflow_landing_request(self, payload: CreateWorkflowLandingRequestPayload) -> WorkflowLandingRequest:
        model = WorkflowLandingRequestModel()
        if payload.workflow_target_type == "stored_workflow":
            model.stored_workflow_id = self.security.decode_id(payload.workflow_id)
        elif payload.workflow_target_type == "workflow":
            model.workflow_id = self.security.decode_id(payload.workflow_id)
        elif payload.workflow_target_type == "trs_url":
            model.workflow_source_type = "trs_url"
            # validate this ?
            model.workflow_source = payload.workflow_id
        elif payload.workflow_target_type == "url":
            model.workflow_source_type = "url"
            model.workflow_source = payload.workflow_id
        model.uuid = uuid4()
        model.client_secret = payload.client_secret

        validated_request_state = self.validate_workflow_request_state(payload.request_state)
        request_state = self._encrypt_headers_in_request_state(validated_request_state, str(model.uuid))
        model.request_state = request_state

        model.public = payload.public
        model.origin = str(payload.origin) if payload.origin else None
        self._save(model)
        return self._workflow_response(model)

    def validate_workflow_request_state(self, request_state: Optional[dict]) -> Optional[dict]:
        # This would ideally be run in the context of a workflow input definition
        if isinstance(request_state, dict):
            for key, value in request_state.items():
                if isinstance(value, dict):
                    try:
                        # persist values after model validators and aliases have been applied
                        validated_value = DataOrCollectionRequestAdapter.validate_python(value)

                        # Validate sample sheet metadata for collections
                        if isinstance(validated_value, DataRequestCollectionUri):
                            has_sample_sheet_metadata = (
                                validated_value.column_definitions is not None or validated_value.rows is not None
                            )
                            if has_sample_sheet_metadata:
                                collection_type = validated_value.collection_type
                                if not collection_type.startswith("sample_sheet"):
                                    raise RequestParameterInvalidException(
                                        f"Sample sheet metadata (column_definitions, rows) can only be used with collection_type 'sample_sheet' or 'sample_sheet:<type>', not '{collection_type}'"
                                    )

                                # Validate column definitions structure
                                if validated_value.column_definitions is not None:
                                    validate_column_definitions(validated_value.column_definitions)

                                # Validate rows against column definitions and element identifiers
                                if validated_value.rows:
                                    element_identifiers = [elem.identifier for elem in validated_value.elements]
                                    for identifier, row in validated_value.rows.items():
                                        if identifier not in element_identifiers:
                                            raise RequestParameterInvalidException(
                                                f"Row identifier '{identifier}' not found in collection elements"
                                            )
                                        validate_row(row, validated_value.column_definitions, element_identifiers)

                        request_state[key] = validated_value.model_dump(by_alias=True, exclude_unset=True, mode="json")
                    except ValidationError:
                        pass
        return request_state

    def claim_tool_landing_request(
        self, trans: ProvidesUserContext, uuid: UUID4, claim: Optional[ClaimLandingPayload]
    ) -> ToolLandingRequest:
        request = self._get_tool_landing_request(uuid)
        self._check_can_claim(trans, request, claim)
        request.user_id = trans.user.id
        self._save(request)
        return self._tool_response(request)

    def claim_workflow_landing_request(
        self, trans: ProvidesUserContext, uuid: UUID4, claim: Optional[ClaimLandingPayload]
    ) -> WorkflowLandingRequest:
        request = self._get_workflow_landing_request(uuid)
        self._check_can_claim(trans, request, claim)
        self._ensure_workflow(trans, request)
        request.user_id = trans.user.id
        self._save(request)
        return self._workflow_response(request)

    def _ensure_workflow(self, trans: ProvidesUserContext, request: WorkflowLandingRequestModel):
        if request.workflow_source_type == "trs_url" and isinstance(trans.app, StructuredApp):
            # trans is always structured app except for unit test
            assert request.workflow_source
            workflow = self.workflow_contents_manager.get_or_create_workflow_from_trs(
                trans, trs_url=request.workflow_source
            )
            request.workflow_id = workflow.latest_workflow_id
        elif request.workflow_source_type == "url" and isinstance(trans.app, StructuredApp):
            # trans is always structured app except for unit test
            assert request.workflow_source
            workflow = self.workflow_contents_manager.get_or_create_workflow_from_url(
                trans, url=request.workflow_source
            )
            request.workflow_id = workflow.latest_workflow_id

    def get_tool_landing_request(self, trans: ProvidesUserContext, uuid: UUID4) -> ToolLandingRequest:
        request = self._get_claimed_tool_landing_request(trans, uuid)
        return self._tool_response(request)

    def get_workflow_landing_request(self, trans: ProvidesUserContext, uuid: UUID4) -> WorkflowLandingRequest:
        request = self._get_claimed_workflow_landing_request(trans, uuid)
        self._ensure_workflow(trans, request)
        return self._workflow_response(request)

    def _check_can_claim(
        self, trans: ProvidesUserContext, request: LandingRequestModel, claim: Optional[ClaimLandingPayload]
    ):
        if request.client_secret is not None:
            if claim is None or not claim.client_secret:
                raise RequestParameterMissingException()
            if not safe_str_cmp(request.client_secret, claim.client_secret):
                raise InsufficientPermissionsException()
        if request.user_id is not None:
            raise ItemAlreadyClaimedException()

    def _get_tool_landing_request(self, uuid: UUID4) -> ToolLandingRequestModel:
        request = self.sa_session.scalars(
            select(ToolLandingRequestModel).where(ToolLandingRequestModel.uuid == str(uuid))
        ).one_or_none()
        if request is None:
            raise ObjectNotFound()
        return request

    def _get_workflow_landing_request(self, uuid: UUID4) -> WorkflowLandingRequestModel:
        request = self.sa_session.scalars(
            select(WorkflowLandingRequestModel).where(WorkflowLandingRequestModel.uuid == str(uuid))
        ).one_or_none()
        if request is None:
            raise ObjectNotFound()
        return request

    def _get_claimed_tool_landing_request(self, trans: ProvidesUserContext, uuid: UUID4) -> ToolLandingRequestModel:
        request = self._get_tool_landing_request(uuid)
        self._check_ownership(trans, request)
        return request

    def _get_claimed_workflow_landing_request(
        self, trans: ProvidesUserContext, uuid: UUID4
    ) -> WorkflowLandingRequestModel:
        request = self._get_workflow_landing_request(uuid)
        self._check_ownership(trans, request)
        return request

    def _tool_response(self, model: ToolLandingRequestModel) -> ToolLandingRequest:
        request_state = self._decrypt_headers_in_request_state(model.request_state, str(model.uuid))

        response_model = ToolLandingRequest(
            tool_id=model.tool_id,
            tool_version=model.tool_version,
            request_state=request_state,
            uuid=model.uuid,
            state=self._state(model),
            origin=model.origin,
        )
        return response_model

    def _workflow_response(self, model: WorkflowLandingRequestModel) -> WorkflowLandingRequest:

        workflow_id: Optional[Union[int, str]] = None
        if model.stored_workflow_id is not None:
            workflow_id = model.stored_workflow_id
            target_type = "stored_workflow"
        elif model.workflow_id is not None:
            workflow_id = model.workflow_id
            target_type = "workflow"
        elif model.workflow_source_type in ("trs_url", "url"):
            target_type = model.workflow_source_type
            workflow_id = model.workflow_source
        assert workflow_id

        request_state = self._decrypt_headers_in_request_state(model.request_state, str(model.uuid))

        response_model = WorkflowLandingRequest(
            workflow_id=self.security.encode_id(workflow_id) if isinstance(workflow_id, int) else workflow_id,
            workflow_target_type=target_type,
            request_state=request_state,
            uuid=model.uuid,
            state=self._state(model),
            origin=model.origin,
        )
        return response_model

    def _check_ownership(self, trans: ProvidesUserContext, model: LandingRequestModel):
        if not model.public and self._state(model) == LandingRequestState.UNCLAIMED:
            raise ItemMustBeClaimed
        if model.user_id and model.user_id != trans.user.id:
            raise InsufficientPermissionsException()

    def _state(self, model: LandingRequestModel) -> LandingRequestState:
        return LandingRequestState.UNCLAIMED if model.user_id is None else LandingRequestState.CLAIMED

    def _save(self, model: LandingRequestModel):
        sa_session = self.sa_session
        sa_session.add(model)
        sa_session.commit()

    def _encrypt_headers_in_request_state(self, request_state: Optional[dict], landing_uuid: str) -> Optional[dict]:
        if request_state is not None:
            if has_sensitive_headers(request_state, self.url_headers_config):
                if not self.vault:
                    raise InvalidVaultConfigException(
                        "Sensitive headers detected in landing request but no vault is configured. "
                        "Configure a vault to securely store sensitive header information."
                    )
                return encrypt_headers_in_data(
                    request_state,
                    landing_uuid,
                    self.vault,
                    key_prefix="landing_request/headers",
                    url_headers_config=self.url_headers_config,
                )
        return request_state

    def _decrypt_headers_in_request_state(self, request_state: Optional[dict], landing_uuid: str):
        if request_state is not None and self.vault:
            return decrypt_headers_in_data(
                request_state,
                landing_uuid,
                self.vault,
                key_prefix="landing_request/headers",
                url_headers_config=self.url_headers_config,
            )
        return request_state
