import json
from io import BytesIO, StringIO
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol

import jmespath
import pandas as pd
from airflow.hooks.base import BaseHook
from airflow.models import BaseOperator

try:
    from airflow.providers.http.operators.http import HttpOperator
except ImportError:
    from airflow.providers.http.operators.http import SimpleHttpOperator as HttpOperator

from airflow.utils.context import Context
from airflow.utils.helpers import merge_dicts
from requests import Response

from airflow_tools.compression_utils import CompressionOptions, compress
from airflow_tools.exceptions import ApiResponseTypeError
from airflow_tools.filesystems.filesystem_factory import FilesystemFactory

if TYPE_CHECKING:
    from requests.auth import AuthBase

SaveFormat = Literal['jsonl']


class Transformation(Protocol):
    def __call__(self, data: bytes, **kwargs) -> bytes:
        ...


class HttpBatchOperator(HttpOperator):
    def execute(
        self, context: Context, use_new_data_parameters_on_pagination=False
    ) -> Any:
        self.log.info("Calling HTTP method")

        response = self.hook.run(
            self.endpoint, self.data, self.headers, self.extra_options
        )
        yield self.process_response(context=context, response=response)
        for response in self.paginate_sync(
            response=response,
            use_new_data_parameters_on_pagination=use_new_data_parameters_on_pagination,
        ):
            yield self.process_response(context=context, response=response)

    def paginate_sync(
        self, response: Response, use_new_data_parameters_on_pagination=False
    ) -> Response | list[Response]:
        if not self.pagination_function:
            return None

        while True:
            next_page_params = self.pagination_function(response)
            if not next_page_params:
                break
            response = self.hook.run(
                **self._merge_next_page_parameters(
                    next_page_params, use_new_data_parameters_on_pagination
                )
            )
            yield response
        return None

    def _merge_next_page_parameters(
        self, next_page_params: dict, use_new_data_parameters_on_pagination=False
    ) -> dict:
        """Merge initial request parameters with next page parameters.

        Merge initial requests parameters with the ones for the next page, generated by
        the pagination function. Items in the 'next_page_params' overrides those defined
        in the previous request.

        :param next_page_params: A dictionary containing the parameters for the next page.
        :return: A dictionary containing the merged parameters.
        """
        data: str | dict | None = None  # makes mypy happy
        next_page_data_param = next_page_params.get("data")
        if use_new_data_parameters_on_pagination and isinstance(
            next_page_data_param, dict
        ):
            data = next_page_data_param
        elif isinstance(self.data, dict) and isinstance(next_page_data_param, dict):
            data = merge_dicts(self.data, next_page_data_param)
        else:
            data = next_page_data_param or self.data

        self.log.info(
            f"Calling HTTP method with endpoint: {next_page_params.get('endpoint') or self.endpoint}"
        )
        self.log.info(f"Calling HTTP method with data: {data}")

        return dict(
            endpoint=next_page_params.get("endpoint") or self.endpoint,
            data=data,
            headers=merge_dicts(self.headers, next_page_params.get("headers", {})),
            extra_options=merge_dicts(
                self.extra_options, next_page_params.get("extra_options", {})
            ),
        )


class HttpToFilesystem(BaseOperator):
    template_fields = list(HttpOperator.template_fields) + [
        'filesystem_path',
        'filesystem_conn_id',
        'jmespath_expression',
        'save_format',
        'source_format',
    ]
    template_fields_renderers = HttpOperator.template_fields_renderers

    json_response_source_format = ['json', 'jsonl']
    binary_response_source_format = ['parquet']

    def __init__(
        self,
        http_conn_id: str,
        filesystem_conn_id: str,
        filesystem_path: str,
        save_format: SaveFormat = 'jsonl',
        source_format: SaveFormat = None,
        compression: CompressionOptions = None,
        endpoint: str | None = None,
        method: str = "POST",
        data: Any = None,
        headers: dict[str, str] | None = None,
        auth_type: type['AuthBase'] | None = None,
        jmespath_expression: str | None = None,
        pagination_function: Callable | None = None,
        use_new_data_parameters_on_pagination: bool = False,
        create_file_on_success: str | None = None,
        data_transformation: Optional[Transformation] = None,
        data_transformation_kwargs: dict[str, Any] | None = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.http_conn_id = http_conn_id
        self.filesystem_conn_id = filesystem_conn_id
        self.filesystem_path = filesystem_path
        self.compression = compression
        self.endpoint = endpoint
        self.method = method
        self.data = data
        self.headers = headers
        self.auth_type = auth_type
        self.jmespath_expression = jmespath_expression
        self.pagination_function = pagination_function
        self.use_new_data_parameters_on_pagination = (
            use_new_data_parameters_on_pagination
        )
        self.create_file_on_success = create_file_on_success
        self.data_transformation = data_transformation
        self.data_transformation_kwargs = data_transformation_kwargs

        self.save_format = save_format
        self.source_format = source_format if source_format else save_format
        self.kwargs = kwargs

        if (
            self.save_format in self.binary_response_source_format
            and self.compression is not None
        ):
            raise ValueError(
                f'Compression is not supported for binary response save formats: {self.binary_response_source_format}'
            )

        if self.data_transformation and not callable(self.data_transformation):
            raise ValueError('data_transformation must be a callable')

        if self.data_transformation is None and self.source_format != self.save_format:
            raise ValueError(
                'data_transformation must be provided if source_format is different from save_format'
            )
        if self.data_transformation_kwargs and self.data_transformation is None:
            raise ValueError(
                'data_transformation must be provided if data_transformation_kwargs is provided'
            )

    def execute(self, context: 'Context') -> Any:
        http_batch_operator = HttpBatchOperator(
            task_id='http-operator',
            http_conn_id=self.http_conn_id,
            endpoint=self.endpoint,
            method=self.method,
            data=self.data,
            headers=self.headers,
            auth_type=self.auth_type,
            response_filter=self._response_filter,
            pagination_function=self.pagination_function,
        )
        for i, data in enumerate(
            http_batch_operator.execute(
                context,
                use_new_data_parameters_on_pagination=self.use_new_data_parameters_on_pagination,
            ),
            start=1,
        ):
            filesystem_protocol = FilesystemFactory.get_data_lake_filesystem(
                connection=BaseHook.get_connection(self.filesystem_conn_id),
            )

            file_path = self.filesystem_path.rstrip('/') + '/' + self._file_name(i)

            filesystem_protocol.write(data, file_path)

            if self.create_file_on_success is not None and isinstance(
                self.create_file_on_success, str
            ):
                success_file_path = (
                    self.filesystem_path.rstrip('/') + '/' + self.create_file_on_success
                )
                filesystem_protocol.write(BytesIO(), success_file_path)

    def _file_name(self, n_part) -> str:
        file_name = f'part{n_part:04}.{self.save_format}'

        if self.compression:
            file_name += f'.{self.compression}'
        return file_name

    def _response_filter(self, response) -> BytesIO:
        if (
            self.jmespath_expression
            and self.source_format in self.json_response_source_format
        ):
            self.data = jmespath.search(self.jmespath_expression, response.json())

        elif (
            self.jmespath_expression
            and self.source_format not in self.json_response_source_format
        ):
            raise ApiResponseTypeError(
                'JMESPath expression is only supported for json and jsonl save formats'
            )
        elif self.source_format in self.json_response_source_format:
            self.data = response.json()

        elif self.source_format in self.binary_response_source_format:
            self.data = response.content
        else:
            self.data = response.text

        # Check if we have a custom data transformation
        if self.data_transformation and self.data_transformation_kwargs:
            return self.data_transformation(self.data, self.data_transformation_kwargs)
        elif self.data_transformation:
            return self.data_transformation(self.data)

        # If we don't have a custom data transformation, use the default one based on the source_format

        match self.source_format:
            case 'json':
                return json_to_binary(self.data, self.compression)

            case 'jsonl':
                if not isinstance(self.data, list):
                    raise ApiResponseTypeError(
                        'Expected response can\'t be transformed to jsonl. It is not  list[dict]'
                    )
                return list_to_jsonl(self.data, self.compression)

            case 'xml':
                return xml_to_binary(self.data, self.compression)

            case 'parquet':
                return self.data

            case 'csv':
                return csv_to_binary(self.data, self.compression)

            case _:
                raise NotImplementedError(
                    f'Unknown source_format/save_format: {self.source_format}'
                )


def list_to_jsonl(data: list[dict], compression: 'CompressionOptions') -> BytesIO:
    out = StringIO()
    df = pd.DataFrame(data)
    df.to_json(out, orient='records', lines=True, compression=compression)
    out.seek(0)
    return BytesIO(out.getvalue().encode())


def json_to_binary(data: dict, compression: 'CompressionOptions') -> BytesIO:
    json_string = json.dumps(data).encode()
    compressed_json = compress(compression, json_string)
    out = BytesIO(compressed_json)
    return out


def csv_to_binary(data: str, compression: 'CompressionOptions') -> BytesIO:
    csv_df = pd.read_csv(StringIO(data), sep=',')
    out = BytesIO()
    csv_df.to_csv(out, compression=compression, index=False)
    out.seek(0)
    return out


def xml_to_binary(data: str, compression: 'CompressionOptions') -> BytesIO:
    compressed_xml = compress(compression, data.encode())
    out = BytesIO(compressed_xml)
    return out
