import asyncio
import socket
from datetime import datetime
from pathlib import Path
from typing import Any

import msgspec
import structlog
from aiohttp import ClientError, TCPConnector
from aiohttp.http_exceptions import HttpProcessingError
from aiohttp.web_exceptions import HTTPException
from gql import Client, gql
from gql.transport.aiohttp import AIOHTTPTransport
from gql.transport.exceptions import TransportError, TransportProtocolError, TransportQueryError, TransportServerError
from graphql import DocumentNode
from jinja2 import Template

from csirtmu.beta_temporal.config import AlphaConfig
from csirtmu.beta_temporal.external_services.alpha.structs import (
    GetAccountingHistoryResult,
    QueryResult,
    RawQueryResult,
)

HTTP_ERR = (
    ClientError,
    HttpProcessingError,
    TransportQueryError,
    TransportProtocolError,
    TransportServerError,
    HTTPException,
    TransportError,
)


class AlphaClient:
    def __init__(self, alpha_config: AlphaConfig) -> None:
        self.alpha_config = alpha_config
        self.logger = structlog.get_logger()
        with (Path(__file__).parent / "queries.graphql").open() as f:
            self._queries = gql(f.read())
        with (Path(__file__).parent / "ip_query.graphql.j2").open() as f:
            self._ip_query = f.read()

    async def get_ip_id(self, ip: str) -> int | None:
        operation = "getIPAddressID"
        result = await self._send_graphql_req(operation, {"ip": ip})

        if not result:
            return None

        if ips := result["getIpAddressList"]["data"]:
            return ips[0]["id"]

        return None

    async def get_user_for_ip_and_timestamp(
        self, ip: str, time_from: datetime, time_to: datetime
    ) -> GetAccountingHistoryResult | None:
        operation = "getUserForIPAndTimestamp"

        res = await self._send_graphql_req(
            operation, {"ip": ip, "timeFrom": time_from.isoformat(), "timeTo": time_to.isoformat()}
        )
        if res is None:
            return None
        data = msgspec.convert(res.get("getRadiusUserAccountingRecordList", {}), type=GetAccountingHistoryResult)
        data.data += msgspec.convert(
            res.get("getEduVpnUserAccountingRecordList", {}), type=GetAccountingHistoryResult
        ).data
        return data

    async def get_ip_address_data(self, ips: list[str]) -> dict[str, QueryResult] | None:
        query = Template(self._ip_query).render(ips=enumerate(ips))
        operation = "getIPData"
        res = await self._send_graphql_req(operation, {}, query=gql(query))
        if res is None:
            return None
        raw_data: dict[str, RawQueryResult] = msgspec.convert(
            {
                ip: {"ip_addresses": res.get(f"ip_{i}", {"data": []}), "prefixes": res.get(f"prefix_{i}", {"data": []})}
                for i, ip in enumerate(ips)
            },
            type=dict[str, RawQueryResult],
        )
        return {k: v.to_query_result() for k, v in raw_data.items()}

    async def _send_graphql_req(
        self, operation: str, variables: dict[str, Any], query: DocumentNode | None = None
    ) -> dict[str, Any] | None:
        if query is None:
            query = self._queries
        client_session_args: dict[str, Any] = {}
        if self.alpha_config.force_ipv4_only:
            client_session_args["connector"] = TCPConnector(loop=asyncio.get_event_loop(), family=socket.AF_INET)
        transport = AIOHTTPTransport(
            url=self.alpha_config.url,
            headers=self.alpha_config.auth_headers,
            client_session_args=client_session_args,
        )
        try:
            async with Client(
                transport=transport,
                fetch_schema_from_transport=False,
                execute_timeout=self.alpha_config.timeout,
            ) as client:
                try:
                    return await client.execute(  # type: ignore
                        query,
                        operation_name=operation,
                        variable_values=variables,
                        parse_result=False,
                    )
                except TimeoutError:
                    self.logger.exception("Alpha connection timeout")
                    return None

        except HTTP_ERR as e:
            self.logger.exception("Alpha connection error")
            self.logger.exception(e)
            return None
