import logging
import re
from typing import Dict, List, Tuple, Callable, Coroutine, AnyStr, Any

import python_multipart
from aiohttp import ClientSession, ClientResponse
from fastapi import HTTPException, UploadFile
from python_multipart.multipart import Field, File

from amqp.adapter.amq_response_factory import AMQResponseFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQMessage, AMQResponse

import rest.v0.endpoints.benchmark as benchmark
import rest.v0.endpoints.jobs as jobs
import rest.v0.endpoints.unstructured as unstructured
from core.deps import get_current_user
from schemas.user_schema import UserSchema

import json
from urllib.parse import parse_qs
from xml.etree import ElementTree
import io


class CleverMultiPartParser:
    def __init__(self, message: AMQMessage):
        self.form_data = {}
        # Convert each list of strings to a single string joined by newline, then to bytes
        headers = {key: "\n".join(value).encode("utf-8") for key, value in message.headers.items()}
        python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)

    def on_field(self, field: Field):
        print(field)
        self.form_data[field.field_name.decode('utf-8')] = field.value
        print(self.form_data)

    def on_file(self, file: File):
        print(file)
        self.form_data[file.field_name.decode('utf-8')] = \
            UploadFile(file.file_object, size=file.size, filename=file.file_name)
        print(self.form_data)


def parse_request_body(message: AMQMessage):
    """
    Parses the HTTP POST request body based on the MIME type.

    Args:
        body (str): The raw request body as a string.
        content_type (str): The Content-Type header of the request.

    Returns:
        dict: Parsed data as a dictionary.
    """
    content_type = message.headers.get('Content-Type', '')
    if not content_type:
        raise ValueError("Content-Type header is required")

    # Parse JSON
    if content_type == "application/json":
        try:
            return json.loads(message.body())
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON: {e}")

    # Parse URL-encoded form data
    elif content_type == "application/x-www-form-urlencoded":
        try:
            return {k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()}
        except Exception as e:
            raise ValueError(f"Invalid URL-encoded form data: {e}")

    # Parse multipart/form-data
    elif content_type.startswith("multipart/form-data"):
        try:
            parser: CleverMultiPartParser = CleverMultiPartParser(message)
            return parser.form_data
        except Exception as e:
            raise ValueError(f"Invalid multipart/form-data: {e}")

    # Parse plain text
    elif content_type == "text/plain":
        return {"text": message.body()}

    # Parse XML
    elif content_type == "application/xml":
        try:
            root = ElementTree.fromstring(message.body())
            return {elem.tag: elem.text for elem in root}
        except ElementTree.ParseError as e:
            raise ValueError(f"Invalid XML: {e}")

    else:
        raise ValueError(f"Unsupported Content-Type: {content_type}")


BENCHMARK = '/benchmark'
JOBS = '/jobs/'
UNSTRUCTURED = '/unstructured'

GET_BENCHMARK_KG = r".*/benchmark/(?P<request_id>[^\?]+)\?index=(?P<index>\d+)"
GET_JOB_STATUS = r".*/jobs/(?P<request_id>.*)"
GET_KG = r".*/unstructured/(?P<request_id>.*)"
DUMMY = r"(?P<request_id>.*)"


# =============  POST Helper function =============

async def call_cleverswarm_post_put_api(
        message: AMQMessage,
        args: Any,
        api_method: Callable[[Any, UserSchema], Coroutine],
        success_code: int,
        authenticated_user: UserSchema
) -> AMQResponse:
    try:
        coroutine_call: Coroutine = api_method(args, authenticated_user)
        result: str = await coroutine_call
        return AMQResponseFactory.of(
            message.id, success_code, message.headers.get("Content-Type", "application/json"),
            result.encode('utf-8'), message.trace_info
        )
    except HTTPException as http_error:
        return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
                                     str(http_error.detail).encode('utf-8'), message.trace_info)
    except Exception as error:
        return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)


#
# Following methods are wrappers to actual CleverSwarm endpoints, so I can refer to the CleverSwarm method and
# pass it as input parameter using general Callable[[Any, UserSchema], str] type
#
async def cleverswarm_create_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await benchmark.create_benchmark_job(**args, authenticated_user=authenticated_user))


async def cleverswarm_create_extract_with_ontology(args: Any, authenticated_user: UserSchema) -> str:
    return str(await unstructured.create_extract_with_ontology(**args, authenticated_user=authenticated_user))


async def cleverswarm_create_extract_with_wildcards(args: Any, authenticated_user: UserSchema) -> str:
    return str(await unstructured.create_extract_with_ontology(**args, authenticated_user=authenticated_user))


async def cleverswarm_update_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await benchmark.update_benchmark_job(**args, authenticated_user=authenticated_user))


async def cleverswarm_retry_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await jobs.retry_job(args[0], authenticated_user=authenticated_user))


async def cleverswarm_delete_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await jobs.delete_job(args[0], authenticated_user=authenticated_user))


# =============  GET Helper function =============
async def call_cleverswarm_get_api(
        message: AMQMessage,
        pattern: str,
        api_method: Callable[[Tuple[AnyStr | Any, ...], UserSchema], Coroutine],
        authenticated_user: UserSchema
) -> AMQResponse:
    try:
        match = re.search(pattern, message.path)
        if match:
            crt: Coroutine = api_method(match.groups(), authenticated_user)
            result: str = await crt
            return AMQResponseFactory.of(message.id, 200, "text/plain", result.encode('utf-8'), message.trace_info)
        else:
            _msg = f"Unexpected URL: {message.path}, expected pattern: {pattern}"
            return AMQResponseFactory.of(message.id, 404, "text/plain", _msg.encode('utf-8'), message.trace_info)
    except HTTPException as http_error:
        return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
                                     str(http_error.detail).encode('utf-8'), message.trace_info)
    except Exception as error:
        return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)


async def cleverswarm_get_benchmark_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await benchmark.get_benchmark_kg(args[0], args[1], authenticated_user=authenticated_user))


async def cleverswarm_get_job_status(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await jobs.get_job_status(args[0], authenticated_user=authenticated_user))


async def cleverswarm_jobs_get_requests(args: Tuple[AnyStr | Any, ...],
                                        authenticated_user: UserSchema) -> str:
    return str(await jobs.get_requests(authenticated_user=authenticated_user))


async def cleverswarm_get_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await unstructured.get_kg(args[0], authenticated_user=authenticated_user))


class CleverSwarmAdapter:
    def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
        self.amq_configuration = amq_configuration
        self.session = session
        self.service_host = self.amq_configuration.dispatch.service_host
        self.service_port = self.amq_configuration.dispatch.service_port

    async def on_message(self, message) -> AMQResponse:
        DEBUG_NO_AUTH = 1
        _auth: str = message.headers.get('Authorization', '')
        _auth_user: UserSchema | None = None
        if 'Bearer' in _auth:
            try:
                token = _auth.split(' ')[1]
                _auth_user = await get_current_user(token)
            except Exception as error:
                logging.error(f"Error getting user: {error}")
        if _auth_user or DEBUG_NO_AUTH == 1:
            method = message.method.lower()
            if method == 'get':
                return await self.get(message, _auth_user)
            elif method == 'put':
                return await self.put(message, _auth_user)
            elif method == 'post':
                return await self.post(message, _auth_user)
            elif method == 'head':
                return await self.head(message, _auth_user)
            elif method == 'delete':
                return await self.delete(message, _auth_user)
            else:
                raise ValueError(f"Unexpected HTTP method: {message.method}")
        else:
            return AMQResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)

    # ==================================================================================================
    # =====================  C l e v e r S w a r m   A P I   m e t h o d s  ============================
    # ==================================================================================================

    async def get(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the GET requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the GET request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        if BENCHMARK in message.path:
            return await call_cleverswarm_get_api(
                message,
                GET_BENCHMARK_KG,
                cleverswarm_get_benchmark_kg,
                authenticated_user=auth_user
            )

        if JOBS in message.path:
            return await call_cleverswarm_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_get_job_status,
                authenticated_user=auth_user
            )

        if message.path.endswith('/jobs'):
            return await call_cleverswarm_get_api(
                message,
                DUMMY,
                cleverswarm_jobs_get_requests,
                authenticated_user=auth_user
            )

        if UNSTRUCTURED in message.path:
            return await call_cleverswarm_get_api(
                message,
                GET_KG,
                cleverswarm_get_kg,
                authenticated_user=auth_user
            )

        # unmatched path - try to invoke the service directly on this path
        url = f"http://{self.service_host}:{self.service_port}{message.path}"
        headers = {**message.headers, **message.trace_info}
        logging.info(f" [*] SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}")
        if self.session:
            _http_resp: ClientResponse = await self.session.get(url, headers=headers)
            return AMQResponseFactory.of(message.id, _http_resp.status, _http_resp.content_type,
                                         await _http_resp.read(),
                                         message.trace_info)
        return AMQResponseFactory.of(message.id, 404, "text/plain",
                                     str("Destination location does not exists").encode('utf-8'), message.trace_info)

    async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the PUT requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the PUT request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        parser: CleverMultiPartParser = CleverMultiPartParser(message)
        if BENCHMARK in message.path:
            return await call_cleverswarm_post_put_api(
                message,
                parser.form_data,
                cleverswarm_create_benchmark_job,
                201,
                authenticated_user=auth_user
            )

        if JOBS in message.path:
            return await call_cleverswarm_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_retry_job,
                authenticated_user=auth_user
            )

        return await self.handle_possible_form(
            self.session.put(
                f"http://{self.service_host}:{self.service_port}{message.path}",
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        )

    async def post(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the POST requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the POST request, including body and headers
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        parser: CleverMultiPartParser = CleverMultiPartParser(message)
        if BENCHMARK in message.path:
            return await call_cleverswarm_post_put_api(
                message,
                parser.form_data,
                cleverswarm_create_benchmark_job,
                201,
                authenticated_user=auth_user
            )
        if UNSTRUCTURED in message.path:
            if 'with_ontology' in message.path:
                return await call_cleverswarm_post_put_api(
                    message,
                    parser.form_data,
                    cleverswarm_create_extract_with_ontology,
                    200,
                    authenticated_user=auth_user
                )
            if 'with_wildcards' in message.path:
                return await call_cleverswarm_post_put_api(
                    message,
                    parser.form_data,
                    cleverswarm_create_extract_with_wildcards,
                    200,
                    authenticated_user=auth_user
                )

        return await self.handle_possible_form(
            self.session.post(
                f"http://{self.service_host}:{self.service_port}{message.path}",
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        )

    async def handle_possible_form(self, request_coroutine, auth_user: UserSchema) -> AMQResponse:
        return await request_coroutine

    async def head(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        raise NotImplementedError

    async def delete(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        parser: CleverMultiPartParser = CleverMultiPartParser(message)
        if JOBS in message.path:
            return await call_cleverswarm_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_retry_job,
                authenticated_user=auth_user
            )

        raise NotImplementedError

    def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
        for header, values in message_headers.items():
            if header.lower() == 'content-type':
                return values[0].split(';')[0]
        return 'application/json'