import datetime
import uuid
from collections.abc import Sequence
from dataclasses import dataclass
from decimal import Decimal

from sqlalchemy import delete, func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import aliased, lazyload, selectinload
from sqlalchemy.orm.attributes import flag_modified

from app.common.data.interfaces.collections import raise_if_data_source_has_references
from app.common.data.interfaces.exceptions import DuplicateDataSourceItemError, flush_and_rollback_on_exceptions
from app.common.data.models import DataSource, DataSourceOrganisationItem
from app.common.data.models_user import User
from app.common.data.types import (
    DataSourceFileMetadata,
    DataSourceSchema,
    DataSourceSchemaColumn,
    DataSourceType,
    NumberTypeEnum,
    QuestionDataOptions,
    QuestionDataType,
    QuestionPresentationOptions,
    TUnvalidatedDataSetRow,
    TUnvalidatedDataSetRows,
)
from app.common.safe_ids import safe_column_id
from app.constants import DATA_SET_EXTERNAL_ID_COLUMN_HEADER, DATA_SET_IDENTIFIER_COLUMN_HEADERS
from app.deliver_grant_funding.session_models import DataSetColumnMapping
from app.extensions import db, s3_service
from app.types import NOT_PROVIDED, TNotProvided


def get_data_source(
    data_source_id: uuid.UUID,
    *,
    with_organisation_items: bool = False,
) -> DataSource:
    stmt = select(DataSource).where(DataSource.id == data_source_id)

    if with_organisation_items:
        stmt = stmt.options(selectinload(DataSource.organisation_items))
        stmt = stmt.options(lazyload(DataSource.items))

    return db.session.execute(stmt).scalar_one()


def _build_schema_from_column_mappings(column_mappings: list[DataSetColumnMapping]) -> DataSourceSchema:
    if not column_mappings:
        raise ValueError("Cannot build a schema from an empty list of column mappings")

    schema_dict = {}
    for index, mapping in enumerate(column_mappings):
        if mapping.data_type == QuestionDataType.NUMBER:
            presentation_options = QuestionPresentationOptions(
                prefix=mapping.prefix if mapping.prefix is not None else "",
                suffix=mapping.suffix if mapping.suffix is not None else "",
            )
            data_options = QuestionDataOptions(
                number_type=mapping.number_type,
                max_decimal_places=mapping.max_decimal_places if mapping.max_decimal_places is not None else None,
            )
        else:
            presentation_options = QuestionPresentationOptions()
            data_options = QuestionDataOptions()

        schema_dict[safe_column_id(mapping.column_name)] = DataSourceSchemaColumn(
            data_type=mapping.data_type,
            presentation_options=presentation_options,
            data_options=data_options,
            original_column_name=mapping.column_name,
            order=index,
        )
    return DataSourceSchema.model_validate(schema_dict)


def _update_schema_column_order(data_source: DataSource, all_headers: list[str]) -> None:
    if not data_source.schema:
        return

    header_order = {header: index for index, header in enumerate(all_headers)}
    for column in data_source.schema.root.values():
        column.order = header_order.get(column.original_column_name)

    flag_modified(data_source, "schema")


def _clean_value(val: str, mapping: DataSetColumnMapping) -> str | int | None:
    if not val:
        return None
    val = val.strip()
    if mapping.data_type != QuestionDataType.NUMBER or not val:
        return val

    if mapping.prefix:
        val = val.removeprefix(mapping.prefix)
    if mapping.suffix:
        val = val.removesuffix(mapping.suffix)
    val = val.replace(",", "").strip()

    if mapping.number_type == NumberTypeEnum.INTEGER:
        return int(val)
    elif mapping.number_type == NumberTypeEnum.DECIMAL:
        return str(Decimal(val))

    return val


def _build_data_blob(
    row: TUnvalidatedDataSetRow,
    mappings: dict[str, DataSetColumnMapping],
    identifier_columns: list[str],
) -> dict[str, str | int | None]:
    return {
        safe_column_id(col): _clean_value(row.get(col, ""), mappings[col])
        for col in row
        if col not in identifier_columns and col in mappings
    }


def _create_organisation_items(
    data_source: DataSource,
    all_rows: TUnvalidatedDataSetRows,
    column_mappings: list[DataSetColumnMapping],
    identifier_columns: list[str],
) -> None:
    mappings = {m.column_name: m for m in column_mappings}
    for row in all_rows:
        item = DataSourceOrganisationItem(
            data_source_id=data_source.id,
            external_id=row.get(DATA_SET_EXTERNAL_ID_COLUMN_HEADER, "").strip(),
            _data=_build_data_blob(row, mappings, identifier_columns),
        )
        data_source.organisation_items.append(item)
        db.session.add(item)


@flush_and_rollback_on_exceptions(coerce_exceptions=[(IntegrityError, DuplicateDataSourceItemError)])
def create_uploaded_data_source(
    *,
    name: str,
    data_source_type: DataSourceType,
    grant_id: uuid.UUID | None,
    collection_id: uuid.UUID | None,
    column_mappings: list["DataSetColumnMapping"],
    all_rows: TUnvalidatedDataSetRows,
    user: User,
    s3_key: str,
    original_filename: str,
    data_source_id: uuid.UUID,
) -> DataSource:
    schema = _build_schema_from_column_mappings(column_mappings)
    file_metadata = DataSourceFileMetadata(s3_key=s3_key, original_filename=original_filename)
    data_source = DataSource(
        id=data_source_id,
        name=name,
        type=data_source_type,
        grant_id=grant_id,
        collection_id=collection_id,
        schema=schema,
        created_by=user,
        file_metadata=file_metadata,
    )
    db.session.add(data_source)

    match data_source_type:
        case DataSourceType.GRANT_RECIPIENT:
            _create_organisation_items(data_source, all_rows, column_mappings, DATA_SET_IDENTIFIER_COLUMN_HEADERS)
        case _:
            raise ValueError(f"Unsupported data source type: {data_source_type}")

    return data_source


@flush_and_rollback_on_exceptions
def replace_uploaded_data_source(
    data_source: DataSource,
    new_columns: list["DataSetColumnMapping"],
    all_headers: list[str],
    all_rows: TUnvalidatedDataSetRows,
    s3_key: str,
    original_filename: str,
    user: User,
    name: str | TNotProvided = NOT_PROVIDED,
) -> DataSource:

    data_source.updated_by = user

    if name is not NOT_PROVIDED:
        data_source.name = name

    data_source.file_metadata = DataSourceFileMetadata(s3_key=s3_key, original_filename=original_filename)
    if new_columns:
        data_source.schema.root.update(_build_schema_from_column_mappings(new_columns).root)  # ty:ignore[unresolved-attribute]
        flag_modified(data_source, "schema")

    for removed_column_id in [
        k
        for k, v in data_source.schema.root.items()  # ty:ignore[unresolved-attribute]
        if v.original_column_name not in all_headers
    ]:
        data_source.schema.root.pop(removed_column_id)  # ty:ignore[unresolved-attribute]
        flag_modified(data_source, "schema")

    _update_schema_column_order(data_source, all_headers)

    all_column_mappings = [
        DataSetColumnMapping.build_from_data_source_schema_column(schema_column)
        for schema_column in data_source.schema.ordered_values()  # ty:ignore[unresolved-attribute]
    ]

    # delete all existing data source org items as the IDs aren't referenced, just column names
    stmt = delete(DataSourceOrganisationItem).where(DataSourceOrganisationItem.data_source_id == data_source.id)
    db.session.execute(stmt)
    data_source.organisation_items.clear()

    # create all new org items with the updated data
    _create_organisation_items(data_source, all_rows, all_column_mappings, DATA_SET_IDENTIFIER_COLUMN_HEADERS)

    return data_source


@flush_and_rollback_on_exceptions
def delete_data_source(data_source: DataSource) -> None:
    raise_if_data_source_has_references(data_source)

    if data_source.file_metadata and data_source.file_metadata.s3_key:
        s3_service.delete_file(data_source.file_metadata.s3_key)

    db.session.delete(data_source)


@dataclass(frozen=True)
class ListDataSourceData:
    id: uuid.UUID
    name: str
    updated_at_utc: datetime.datetime
    uploaded_by_name: str | None
    has_missing_data: bool


def get_data_source_list_for_collection(collection_id: uuid.UUID) -> Sequence[ListDataSourceData]:
    user_updated = aliased(User)
    user_created = aliased(User)

    stmt = (
        select(
            DataSource.id,
            DataSource.name,
            DataSource.updated_at_utc,
            func.coalesce(user_updated.name, user_created.name).label("uploaded_by_name"),
            DataSource.has_missing_data().label("has_missing_data"),  # ty:ignore[unresolved-attribute]
        )
        .outerjoin(user_created, DataSource.created_by_id == user_created.id)
        .outerjoin(user_updated, DataSource.updated_by_id == user_updated.id)
        .where(
            DataSource.collection_id == collection_id,
            DataSource.type == DataSourceType.GRANT_RECIPIENT,
        )
        .order_by(DataSource.name)
    )
    return [
        ListDataSourceData(
            id=row.id,
            name=row.name,
            updated_at_utc=row.updated_at_utc,
            uploaded_by_name=row.uploaded_by_name,
            has_missing_data=row.has_missing_data,
        )
        for row in db.session.execute(stmt).all()
    ]


def get_collection_ids_with_missing_data_data_sets(grant_id: uuid.UUID) -> set[uuid.UUID]:
    stmt = (
        select(DataSource.collection_id)
        .where(
            DataSource.grant_id == grant_id,
            DataSource.type == DataSourceType.GRANT_RECIPIENT,
            DataSource.collection_id.is_not(None),
            DataSource.has_missing_data(),
        )
        .distinct()
    )
    return {id for id in db.session.scalars(stmt).all() if id is not None}


def get_referenced_grant_recipient_data_sources_for_collection(
    collection_id: uuid.UUID,
    *,
    only_with_missing_data: bool = False,
    with_organisation_items: bool = False,
) -> Sequence[DataSource]:
    """
    Get data sets belonging to a collection that are referenced by a component in its forms.

    Pass `only_with_missing_data=True` to further filter down to those with missing data.

    Pass `with_organisation_items=True` to eager-load `organisation_items` - needed by callers that go on to inspect
    exactly which organisations have missing data.
    """
    stmt = (
        select(DataSource)
        .where(
            DataSource.collection_id == collection_id,
            DataSource.type == DataSourceType.GRANT_RECIPIENT,
            DataSource.depended_on_by_columns.any(),
        )
        .order_by(DataSource.name)
    )
    if only_with_missing_data:
        stmt = stmt.where(DataSource.has_missing_data())
    if with_organisation_items:
        stmt = stmt.options(selectinload(DataSource.organisation_items))

    return db.session.scalars(stmt).all()
