from kafka.sasl.oauth import AbstractTokenProvider
from aws_msk_iam_sasl_signer import MSKAuthTokenProvider

from app.core.config import settings
from app.core.logger import logger

class MSKTokenProvider(AbstractTokenProvider):
    def token(self):
        oauth2_token, expiry_millis = MSKAuthTokenProvider.generate_auth_token(
            region=settings.AWS_DEFAULT_REGION,
        )
        logger.info(
            f"MSKTOKENPROVIDER: Token generated successfully. Expires in approx {expiry_millis / 1000 / 60:.1f} minutes."
        )
        return oauth2_token


class KafkaTransactionEventConsumer:
    def __init__(self, loop=None):
        self.consumer = None
        self.admin_client = None
        self.running = False
        self._stop_event = threading.Event()
        self.consumer_thread = None
        self.loop = loop or asyncio.get_event_loop()
        self.token_provider = MSKTokenProvider()

    async def start(self):
        bootstrap_servers = get_bootstrap_servers()
        logger.info(f"Bootstrap servers: {bootstrap_servers}") # retrieving  the brokers
        consumer_config = {
                "bootstrap_servers": bootstrap_servers,
                "security_protocol": "SASL_SSL",
                "sasl_mechanism": "AWS_MSK_IAM",
                "client_id": (settings.CLIENT_ID or "default_client_id")
                + "-main-consumer",
                "group_id": (settings.CLIENT_ID or "default_client_id")
                + "-consumer-group",
                "value_deserializer": lambda m: json.loads(m.decode("utf-8")),
                "auto_offset_reset": "earliest",
                "consumer_timeout_ms": 1000,
                "sasl_oauth_token_provider": self.token_provider,
                "api_version": (2, 6, 0),  
                "api_version_auto_timeout_ms": 10000,
                "request_timeout_ms": 30000,
                "connections_max_idle_ms": 540000,
                "session_timeout_ms": 30000,
                "heartbeat_interval_ms": 10000,
                "max_poll_records": 100,
                "enable_auto_commit": True,
                "auto_commit_interval_ms": 5000,
            }

         logger.info("Creating Kafka consumer with config...")
         self.consumer = KafkaConsumer(
                settings.KAFKA_TOPIC_TRANSACTION_EVENTS,
                **consumer_config,
            )

def get_bootstrap_servers() -> list:
    logger.info(settings.AWS_KAFKA_ARN)
    try:
        config = Config(
            region_name=settings.AWS_DEFAULT_REGION,
        )
        client = boto3.client(
            "kafka",
            region_name=settings.AWS_DEFAULT_REGION,
            aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
            config=config,
        )

        bootstrap_response = client.get_bootstrap_brokers(
            ClusterArn=settings.AWS_KAFKA_ARN
        )

        logger.info(bootstrap_response)

        if "BootstrapBrokerStringSaslIam" in bootstrap_response:
            broker_string = bootstrap_response["BootstrapBrokerStringSaslIam"]
            return broker_string.split(",")
        elif "BootstrapBrokerStringPublicTls" in bootstrap_response:
            broker_string = bootstrap_response["BootstrapBrokerStringPublicTls"]
            return broker_string.split(",")
        elif "BootstrapBrokerString" in bootstrap_response:
            broker_string = bootstrap_response["BootstrapBrokerString"]
            return broker_string.split(",")
        elif "BootstrapBrokerStringTls" in bootstrap_response:
            broker_string = bootstrap_response["BootstrapBrokerStringTls"]
            return broker_string.split(",")

        logger.warning(
            "No bootstrap servers found in MSK response, falling back to configured servers"
        )
    except Exception as e:
        logger.error(f"Error connecting to MSK: {e}")
        logger.warning("Falling back to configured bootstrap servers")

    if settings.KAFKA_BOOTSTRAP_SERVERS:
        return settings.KAFKA_BOOTSTRAP_SERVERS.split(",")
    else:
        logger.error("No bootstrap servers available from MSK or environment")
        return []