# ws_package/controllers/server_controller.py

import asyncio
import base64
from typing import Any

from ws_package.services.server_service import ServerService
from ws_package.entities.s3bucket_entity import S3BucketEntity
from ws_package.entities.user_entity import UserEntity
from ws_package.log.writer import LogWriter
from ws_package.configs import config
from ws_package.configs.security import Security_config
from ws_package.utils.encryption import decrypt_data, decrypt_data_key_api_payload
from tzlocal import get_localzone

logger = LogWriter.file(__name__)

class ServerController:
    def __init__(self):
        self.server_service = ServerService()  # Instance-specific state to prevent shared mutable state
        self.encryption_salt = Security_config('encryption_salt', 'Security')
    def decrypt(self, encrypted_key) -> str:
        try:
            position_of_dollar = encrypted_key.find('$')
            new_string = encrypted_key[position_of_dollar + 1:]
            new_string = new_string[::-1]
            decrypted_data = base64.b64decode(new_string).decode()
            return decrypted_data
        except Exception as e:
            logger.error("Error in Encrypted Salt: %s", e, exc_info=True)

    def _parse_s3_api_response(self, response_data: Any) -> dict | None:
        """Normalize sts/keys API output to a credentials dict."""
        if isinstance(response_data, str):
            return decrypt_data_key_api_payload(response_data)

        if isinstance(response_data, dict):
            code = response_data.get("code")
            if code is not None and str(code) not in ("200", "201"):
                logger.error(
                    "[S3Bucket] API returned error code=%s message=%s",
                    code,
                    response_data.get("message") or response_data.get("response"),
                )
                return None

            if "ACCESS_KEY" in response_data:
                return response_data

            nested = response_data.get("data")
            if nested is not None:
                return self._parse_s3_api_response(nested)

            nested_response = response_data.get("response")
            if isinstance(nested_response, dict):
                return self._parse_s3_api_response(nested_response)

        logger.error(
            "[S3Bucket] Unrecognized API response type: %s",
            type(response_data).__name__,
        )
        return None

    def _extract_expire_time(self, response_data: dict) -> str:
        for key in (
            "AWS_EXPIRY_TIME_FORMATTED",
            "expire_time",
            "EXPIRE_TIME",
            "AWS_EXPIRY_TIME",
        ):
            value = response_data.get(key)
            if value:
                return str(value).strip()
        return ""

    def _extract_session_token(self, response_data: dict, bucket: str) -> str:
        raw = (
            response_data.get("SESSION_TOKEN")
            or response_data.get("session_token")
            or response_data.get("AWS_SESSION_TOKEN")
            or ""
        )
        if not raw:
            return ""
        text = str(raw).strip()
        if bucket == "hetzner" and "$" in text:
            try:
                return str(self.decrypt(text)).strip()
            except Exception:
                return self._normalize_cred_value(text)
        return self._normalize_cred_value(text)

    def build_s3_entity_from_api_response(
        self, response_data: Any, bucket: str
    ) -> S3BucketEntity | None:
        credentials = self._parse_s3_api_response(response_data)
        if credentials is None:
            return None
        return self._entity_from_api_response(credentials, bucket)

    def _entity_from_api_response(self, response_data: dict, bucket: str) -> S3BucketEntity:
        session_token = self._extract_session_token(response_data, bucket)
        expire_time = self._extract_expire_time(response_data)
        if bucket == "hetzner":
            entity = S3BucketEntity(
                access_key=self.decrypt(response_data['ACCESS_KEY']),
                secret_access_key=self.decrypt(response_data['SECRET_ACCESS_KEY']),
                region=self.decrypt(response_data['DEFAULT_REGION']),
                bucket=self.decrypt(response_data['BUCKET']),
                session_token=session_token,
                expire_time=expire_time,
            )
        else:
            entity = S3BucketEntity(
                access_key=response_data['ACCESS_KEY'],
                secret_access_key=response_data['SECRET_ACCESS_KEY'],
                region=response_data['DEFAULT_REGION'],
                bucket=response_data['BUCKET'],
                session_token=session_token,
                expire_time=expire_time,
            )
        return self._normalize_s3_entity(entity)

    def _normalize_cred_value(self, value) -> str:
        """Normalize a credential field to plaintext for reliable compare/persist."""
        if value is None:
            return ""
        text = str(value).strip()
        if not text:
            return ""
        if "$" in text:
            try:
                return decrypt_data(text).strip()
            except Exception:
                try:
                    decrypted = self.decrypt(text)
                    if decrypted:
                        return str(decrypted).strip()
                except Exception:
                    pass
        return text

    def _normalize_s3_entity(self, entity: S3BucketEntity) -> S3BucketEntity:
        return S3BucketEntity(
            access_key=self._normalize_cred_value(entity.access_key),
            secret_access_key=self._normalize_cred_value(entity.secret_access_key),
            region=self._normalize_cred_value(entity.region),
            bucket=self._normalize_cred_value(entity.bucket),
            session_token=self._normalize_cred_value(entity.session_token),
            expire_time=str(entity.expire_time or "").strip(),
        )

    def _config_to_s3_entity(self, stored) -> S3BucketEntity:
        """Convert AppConfig row to normalized S3BucketEntity for comparison."""
        from ws_package.models.base import _decrypt_app_config_fields

        if stored is None:
            return S3BucketEntity(
                access_key="",
                secret_access_key="",
                region="",
                bucket="",
                session_token="",
                expire_time="",
            )
        try:
            _decrypt_app_config_fields(stored)
        except Exception as e:
            logger.warning("Could not decrypt stored app_config for compare: %s", e)
        return S3BucketEntity(
            access_key=self._normalize_cred_value(stored.access_key),
            secret_access_key=self._normalize_cred_value(stored.secret_access_key),
            region=self._normalize_cred_value(stored.region),
            bucket=self._normalize_cred_value(stored.bucket),
            session_token=self._normalize_cred_value(getattr(stored, "session_token", "") or ""),
            expire_time=str(getattr(stored, "expire_time", "") or "").strip(),
        )

    def _s3_credentials_changed(self, stored, s3bucket_entity: S3BucketEntity) -> bool:
        stored_entity = self._config_to_s3_entity(stored)
        incoming = self._normalize_s3_entity(s3bucket_entity)
        changed = (
            stored_entity.access_key != incoming.access_key
            or stored_entity.secret_access_key != incoming.secret_access_key
            or stored_entity.region != incoming.region
            or stored_entity.bucket != incoming.bucket
            or stored_entity.session_token != incoming.session_token
            or stored_entity.expire_time != incoming.expire_time
        )
        if changed:
            logger.info(
                "S3 credential diff for org compare — access_key=%s region=%s bucket=%s session=%s expire=%s",
                "changed" if stored_entity.access_key != incoming.access_key else "same",
                "changed" if stored_entity.region != incoming.region else "same",
                "changed" if stored_entity.bucket != incoming.bucket else "same",
                "changed" if stored_entity.session_token != incoming.session_token else "same",
                "changed" if stored_entity.expire_time != incoming.expire_time else "same",
            )
        return changed

    async def persist_s3bucket_if_changed(
        self, s3bucket_entity: S3BucketEntity, org_id: int
    ) -> S3BucketEntity:
        org_id = int(org_id)
        s3bucket_entity = self._normalize_s3_entity(s3bucket_entity)
        existing = await self.server_service.get_s3bucket(org_id)
        if existing is None:
            return await self.__save_s3bucket_entity(s3bucket_entity, org_id)

        if not self._s3_credentials_changed(existing, s3bucket_entity):
            logger.debug(
                "S3 credentials unchanged for org_id %s — skipping database update",
                org_id,
            )
            return s3bucket_entity

        logger.info(
            "S3 credentials changed for org_id %s — updating local database (expire_time=%s)",
            org_id,
            s3bucket_entity.expire_time or "n/a",
        )
        saved = await self.__save_s3bucket_entity(s3bucket_entity, org_id)
        try:
            from ws_package.utils.s3 import clear_s3_credential_cache
            clear_s3_credential_cache(org_id)
        except Exception as cache_err:
            logger.warning(
                "Could not clear S3 credential cache for org_id %s: %s",
                org_id,
                cache_err,
            )
        return saved

    async def ensure_fresh_s3_bucket(self, user_entity: UserEntity) -> S3BucketEntity | None:
        """Return S3 credentials, calling sts/keys when missing or past expire_time."""
        from ws_package.utils.s3 import (
            clear_s3_credential_cache,
            s3_credentials_expired,
            warm_s3_credential_cache,
        )

        if user_entity is None:
            logger.warning("[S3Bucket] Cannot ensure fresh credentials: user_entity is None")
            return None

        org_id = int(user_entity.org_id)
        existing = await self.server_service.get_s3bucket(org_id)
        if existing is not None and not s3_credentials_expired(existing):
            warm_s3_credential_cache(org_id, existing)
            return self._config_to_s3_entity(existing)

        expire_label = getattr(existing, "expire_time", "") if existing else "missing"
        logger.info(
            "[S3Bucket] STS credentials expired or missing for org_id=%s (expire_time=%s) — calling sts/keys",
            org_id,
            expire_label or "n/a",
        )
        clear_s3_credential_cache(org_id)
        entity = await self.get_s3_bucket_details(user_entity)
        if entity is None:
            return None

        stored = await self.server_service.get_s3bucket(org_id)
        if stored is not None:
            warm_s3_credential_cache(org_id, stored)
        return entity

    async def get_s3_bucket_details(self, user_entity: UserEntity) -> S3BucketEntity | None:
        self.server_service.application_s3bucket_detail_url = config('application_s3bucket_detail_url', 'Application')
        sd_token = user_entity.sd_token
        org_id = user_entity.org_id
        user_id = user_entity.user_id
        bucket = user_entity.bucket
        self.server_service.headers = {
            'SDToken': str(sd_token), 
            'OrgID': str(org_id), 
            'UserID': str(user_id), 
            'bucket': str(bucket),
            'Authorization': f'Bearer {sd_token}', 
            'deviceType': 'Linux'
        }
        current_timezone = str(get_localzone())
        # Match activity/log: deprecated Asia/Calcutta → Asia/Kolkata
        if current_timezone == "Asia/Calcutta":
            current_timezone = "Asia/Kolkata"
        payload = {
            "org_id": org_id,
            "appViewType": "ws_py_tracker",
            "current_timezone": current_timezone,
        }
        response_data = self.server_service.fetch_s3bucket_details(payload)
        if response_data is None:
            return None

        s3bucket_entity = self.build_s3_entity_from_api_response(response_data, bucket)
        if s3bucket_entity is None:
            logger.error("[S3Bucket] Failed to parse credentials for org_id=%s", org_id)
            return None
        logger.info(
            "[S3Bucket] Parsed STS credentials for org_id=%s expire_time=%s has_session_token=%s",
            org_id,
            s3bucket_entity.expire_time or "n/a",
            bool(s3bucket_entity.session_token),
        )
        return await self.persist_s3bucket_if_changed(s3bucket_entity, org_id)
    
    async def __save_s3bucket_entity(self, s3bucket_entity: S3BucketEntity, org_id: int):
        return await self.server_service.save_s3bucket(s3bucket_entity, org_id)



