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 and_, asc, cast, desc
from sqlalchemy.orm import Session
from sqlalchemy.types import Text

from db_models import DatabaseDataset
from db_models.dataset_models import (
    DatabaseDatasetVersion,
    DatabaseDatasetVersionRow,
)
from db_models.task_models import DatabaseTask
from schemas.internal_schemas import (
    Dataset,
    DatasetVersion,
    ListDatasetVersions,
)
from schemas.request_schemas import (
    DatasetUpdateRequest,
    NewDatasetVersionRequest,
)

logger = logging.getLogger(__name__)


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

    def create_dataset(self, dataset: Dataset, commit: bool = True) -> None:
        db_dataset = dataset._to_database_model()
        self.db_session.add(db_dataset)
        if commit:
            self.db_session.commit()
        else:
            self.db_session.flush()

    def _get_db_dataset(
        self,
        dataset_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatabaseDataset:
        q = self.db_session.query(DatabaseDataset).filter(
            DatabaseDataset.id == dataset_id,
        )
        # Datasets carry only task_id — join through tasks for the org filter.
        if org_scope is not None:
            q = q.join(DatabaseTask, DatabaseTask.id == DatabaseDataset.task_id).filter(
                DatabaseTask.org_id == org_scope,
            )
        db_dataset = q.first()
        if not db_dataset:
            raise HTTPException(
                status_code=404,
                detail="Dataset %s not found." % dataset_id,
                headers={"full_stacktrace": "false"},
            )
        return db_dataset

    def get_dataset(self, dataset_id: UUID, org_scope: UUID | None = None) -> Dataset:
        db_dataset = self._get_db_dataset(dataset_id, org_scope=org_scope)
        return Dataset._from_database_model(db_dataset)

    def update_dataset(
        self,
        dataset_id: UUID,
        update_dataset_request: DatasetUpdateRequest,
        org_scope: UUID | None = None,
    ) -> None:
        db_dataset = self._get_db_dataset(dataset_id, org_scope=org_scope)
        if update_dataset_request.name:
            db_dataset.name = update_dataset_request.name
        if update_dataset_request.description is not None:
            db_dataset.description = update_dataset_request.description
        if update_dataset_request.metadata is not None:
            db_dataset.dataset_metadata = update_dataset_request.metadata

        db_dataset.updated_at = datetime.now()

        self.db_session.commit()

    def query_datasets(
        self,
        task_id: str,
        pagination_params: PaginationParameters,
        dataset_ids: Optional[list[UUID]] = None,
        dataset_name: Optional[str] = None,
    ) -> Tuple[List[Dataset], int]:
        base_query = self.db_session.query(DatabaseDataset).where(
            DatabaseDataset.task_id == task_id,
        )

        # apply filters
        if dataset_ids is not None:
            base_query = base_query.where(DatabaseDataset.id.in_(dataset_ids))
        if dataset_name:
            base_query = base_query.where(
                DatabaseDataset.name.ilike(f"%{dataset_name}%"),
            )

        # apply sorting - sort by updated_at field
        if pagination_params.sort == PaginationSortMethod.DESCENDING:
            base_query = base_query.order_by(desc(DatabaseDataset.updated_at))
        elif pagination_params.sort == PaginationSortMethod.ASCENDING:
            base_query = base_query.order_by(asc(DatabaseDataset.updated_at))

        # calculate total count before offset to apply pagination
        count = base_query.count()

        base_query = base_query.offset(
            pagination_params.page * pagination_params.page_size,
        )
        db_datasets = base_query.limit(pagination_params.page_size).all()

        return [
            Dataset._from_database_model(db_dataset) for db_dataset in db_datasets
        ], count

    def delete_dataset(self, dataset_id: UUID, org_scope: UUID | None = None) -> None:
        db_dataset = self._get_db_dataset(dataset_id, org_scope=org_scope)
        self.db_session.delete(db_dataset)
        self.db_session.commit()

    def _get_latest_db_dataset_version(
        self,
        dataset_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatabaseDatasetVersion:
        # Validate ownership via the parent dataset; raises 404 if cross-org.
        if org_scope is not None:
            self._get_db_dataset(dataset_id, org_scope=org_scope)
        db_dataset_version = (
            self.db_session.query(DatabaseDatasetVersion)
            .filter(DatabaseDatasetVersion.dataset_id == dataset_id)
            .order_by(DatabaseDatasetVersion.version_number.desc())
            .first()
        )
        if not db_dataset_version:
            raise HTTPException(
                status_code=404,
                detail="Dataset version for dataset %s not found." % dataset_id,
                headers={"full_stacktrace": "false"},
            )
        return db_dataset_version

    def get_latest_dataset_version(
        self,
        dataset_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatasetVersion:
        db_dataset_version = self._get_latest_db_dataset_version(
            dataset_id,
            org_scope=org_scope,
        )
        # return page params representing that all rows have been fetched
        return DatasetVersion._from_database_model(
            db_dataset_version,
            len(db_dataset_version.version_rows),
            PaginationParameters(
                page=0,
                page_size=len(db_dataset_version.version_rows),
            ),
        )

    def get_dataset_version(
        self,
        dataset_id: UUID,
        dataset_version: int,
        pagination_params: PaginationParameters,
        search_query: Optional[str] = None,
        org_scope: UUID | None = None,
    ) -> DatasetVersion:
        # Validate ownership via the parent dataset; raises 404 if cross-org.
        if org_scope is not None:
            self._get_db_dataset(dataset_id, org_scope=org_scope)
        base_query = (
            self.db_session.query(DatabaseDatasetVersion, DatabaseDatasetVersionRow)
            .join(
                DatabaseDatasetVersionRow,
                and_(
                    DatabaseDatasetVersionRow.version_number
                    == DatabaseDatasetVersion.version_number,
                    DatabaseDatasetVersionRow.dataset_id
                    == DatabaseDatasetVersion.dataset_id,
                ),
                isouter=True,
            )
            .filter(DatabaseDatasetVersion.dataset_id == dataset_id)
            .filter(DatabaseDatasetVersion.version_number == dataset_version)
        )

        if search_query and search_query.strip():
            search_term = f"%{search_query.strip()}%"
            base_query = base_query.filter(
                cast(DatabaseDatasetVersionRow.data, Text).ilike(search_term),
            )

        # apply pagination
        total_count = base_query.count()

        if pagination_params.sort == PaginationSortMethod.ASCENDING:
            base_query = base_query.order_by(DatabaseDatasetVersionRow.created_at.asc())
        else:
            base_query = base_query.order_by(
                DatabaseDatasetVersionRow.created_at.desc(),
            )

        offset = pagination_params.page * pagination_params.page_size
        paginated_results = (
            base_query.offset(offset).limit(pagination_params.page_size).all()
        )

        if not paginated_results:
            # If a search query was provided and returned no results, the version
            # still exists — return it with an empty row set instead of 404.
            if search_query and search_query.strip():
                db_dataset_version = (
                    self.db_session.query(DatabaseDatasetVersion)
                    .filter(DatabaseDatasetVersion.dataset_id == dataset_id)
                    .filter(DatabaseDatasetVersion.version_number == dataset_version)
                    .first()
                )
                if db_dataset_version is None:
                    raise HTTPException(
                        status_code=404,
                        detail="Dataset version for dataset %s not found." % dataset_id,
                        headers={"full_stacktrace": "false"},
                    )
                return DatasetVersion._from_database_model(
                    db_dataset_version,
                    total_count,
                    pagination_params,
                    paginated_rows=[],
                )

            raise HTTPException(
                status_code=404,
                detail="Dataset version for dataset %s not found." % dataset_id,
                headers={"full_stacktrace": "false"},
            )

        # every row will have the same version because of how the query was constructed, so we can just select
        # the first element of the first row as the dataset version
        db_dataset_version = paginated_results[0][0]

        # extract paginated rows from the version (in the second element for every row)
        # if the version has no rows, tuple will be [db_version, None]
        paginated_rows = [
            result[1] for result in paginated_results if result[1] is not None
        ]
        return DatasetVersion._from_database_model(
            db_dataset_version,
            total_count,
            pagination_params,
            paginated_rows=paginated_rows,
        )

    def get_dataset_version_row(
        self,
        dataset_id: UUID,
        dataset_version: int,
        row_id: UUID,
        org_scope: UUID | None = None,
    ) -> DatabaseDatasetVersionRow:
        """Get a specific row by ID from a dataset version"""
        if org_scope is not None:
            self._get_db_dataset(dataset_id, org_scope=org_scope)
        db_row = (
            self.db_session.query(DatabaseDatasetVersionRow)
            .filter(
                DatabaseDatasetVersionRow.dataset_id == dataset_id,
                DatabaseDatasetVersionRow.version_number == dataset_version,
                DatabaseDatasetVersionRow.id == row_id,
            )
            .first()
        )

        if not db_row:
            raise HTTPException(
                status_code=404,
                detail=f"Row {row_id} not found in dataset {dataset_id} version {dataset_version}",
                headers={"full_stacktrace": "false"},
            )

        return db_row

    def create_dataset_version(
        self,
        dataset_id: UUID,
        dataset_version: NewDatasetVersionRequest,
        org_scope: UUID | None = None,
        commit: bool = True,
    ) -> None:
        db_dataset = self._get_db_dataset(dataset_id, org_scope=org_scope)
        try:
            latest_version = self._get_latest_db_dataset_version(
                dataset_id,
                org_scope=org_scope,
            )
        except HTTPException:
            # no version exists for this dataset yet
            latest_version = None

        internal_dataset_version = DatasetVersion._from_request_model(
            dataset_id,
            latest_version,
            dataset_version,
        )
        self.db_session.add(internal_dataset_version._to_database_model())
        db_dataset.updated_at = datetime.now()
        db_dataset.latest_version_number = internal_dataset_version.version_number
        if commit:
            self.db_session.commit()
        else:
            self.db_session.flush()

    def get_dataset_versions(
        self,
        dataset_id: UUID,
        latest_version_only: bool,
        pagination_params: PaginationParameters,
        org_scope: UUID | None = None,
    ) -> ListDatasetVersions:
        if org_scope is not None:
            self._get_db_dataset(dataset_id, org_scope=org_scope)
        base_query = self.db_session.query(DatabaseDatasetVersion).filter(
            DatabaseDatasetVersion.dataset_id == dataset_id,
        )

        # get total count before applying additional filters
        total_count = base_query.count()

        # get first dataset in descending order if latest_version_only filter is set
        applied_page_params = (
            PaginationParameters(
                sort=PaginationSortMethod.DESCENDING,
                page_size=1,
                page=0,
            )
            if latest_version_only
            else pagination_params
        )

        total_count = 1 if latest_version_only and total_count != 0 else total_count

        # apply pagination
        if applied_page_params.sort == PaginationSortMethod.ASCENDING:
            base_query = base_query.order_by(
                DatabaseDatasetVersion.version_number.asc(),
            )
        else:
            base_query = base_query.order_by(
                DatabaseDatasetVersion.version_number.desc(),
            )

        offset = applied_page_params.page * applied_page_params.page_size
        dataset_versions = (
            base_query.offset(offset).limit(applied_page_params.page_size).all()
        )

        return ListDatasetVersions._from_database_model(
            dataset_versions,
            total_count,
            pagination_params,
        )
