import logging
from datetime import datetime
from typing import List, Optional, Tuple
from uuid import UUID

from arthur_common.models.common_schemas import PaginationParameters
from arthur_common.models.enums import PaginationSortMethod
from fastapi import HTTPException
from sqlalchemy import asc, desc, func
from sqlalchemy.orm import Session

from db_models.rag_provider_models import (
    DatabaseApiKeyRagProviderConfiguration,
    DatabaseRagProviderConfiguration,
    DatabaseRagSearchSettingConfiguration,
    DatabaseRagSearchSettingConfigurationVersion,
    DatabaseRagSearchVersionTag,
)
from db_models.task_models import DatabaseTask
from schemas.enums import (
    RagAPIKeyAuthenticationProviderEnum,
    RagProviderAuthenticationMethodEnum,
)
from schemas.internal_schemas import (
    ApiKeyRagProviderSecretValue,
    RagProviderConfiguration,
    RagSearchSettingConfiguration,
    RagSearchSettingConfigurationVersion,
    RagSearchSettingTag,
)
from schemas.request_schemas import (
    ApiKeyRagAuthenticationConfigUpdateRequest,
    RagProviderConfigurationUpdateRequest,
    RagSearchSettingConfigurationUpdateRequest,
    RagSearchSettingConfigurationVersionUpdateRequest,
)

logger = logging.getLogger(__name__)


class RagProvidersRepository:
    def __init__(self, db_session: Session):
        self.db_session = db_session

    def create_rag_provider_configuration(
        self,
        rag_provider_config: RagProviderConfiguration,
    ) -> None:
        """Create a new RAG provider configuration with polymorphic support"""
        db_config = rag_provider_config._to_database_model()
        self.db_session.add(db_config)
        self.db_session.commit()

    def _get_db_rag_provider_config(
        self,
        config_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatabaseRagProviderConfiguration | DatabaseApiKeyRagProviderConfiguration:
        q = self.db_session.query(DatabaseRagProviderConfiguration).filter(
            DatabaseRagProviderConfiguration.id == config_id,
        )
        if org_scope is not None:
            q = q.join(
                DatabaseTask,
                DatabaseTask.id == DatabaseRagProviderConfiguration.task_id,
            ).filter(DatabaseTask.org_id == org_scope)
        db_config = q.first()

        if not db_config:
            raise HTTPException(
                status_code=404,
                detail="RAG provider configuration not found",
            )
        return db_config

    def get_rag_provider_configuration(
        self,
        config_id: UUID,
        org_scope: UUID | None = None,
    ) -> RagProviderConfiguration:
        """Get a RAG provider configuration by ID with polymorphic loading"""
        db_config = self._get_db_rag_provider_config(config_id, org_scope=org_scope)
        return RagProviderConfiguration._from_database_model(db_config)

    def get_rag_provider_configurations_by_task(
        self,
        task_id: str,
        pagination_params: PaginationParameters,
        config_name: Optional[str],
        authentication_method: Optional[RagProviderAuthenticationMethodEnum],
        rag_provider_name: Optional[RagAPIKeyAuthenticationProviderEnum],
    ) -> Tuple[List[RagProviderConfiguration], int]:
        """Get RAG provider configurations for a task with pagination"""
        query = self.db_session.query(DatabaseRagProviderConfiguration).filter(
            DatabaseRagProviderConfiguration.task_id == task_id,
        )

        # apply filters
        if config_name:
            query = query.where(
                DatabaseRagProviderConfiguration.name.ilike(f"%{config_name}%"),
            )
        if authentication_method:
            query = query.where(
                DatabaseRagProviderConfiguration.authentication_method
                == authentication_method,
            )
        if rag_provider_name:
            query = query.where(
                DatabaseApiKeyRagProviderConfiguration.rag_provider
                == rag_provider_name,
            )

        # apply sorting
        if pagination_params.sort == PaginationSortMethod.DESCENDING:
            query = query.order_by(desc(DatabaseRagProviderConfiguration.updated_at))
        elif pagination_params.sort == PaginationSortMethod.ASCENDING:
            query = query.order_by(asc(DatabaseRagProviderConfiguration.updated_at))

        total_count = query.count()

        # Apply pagination
        offset = pagination_params.page * pagination_params.page_size
        db_configs = query.offset(offset).limit(pagination_params.page_size).all()

        configs = [
            RagProviderConfiguration._from_database_model(db_config)
            for db_config in db_configs
        ]
        return configs, total_count

    def update_rag_provider_configuration(
        self,
        config_id: UUID,
        update_config: RagProviderConfigurationUpdateRequest,
        org_scope: UUID | None = None,
    ) -> None:
        """Update a RAG provider configuration"""
        db_provider = self._get_db_rag_provider_config(config_id, org_scope=org_scope)

        if update_config.name:
            db_provider.name = update_config.name
        if update_config.description is not None:
            db_provider.description = update_config.description
        if update_config.authentication_config is not None:
            if isinstance(
                update_config.authentication_config,
                ApiKeyRagAuthenticationConfigUpdateRequest,
            ):
                if not isinstance(db_provider, DatabaseApiKeyRagProviderConfiguration):
                    raise HTTPException(
                        status_code=400,
                        detail="RAG provider configuration update not supported for this config type.",
                    )
                if update_config.authentication_config.api_key:
                    new_secret_value = ApiKeyRagProviderSecretValue(
                        api_key=update_config.authentication_config.api_key,
                    )
                    db_provider.api_key.value = new_secret_value.model_dump(mode="json")
                    db_provider.api_key.updated_at = datetime.now()
                if update_config.authentication_config.host_url:
                    db_provider.host_url = str(
                        update_config.authentication_config.host_url,
                    )
                if update_config.authentication_config.rag_provider:
                    db_provider.rag_provider = (
                        update_config.authentication_config.rag_provider
                    )
                if (
                    update_config.authentication_config.authentication_method
                    is not None
                ):
                    db_provider.authentication_method = (
                        update_config.authentication_config.authentication_method
                    )
            else:
                raise HTTPException(
                    status_code=404,
                    detail="RAG provider configuration update not supported for this config type.",
                )
        db_provider.updated_at = datetime.now()

        self.db_session.commit()

    def delete_rag_provider_configuration(
        self, config_id: UUID, org_scope: UUID | None = None
    ) -> None:
        """Delete a RAG provider configuration"""
        db_config = self._get_db_rag_provider_config(config_id, org_scope=org_scope)
        self.db_session.delete(db_config)
        self.db_session.commit()

    def create_rag_setting_configuration(
        self,
        rag_setting_config: RagSearchSettingConfiguration,
    ) -> None:
        """Create a new RAG setting configuration"""
        db_config = rag_setting_config._to_database_model()
        self.db_session.add(db_config)
        self.db_session.commit()

    def _get_db_rag_setting_config(
        self,
        setting_config_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatabaseRagSearchSettingConfiguration:
        q = self.db_session.query(DatabaseRagSearchSettingConfiguration).filter(
            DatabaseRagSearchSettingConfiguration.id == setting_config_id,
        )
        if org_scope is not None:
            q = q.join(
                DatabaseTask,
                DatabaseTask.id == DatabaseRagSearchSettingConfiguration.task_id,
            ).filter(DatabaseTask.org_id == org_scope)
        db_config = q.first()

        if not db_config:
            raise HTTPException(
                status_code=404,
                detail="RAG setting configuration not found",
            )
        return db_config

    def get_rag_setting_configuration(
        self,
        config_id: UUID,
        org_scope: UUID | None = None,
    ) -> RagSearchSettingConfiguration:
        """Get a RAG provider configuration by ID with polymorphic loading"""
        db_config = self._get_db_rag_setting_config(config_id, org_scope=org_scope)
        return RagSearchSettingConfiguration._from_database_model(db_config)

    def delete_rag_setting_configuration(
        self, config_id: UUID, org_scope: UUID | None = None
    ) -> None:
        """Delete a RAG setting configuration"""
        db_config = self._get_db_rag_setting_config(config_id, org_scope=org_scope)
        self.db_session.delete(db_config)
        self.db_session.commit()

    def update_rag_provider_setting_configuration(
        self,
        config_id: UUID,
        update_config: RagSearchSettingConfigurationUpdateRequest,
        org_scope: UUID | None = None,
    ) -> None:
        """Update a RAG provider setting configuration"""
        db_setting_config = self._get_db_rag_setting_config(
            config_id, org_scope=org_scope
        )

        if update_config.name:
            db_setting_config.name = update_config.name
        if update_config.description is not None:
            db_setting_config.description = update_config.description
        if update_config.rag_provider_id:
            # check rag provider exists & lives in caller's org - will raise 404 otherwise
            self._get_db_rag_provider_config(
                update_config.rag_provider_id, org_scope=org_scope
            )
            # set new field
            db_setting_config.rag_provider_id = update_config.rag_provider_id

        db_setting_config.updated_at = datetime.now()

        self.db_session.commit()

    def get_rag_search_setting_configurations_by_task(
        self,
        task_id: str,
        pagination_params: PaginationParameters,
        config_name: Optional[str],
        rag_provider_ids: Optional[list[UUID]],
    ) -> Tuple[List[RagSearchSettingConfiguration], int]:
        """Get RAG provider setting configurations for a task with pagination"""
        query = self.db_session.query(DatabaseRagSearchSettingConfiguration).filter(
            DatabaseRagSearchSettingConfiguration.task_id == task_id,
        )

        # apply filters
        if config_name:
            query = query.where(
                DatabaseRagSearchSettingConfiguration.name.ilike(f"%{config_name}%"),
            )
        if rag_provider_ids:
            query = query.where(
                DatabaseRagSearchSettingConfiguration.rag_provider_id.in_(
                    rag_provider_ids,
                ),
            )

        # apply sorting
        if pagination_params.sort == PaginationSortMethod.DESCENDING:
            query = query.order_by(
                desc(DatabaseRagSearchSettingConfiguration.updated_at),
            )
        elif pagination_params.sort == PaginationSortMethod.ASCENDING:
            query = query.order_by(
                asc(DatabaseRagSearchSettingConfiguration.updated_at),
            )

        total_count = query.count()

        # Apply pagination
        offset = pagination_params.page * pagination_params.page_size
        db_configs = query.offset(offset).limit(pagination_params.page_size).all()

        configs = [
            RagSearchSettingConfiguration._from_database_model(db_config)
            for db_config in db_configs
        ]
        return configs, total_count

    def create_rag_setting_configuration_version(
        self,
        rag_setting_version: RagSearchSettingConfigurationVersion,
        org_scope: UUID | None = None,
    ) -> None:
        """Create a new RAG setting configuration version. Updates parent model metadata as needed."""
        # create new version
        db_version = rag_setting_version._to_database_model()

        # update parent model — also validates ownership for tenant callers
        db_parent_config = self._get_db_rag_setting_config(
            rag_setting_version.setting_configuration_id,
            org_scope=org_scope,
        )
        db_parent_config.updated_at = db_version.created_at
        db_parent_config.latest_version_number = db_version.version_number

        # validate tags must be unique across versions - want to raise clean error instead of 5XX
        existing_tags = db_parent_config.all_possible_tags or []
        existing_tag_strings = {tag.tag for tag in existing_tags}
        for new_tag_db in db_version.tags or []:
            if new_tag_db.tag in existing_tag_strings:
                raise HTTPException(
                    status_code=404,
                    detail=f"RAG setting version tag {new_tag_db.tag} already exists for this configuration.",
                )
        # update all_possible_tags to include the new tags as well
        # reuse the tag DB objects from db_version to avoid creating duplicates
        tags_to_add = [tag_db for tag_db in db_version.tags or []]
        db_parent_config.all_possible_tags = existing_tags + tags_to_add

        # add objects to DB
        self.db_session.add(db_version)
        self.db_session.commit()

    def _get_db_rag_setting_config_version(
        self,
        setting_config_id: UUID,
        version_number: int,
        include_deleted_versions: bool = False,
        org_scope: UUID | None = None,
    ) -> DatabaseRagSearchSettingConfigurationVersion:
        # Validate ownership via the parent config; raises 404 if cross-org.
        if org_scope is not None:
            self._get_db_rag_setting_config(setting_config_id, org_scope=org_scope)
        query = (
            self.db_session.query(DatabaseRagSearchSettingConfigurationVersion)
            .filter(
                DatabaseRagSearchSettingConfigurationVersion.setting_configuration_id
                == setting_config_id,
            )
            .filter(
                DatabaseRagSearchSettingConfigurationVersion.version_number
                == version_number,
            )
        )
        if not include_deleted_versions:
            query = query.filter(
                DatabaseRagSearchSettingConfigurationVersion.deleted_at.is_(None),
            )

        db_config = query.first()

        if not db_config:
            raise HTTPException(
                status_code=404,
                detail="RAG setting configuration version not found",
            )
        return db_config

    def get_rag_setting_configuration_version(
        self,
        config_id: UUID | None,
        version_number: int | None,
        include_deleted_versions: bool = False,
        org_scope: UUID | None = None,
    ) -> RagSearchSettingConfigurationVersion:
        """Get a RAG provider configuration version by ID and version number"""
        if config_id is None or version_number is None:
            raise HTTPException(
                status_code=404,
                detail="Config ID or version number is required",
            )
        db_config = self._get_db_rag_setting_config_version(
            config_id,
            version_number,
            include_deleted_versions=include_deleted_versions,
            org_scope=org_scope,
        )
        return RagSearchSettingConfigurationVersion._from_database_model(db_config)

    def _get_db_rag_setting_configuration_versions(
        self,
        setting_config_id: UUID,
        pagination_params: PaginationParameters,
        tags: Optional[list[str]] = None,
        version_numbers: Optional[list[int]] = None,
        org_scope: UUID | None = None,
    ) -> Tuple[list[DatabaseRagSearchSettingConfigurationVersion], int]:
        """Gets list of DB RAG search configuration versions by ID"""
        # Validate ownership via the parent config; raises 404 if cross-org.
        if org_scope is not None:
            self._get_db_rag_setting_config(setting_config_id, org_scope=org_scope)
        # filters out soft-deleted versions
        base_query = self.db_session.query(
            DatabaseRagSearchSettingConfigurationVersion,
        ).filter(
            DatabaseRagSearchSettingConfigurationVersion.setting_configuration_id
            == setting_config_id,
            DatabaseRagSearchSettingConfigurationVersion.deleted_at.is_(None),
        )

        # join with tags if we're filtering by tags - otherwise relationship will handle returning them
        if tags:
            # case-insensitive in_ search - force list of tags and tags in the DB to be lowercase
            lowercase_tag_search = [tag.lower() for tag in tags]
            base_query = base_query.join(
                DatabaseRagSearchSettingConfigurationVersion.tags,
            ).filter(
                func.lower(DatabaseRagSearchVersionTag.tag).in_(lowercase_tag_search),
            )

        if version_numbers:
            base_query = base_query.filter(
                DatabaseRagSearchSettingConfigurationVersion.version_number.in_(
                    version_numbers,
                ),
            )

        # get count of versions before applying sorting/pagination
        total_count = base_query.count()

        # apply sorting
        if pagination_params.sort == PaginationSortMethod.DESCENDING:
            base_query = base_query.order_by(
                desc(DatabaseRagSearchSettingConfigurationVersion.version_number),
            )
        elif pagination_params.sort == PaginationSortMethod.ASCENDING:
            base_query = base_query.order_by(
                asc(DatabaseRagSearchSettingConfigurationVersion.version_number),
            )

        # Apply pagination
        offset = pagination_params.page * pagination_params.page_size
        db_configs = base_query.offset(offset).limit(pagination_params.page_size).all()

        return db_configs, total_count

    def get_rag_setting_configuration_versions(
        self,
        setting_config_id: UUID,
        pagination_params: PaginationParameters,
        tags: Optional[list[str]] = None,
        version_numbers: Optional[list[int]] = None,
        org_scope: UUID | None = None,
    ) -> Tuple[list[RagSearchSettingConfigurationVersion], int]:
        db_configs, count = self._get_db_rag_setting_configuration_versions(
            setting_config_id,
            pagination_params,
            tags=tags,
            version_numbers=version_numbers,
            org_scope=org_scope,
        )
        return [
            RagSearchSettingConfigurationVersion._from_database_model(db_config)
            for db_config in db_configs
        ], count

    def soft_delete_rag_setting_configuration_version(
        self,
        config_id: UUID,
        version_number: int,
        org_scope: UUID | None = None,
    ) -> None:
        db_version_config = self._get_db_rag_setting_config_version(
            config_id,
            version_number,
            org_scope=org_scope,
        )
        db_parent_config = self._get_db_rag_setting_config(
            config_id, org_scope=org_scope
        )
        db_version_config.deleted_at = datetime.now()

        # empty out all other fields in the version except for the PK fields and the created/updated fields
        db_version_config.settings = None
        # delete the tag records from the database and clear the relationship
        for tag in db_version_config.tags or []:
            self.db_session.delete(tag)
        db_version_config.tags = []

        # update db_parent_config updated_at time since one of its versions was affected
        db_parent_config.updated_at = datetime.now()

        self.db_session.commit()

    def _get_db_rag_setting_config_version_by_tag(
        self,
        setting_config_id: UUID,
        tag: str,
        org_scope: UUID | None = None,
    ) -> DatabaseRagSearchSettingConfigurationVersion:
        # Validate ownership via the parent config; raises 404 if cross-org.
        if org_scope is not None:
            self._get_db_rag_setting_config(setting_config_id, org_scope=org_scope)
        db_config = (
            self.db_session.query(DatabaseRagSearchSettingConfigurationVersion)
            .join(DatabaseRagSearchSettingConfigurationVersion.tags)
            .filter(
                DatabaseRagSearchSettingConfigurationVersion.setting_configuration_id
                == setting_config_id,
            )
            # do case-insensitive tag search
            .filter(DatabaseRagSearchVersionTag.tag.ilike(tag))
            .first()
        )
        # we don't need to handle soft-deletes here because if a version was soft-deleted, the tags will be
        # cleared and the version won't be returned

        if not db_config:
            raise HTTPException(
                status_code=404,
                detail=f"RAG setting configuration version with tag {tag} not found",
            )
        return db_config

    def get_rag_setting_configuration_version_by_tag(
        self,
        config_id: UUID,
        tag: str,
        org_scope: UUID | None = None,
    ) -> RagSearchSettingConfigurationVersion:
        """Get a RAG provider configuration version by ID and version number"""
        db_config = self._get_db_rag_setting_config_version_by_tag(
            config_id, tag, org_scope=org_scope
        )
        return RagSearchSettingConfigurationVersion._from_database_model(db_config)

    def update_rag_provider_setting_configuration_version(
        self,
        config_id: UUID,
        version_number: int,
        update_config: RagSearchSettingConfigurationVersionUpdateRequest,
        org_scope: UUID | None = None,
    ) -> None:
        """Update a RAG provider setting configuration version metadata"""
        db_setting_version = self._get_db_rag_setting_config_version(
            config_id,
            version_number,
            org_scope=org_scope,
        )
        curr_time = datetime.now()

        # delete old tag objects
        for tag_db in db_setting_version.tags or []:
            self.db_session.delete(tag_db)

        # create new tag objects from the string list
        new_tags = [
            RagSearchSettingTag._from_request_model(config_id, version_number, tag)
            for tag in update_config.tags
        ]
        new_db_tags = [tag._to_database_model() for tag in new_tags]

        # adding the tags to the version object will also create them in the DB on commit
        db_setting_version.tags = new_db_tags
        db_setting_version.updated_at = curr_time

        # parent config all_possible_tags will be updated via the join condition
        # just update the updated_at time to indicate something about its versions changed
        db_parent_config = self._get_db_rag_setting_config(config_id)
        db_parent_config.updated_at = curr_time

        self.db_session.commit()
