import os
import json
import logging
import uuid
import boto3
import jwt  # PyJWT

from typing import Any, Dict, Optional
from urllib.parse import unquote_plus
from datetime import datetime, timedelta
from ..aws.services.AWSSSM import SSM
from ..aws.services.AWSSTS import STS
from ..config.SSMParameters import ssmParameter


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

ENVIRONMENT = os.getenv("NODE_ENVIRONMENT", "dev")
REGION = os.getenv("AWS_DEFAULT_REGION", "us-east-1")

# NOTE: boto3 is synchronous. If you want non-blocking calls inside asyncio,
# consider using aiobotocore. This implementation mirrors your TS behavior with boto3.
_lambda_client = boto3.client("lambda", region_name=REGION)


def _parse_query_string(query: str) -> Dict[str, str]:
    result: Dict[str, str] = {}
    if not query:
        return result
    q = query.lstrip("?")
    if not q:
        return result
    for part in q.split("&"):
        if not part:
            continue
        if "=" in part:
            k, v = part.split("=", 1)
            result[k.strip()] = unquote_plus(v).strip()
        else:
            result[part.strip()] = ""
    return result


class LambdaCommunicationService:
    """
    Python equivalent of your TypeScript LambdaCommunicationService.
    Methods are async to match TS API, but boto3 calls are blocking.
    """

    def __init__(self, environment: Optional[str] = None, region: Optional[str] = None):
        self.environment = environment or ENVIRONMENT
        self.region = region or REGION
        self.ssm = SSM()  # optional injection, expected to have get_ssm_parameter(name)
        self.sts = STS()  # optional injection, expected to have get_identity() or get_caller_identity()
        # If not provided, module-level _lambda_client is used
        self.lambda_client = _lambda_client
        self.ssm_param = ssmParameter


    # -------------------------
    # Validation
    # -------------------------
    @staticmethod
    def _validate_request(req: Dict[str, Any]) -> None:
        if not req:
            logger.error("communication request is required.")
            raise ValueError("communication request is required.")
        for name in ("moduleName", "method", "path", "userType"):
            if not req.get(name):
                logger.error("%s is required.", name)
                raise ValueError(f"{name} is required.")

    # -------------------------
    # Payload generation
    # -------------------------
    async def generate_payload(
        self,
        method: str,
        path: str,
        json_data: Any,
        user_type: str,
        raw_token: Optional[str],
        id_: int,
        parent_id: Any = "0",
        reseller_id: str = "0",
        merchant_id: str = "0",
    ) -> str:
        logger.info("inside generate_payload start")

        if not raw_token or raw_token == "":
            raw_token = await self.generate_jwt_token(user_type, id_, parent_id, reseller_id, merchant_id)

        query_string_parameters: Optional[Dict[str, str]] = None
        multi_value_query_string_parameters: Dict[str, list] = {}

        query_string = path[path.find("?") :] if "?" in path else ""
        if query_string:
            query_string_parameters = _parse_query_string(query_string)
            for k, v in (query_string_parameters or {}).items():
                multi_value_query_string_parameters[k] = [v]

        # get identity from injected STS module or from global if available
        identity = {}
        try:
            sts_client = self.sts if self.sts is not None else globals().get("STS")
            if sts_client:
                if hasattr(sts_client, "get_identity"):
                    identity = sts_client.get_identity()
                elif hasattr(sts_client, "get_caller_identity"):
                    identity = sts_client.get_caller_identity()
                elif hasattr(sts_client, "getCallerIdentity"):
                    identity = sts_client.getCallerIdentity()
                else:
                    identity = {}
            else:
                identity = {}
        except Exception:
            logger.exception("Failed to fetch identity from STS; continuing with empty identity")
            identity = {}

        api_event: Dict[str, Any] = {
            "httpMethod": method,
            "path": path.split("?")[0] if "?" in path else path,
            "pathParameters": {"proxy": path.split("?")[0] if "?" in path else path},
            "body": json.dumps(json_data) if json_data not in (None, "", {}) else "",
            "requestContext": {
                "requestId": str(uuid.uuid4()),
                "identity": {"caller": "", "userAgent": "AuroChat-LambdaCommunicator"},
                "accountId": identity.get("Account") if isinstance(identity, dict) else identity,
            },
            "headers": {
                "via": "",
                "X-Forwarded-For": "",
                "Accept": "application/json",
                "Content-Type": "application/json",
                "User-Agent": "AuroChat-LambdaCommunicator",
                "Authorization": f"Bearer {raw_token}",
                "Accept-Encoding": "gzip, deflate, br",
            },
            "stageVariables": {"env": self.environment},
            "queryStringParameters": query_string_parameters,
        }

        if multi_value_query_string_parameters:
            api_event["multiValueQueryStringParameters"] = multi_value_query_string_parameters

        logger.info("inside generate_payload end")
        return json.dumps(api_event)

    # -------------------------
    # Lambda invocation
    # -------------------------
    async def invoke_api(self, communication_request: Dict[str, Any]) -> Any:
        logger.info("LambdaCommunicationService : invokeAPI : Start")
        self._validate_request(communication_request)

        if not communication_request.get("rawToken"):
            logger.error("raw token is required.")
            raise ValueError("raw token is required.")

        try:
            invoke_payload = await self.generate_payload(
                communication_request["method"],
                communication_request["path"],
                communication_request.get("jsonData"),
                communication_request["userType"],
                communication_request.get("rawToken"),
                communication_request["id"],
                communication_request.get("parentId", "0"),
            )

            function_name = f"{self.environment}-{communication_request['moduleName']}-API"
            params = {
                "FunctionName": function_name,
                "InvocationType": "RequestResponse",
                "Payload": invoke_payload.encode("utf-8"),
            }

            response = self.lambda_client.invoke(**params)
            logger.info("LambdaCommunicationService : invokeAPI : End")

            payload_stream = response.get("Payload")
            if payload_stream:
                # boto3 returns a StreamingBody; read and try parse
                raw = payload_stream.read()
                try:
                    return json.loads(raw) if raw else raw
                except Exception:
                    return raw
            return None
        except Exception:
            logger.exception("Error occured while invoking api")
            raise

    async def invoke_lambda_api(self, communication_request: Dict[str, Any]) -> Any:
        logger.info("LambdaCommunicationService : invokeLambdaAPI : Start")
        self._validate_request(communication_request)

        try:
            invoke_payload = await self.generate_payload(
                communication_request["method"],
                communication_request["path"],
                communication_request.get("jsonData"),
                communication_request["userType"],
                communication_request.get("rawToken"),
                communication_request["id"],
                communication_request.get("parentId", "0"),
            )

            function_name = f"{self.environment}-{communication_request['moduleName']}-API"
            params = {
                "FunctionName": function_name,
                "InvocationType": "RequestResponse",
                "Payload": invoke_payload.encode("utf-8"),
            }

            response = self.lambda_client.invoke(**params)
            logger.info("LambdaCommunicationService : invokeLambdaAPI : End")

            payload_stream = response.get("Payload")
            if payload_stream:
                raw = payload_stream.read()
                try:
                    return json.loads(raw) if raw else raw
                except Exception:
                    return raw
            return None
        except Exception:
            logger.exception("Error occured while invoking api")
            raise

    async def invoke_lambda(self, module_name: str, invocation_type: str, events: Any) -> Any:
        logger.info("LambdaCommunicationService : invokeLambda : Start")
        try:
            params = {
                "FunctionName": module_name,
                "InvocationType": invocation_type or "RequestResponse",
                "Payload": json.dumps(events).encode("utf-8"),
            }
            response = self.lambda_client.invoke(**params)
            logger.info("LambdaCommunicationService : invokeLambda : End")
            payload_stream = response.get("Payload")
            if payload_stream:
                raw = payload_stream.read()
                try:
                    return json.loads(raw) if raw else raw
                except Exception:
                    return raw
            return None
        except Exception:
            logger.exception("Error occured while invoking lambda")
            raise

    # -------------------------
    # Token generation & verification
    # -------------------------
    async def generate_jwt_token(
        self,
        user_type: str,
        id_: int,
        parent_id: str = "0",
        reseller_id: str = "0",
        merchant_id: str = "0",
    ) -> str:
        return await self.generate_token(user_type, id_, parent_id, reseller_id, merchant_id)

    async def generate_token(
        self,
        user_type: str,
        id_: int,
        parent_id: str = "0",
        reseller_id: str = "0",
        merchant_id: str = "0",
    ) -> str:
        logger.info("inside generate_token start")

        token_ssm_key = self.ssm_param.get("tokenSSM")
        if token_ssm_key is None and not self.ssm:
            raise RuntimeError("SSM client or ssmParameter.tokenSSM must be available")

        try:
            ssm_config = self.ssm.get_ssm_parameter(token_ssm_key) if self.ssm else {}
        except Exception:
            logger.exception("Failed to fetch SSM config for tokenSSM")
            ssm_config = {}

        now = datetime.utcnow()
        payload: Dict[str, Any] = {
            "Id": id_,
            "sub": "System",
            "UserType": user_type,
            "jit": "0",
            "Role": "0",
            "ParentId": "0" if parent_id in (None, "0") else str(parent_id),
            "Service": "Service",
        }

        secret_signing_key = ssm_config.get("SigningKey") if isinstance(ssm_config, dict) else None
        issuer = ssm_config.get("Issuer") if isinstance(ssm_config, dict) else None
        audience = ssm_config.get("Audience") if isinstance(ssm_config, dict) else None

        if issuer:
            payload["iss"] = issuer
        if audience:
            payload["aud"] = audience

        payload["exp"] = now + timedelta(hours=1)
        payload["nbf"] = now

        if not secret_signing_key:
            raise RuntimeError("SigningKey for JWT not found in SSM")

        token = jwt.encode(payload, secret_signing_key, algorithm="HS256")
        if isinstance(token, bytes):
            token = token.decode("utf-8")
        logger.info("inside generate_token end")
        return token

    async def generate_refresh_jwt_token(
        self,
        user_type: str,
        id_: int,
        parent_id: str = "0",
        reseller_id: str = "0",
        merchant_id: str = "0",
    ) -> str:
        logger.info("inside generate_refresh_jwt_token start")

        refresh_key_obj = self.ssm_param.get("refreshTokenSSM")

        if refresh_key_obj is None and not self.ssm:
            raise RuntimeError("SSM client or ssmParameter.refreshTokenSSM must be available")

        try:
            ssm_config = self.ssm.get_ssm_parameter(refresh_key_obj) if self.ssm else {}
        except Exception:
            logger.exception("Failed to fetch SSM config for refreshTokenSSM")
            ssm_config = {}

        now = datetime.utcnow()
        payload = {
            "Id": id_,
            "sub": "System",
            "UserType": user_type,
            "jit": "0",
            "Role": "0",
            "ParentId": "0" if parent_id in (None, "0") else str(parent_id),
            "Service": "Service",
            "exp": now + timedelta(days=30),
        }

        signing_key = ssm_config.get("SigningKey") if isinstance(ssm_config, dict) else None
        if not signing_key:
            raise RuntimeError("SigningKey for refresh token not found in SSM")

        token = jwt.encode(payload, signing_key, algorithm="HS256")
        if isinstance(token, bytes):
            token = token.decode("utf-8")
        logger.info("inside generate_refresh_jwt_token end")
        return token

    async def verify_refresh_jwt_token(self, token: str) -> Any:
        try:
            refresh_key_obj = self.ssm_param.get("refreshTokenSSM")

            if refresh_key_obj is None and not self.ssm:
                raise RuntimeError("SSM client or ssmParameter.refreshTokenSSM must be available")

            ssm_config = self.ssm.get_ssm_parameter(refresh_key_obj) if self.ssm else {}
            signing_key = ssm_config.get("SigningKey") if isinstance(ssm_config, dict) else None
            if not signing_key:
                raise RuntimeError("SigningKey for refresh token not found in SSM")

            payload = jwt.decode(token, signing_key, algorithms=["HS256"])
            logger.info("inside verify_refresh_jwt_token end")
            return payload
        except Exception:
            logger.exception("Error in verify_refresh_jwt_token")
            raise



