from __future__ import annotations

import json
import re as _re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from mockmesh.core.logger import get_logger
from mockmesh.storage.blob import BlobStorage
from mockmesh.storage.cache import CacheStorage
from mockmesh.storage.nosql import NoSQLStorage
from mockmesh.storage.sql import SQLStorage
from mockmesh.storage.utils import prefixed_name

_log = get_logger("storage.manager")

_DEFAULT_WORKSPACE = Path.home() / ".mockmesh"

_RDS_RESOURCE_TYPES = {
    "instance":     ("rds_instances",      "rds"),
    "cluster":      ("rds_clusters",       "rds"),
    "snapshot":     ("rds_snapshots",      "rds"),
    "param_group":  ("rds_param_groups",   None),
    "subnet_group": ("rds_subnet_groups",  None),
    "option_group": ("rds_option_groups",  None),
    "proxy":        ("rds_proxies",        None),
}


def _make_rds_put(collection: str, audit_service: Optional[str], audit_op: str):
    def rds_put(self, identifier: str, attrs: Dict[str, Any]) -> None:
        self.nosql.put(collection, identifier, attrs)
        if audit_service:
            self.audit(audit_service, audit_op, identifier=identifier)
    return rds_put


def _make_rds_get(collection: str):
    def rds_get(self, identifier: str) -> Optional[Dict[str, Any]]:
        return self.nosql.get(collection, identifier)
    return rds_get


def _make_rds_delete(collection: str):
    def rds_delete(self, identifier: str) -> bool:
        return self.nosql.delete(collection, identifier)
    return rds_delete


def _make_rds_list(collection: str):
    def rds_list(self, prefix: str = "") -> List[Dict[str, Any]]:
        return [r["value"] for r in self.nosql.list(collection, prefix=prefix)]
    return rds_list


class StorageManager:

    def __init__(self, workspace: Optional[Path | str] = None) -> None:
        self.workspace = Path(workspace) if workspace else _DEFAULT_WORKSPACE
        self.sql   = SQLStorage(self.workspace)
        self.nosql = NoSQLStorage(self.workspace)
        self.blob  = BlobStorage(self.workspace)
        self.cache = CacheStorage(self.workspace)
        self._audit_path = self.workspace / "audit.log"
        self._ready = False

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def initialize(self) -> None:
        self.workspace.mkdir(parents=True, exist_ok=True)
        self.sql.initialize()
        self.nosql.initialize()
        self.blob.initialize()
        self.cache.initialize()
        self._ready = True
        _log.info(
            "StorageManager ready",
            extra={"mm_workspace": str(self.workspace)},
        )

    def close(self) -> None:
        self.sql.close()
        self._ready = False
        _log.info("StorageManager closed")

    # ------------------------------------------------------------------
    # Audit log
    # ------------------------------------------------------------------

    def audit(self, service: str, operation: str, **context: Any) -> None:
        entry = {
            "ts":        datetime.now(tz=timezone.utc).isoformat(),
            "service":   service,
            "operation": operation,
            **context,
        }
        try:
            with self._audit_path.open("a", encoding="utf-8") as f:
                f.write(json.dumps(entry) + "\n")
        except OSError:
            pass

    # ------------------------------------------------------------------
    # S3-compatible helpers (blob backend)
    # ------------------------------------------------------------------

    @staticmethod
    def _s3_bucket(bucket: str) -> str:
        return prefixed_name(bucket, "s3")

    def s3_put(self, bucket: str, key: str, body: bytes,
               content_type: str = "application/octet-stream",
               metadata: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
        result = self.blob.put_object(self._s3_bucket(bucket), key, body, content_type, metadata)
        self.audit("s3", "PutObject", bucket=bucket, key=key, size=len(body))
        return result

    def s3_get(self, bucket: str, key: str) -> Optional[Dict[str, Any]]:
        result = self.blob.get_object(self._s3_bucket(bucket), key)
        self.audit("s3", "GetObject", bucket=bucket, key=key,
                   found=result is not None)
        return result

    def s3_delete(self, bucket: str, key: str) -> bool:
        ok = self.blob.delete_object(self._s3_bucket(bucket), key)
        self.audit("s3", "DeleteObject", bucket=bucket, key=key, deleted=ok)
        return ok

    def s3_list(self, bucket: str, prefix: str = "") -> List[Dict[str, Any]]:
        objects = self.blob.list_objects(self._s3_bucket(bucket), prefix)
        self.audit("s3", "ListObjectsV2", bucket=bucket, prefix=prefix,
                   count=len(objects))
        return objects

    # ------------------------------------------------------------------
    # DynamoDB-compatible helpers (nosql backend)
    # ------------------------------------------------------------------

    @staticmethod
    def _dynamo_col(table: str) -> str:
        return prefixed_name(table, "dynamodb", "--")

    def dynamo_put(self, table: str, item: Dict[str, Any]) -> None:
        flat = _flatten_dynamo(item)
        key = _dynamo_key(flat)
        self.nosql.put(self._dynamo_col(table), key, flat)
        self.audit("dynamodb", "PutItem", table=table, key=key)

    def dynamo_get(self, table: str, key_attrs: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        flat_key = _dynamo_key(_flatten_dynamo(key_attrs))
        item = self.nosql.get(self._dynamo_col(table), flat_key)
        self.audit("dynamodb", "GetItem", table=table, key=flat_key,
                   found=item is not None)
        if item is None:
            return None
        return _unflatten_dynamo(item)

    def dynamo_delete(self, table: str, key_attrs: Dict[str, Any]) -> bool:
        flat_key = _dynamo_key(_flatten_dynamo(key_attrs))
        ok = self.nosql.delete(self._dynamo_col(table), flat_key)
        self.audit("dynamodb", "DeleteItem", table=table, key=flat_key, deleted=ok)
        return ok

    def dynamo_update(
        self,
        table: str,
        key_attrs: Dict[str, Any],
        expression: str,
        attr_values: Optional[Dict[str, Any]] = None,
        attr_names:  Optional[Dict[str, str]] = None,
    ) -> None:
        existing = self.dynamo_get(table, key_attrs) or {}
        flat     = _flatten_dynamo(existing)
        updated  = apply_update_expression(
            flat,
            expression,
            attr_values or {},
            attr_names  or {},
        )
        self.nosql.put(self._dynamo_col(table), _dynamo_key(flat), updated)
        self.audit("dynamodb", "UpdateItem", table=table)

    def dynamo_scan(self, table: str) -> List[Dict[str, Any]]:
        records = self.nosql.list(self._dynamo_col(table))
        self.audit("dynamodb", "Scan", table=table, count=len(records))
        return [_unflatten_dynamo(r["value"]) for r in records]

    # ------------------------------------------------------------------
    # Azure Blob helpers (blob backend, same backend different naming)
    # ------------------------------------------------------------------

    @staticmethod
    def _az_container(container: str) -> str:
        return prefixed_name(container, "az")

    def az_blob_put(self, container: str, blob_name: str, data: bytes,
                    content_type: str = "application/octet-stream") -> Dict[str, Any]:
        result = self.blob.put_object(self._az_container(container), blob_name, data, content_type)
        self.audit("azure_blob", "upload_blob", container=container,
                   blob=blob_name, size=len(data))
        return result

    def az_blob_get(self, container: str, blob_name: str) -> Optional[Dict[str, Any]]:
        result = self.blob.get_object(self._az_container(container), blob_name)
        self.audit("azure_blob", "download_blob", container=container,
                   blob=blob_name, found=result is not None)
        return result

    def az_blob_delete(self, container: str, blob_name: str) -> bool:
        ok = self.blob.delete_object(self._az_container(container), blob_name)
        self.audit("azure_blob", "delete_blob", container=container,
                   blob=blob_name, deleted=ok)
        return ok

    # ------------------------------------------------------------------
    # GCS-compatible helpers (blob backend)
    # ------------------------------------------------------------------

    @staticmethod
    def _gcs_bucket(bucket: str) -> str:
        return prefixed_name(bucket, "gcs")

    def gcs_put(self, bucket: str, key: str, body: bytes,
                content_type: str = "application/octet-stream") -> Dict[str, Any]:
        result = self.blob.put_object(self._gcs_bucket(bucket), key, body, content_type)
        self.audit("gcs", "insert", bucket=bucket, key=key, size=len(body))
        return result

    def gcs_get(self, bucket: str, key: str) -> Optional[Dict[str, Any]]:
        result = self.blob.get_object(self._gcs_bucket(bucket), key)
        self.audit("gcs", "get", bucket=bucket, key=key, found=result is not None)
        return result

    def gcs_delete(self, bucket: str, key: str) -> bool:
        ok = self.blob.delete_object(self._gcs_bucket(bucket), key)
        self.audit("gcs", "delete", bucket=bucket, key=key, deleted=ok)
        return ok

    def gcs_list(self, bucket: str, prefix: str = "") -> List[Dict[str, Any]]:
        objects = self.blob.list_objects(self._gcs_bucket(bucket), prefix)
        self.audit("gcs", "list", bucket=bucket, prefix=prefix, count=len(objects))
        return objects

    # ------------------------------------------------------------------
    # CosmosDB-compatible helpers (nosql backend)
    # ------------------------------------------------------------------

    @staticmethod
    def _cosmos_col(container: str) -> str:
        return prefixed_name(container, "cosmos", "--")

    def cosmos_upsert(self, container: str, item: Dict[str, Any]) -> Dict[str, Any]:
        key = str(item.get("id", item.get("_id", "unknown")))
        self.nosql.put(self._cosmos_col(container), key, item)
        self.audit("cosmosdb", "upsert_item", container=container, id=key)
        return {**item, "_etag": f'"{key[:8]}-mock"'}

    def cosmos_get(self, container: str, item_id: str) -> Optional[Dict[str, Any]]:
        item = self.nosql.get(self._cosmos_col(container), item_id)
        self.audit("cosmosdb", "read_item", container=container, id=item_id,
                   found=item is not None)
        return item

    # ------------------------------------------------------------------
    # SQS / messaging helper (nosql backend, queue semantics)
    # ------------------------------------------------------------------

    def sqs_enqueue(self, queue_name: str, message_id: str,
                    body: str) -> None:
        self.nosql.put(f"sqs_{queue_name}", message_id,
                       {"MessageId": message_id, "Body": body, "status": "queued"})
        self.audit("sqs", "SendMessage", queue=queue_name, message_id=message_id)

    def sqs_dequeue(self, queue_name: str) -> Optional[Dict[str, Any]]:
        records = self.nosql.query(f"sqs_{queue_name}", status="queued")
        if not records:
            return None
        record = records[0]
        updated = {**record["value"], "status": "in-flight"}
        self.nosql.put(f"sqs_{queue_name}", record["key"], updated)
        self.audit("sqs", "ReceiveMessage", queue=queue_name,
                   message_id=record["key"])
        return updated

    # ------------------------------------------------------------------
    # MongoDB helpers (nosql backend)
    # ------------------------------------------------------------------

    def mongo_insert(self, db: str, collection: str, doc_id: str,
                     document: Dict[str, Any]) -> None:
        col_key = f"mongo--{db}--{collection}"
        self.nosql.put(col_key, doc_id, document)
        self.audit("mongodb", "insert", db=db, collection=collection, id=doc_id)

    def mongo_find(self, db: str, collection: str,
                   key: str = "") -> Optional[Dict[str, Any]]:
        col_key = f"mongo--{db}--{collection}"
        if key:
            return self.nosql.get(col_key, key)
        return None

    def mongo_find_all(self, db: str, collection: str) -> List[Dict[str, Any]]:
        col_key = f"mongo--{db}--{collection}"
        return [r["value"] for r in self.nosql.list(col_key)]

    def mongo_delete(self, db: str, collection: str, key: str) -> bool:
        col_key = f"mongo--{db}--{collection}"
        ok = self.nosql.delete(col_key, key)
        self.audit("mongodb", "delete", db=db, collection=collection, id=key)
        return ok

    # ------------------------------------------------------------------
    # Redis helpers (cache backend)
    # ------------------------------------------------------------------

    def redis_set(self, db: int, key: str, value: Dict[str, Any]) -> None:
        col = f"redis_{db}"
        self.cache.put(col, key, value)

    def redis_get(self, db: int, key: str) -> Optional[Dict[str, Any]]:
        col = f"redis_{db}"
        return self.cache.get(col, key)

    def redis_delete(self, db: int, key: str) -> bool:
        col = f"redis_{db}"
        return self.cache.delete(col, key)

    def redis_keys(self, db: int) -> List[str]:
        col = f"redis_{db}"
        return [r["key"] for r in self.cache.list(col)]

    def redis_flush(self, db: int) -> None:
        col = f"redis_{db}"
        for r in self.cache.list(col):
            self.cache.delete(col, r["key"])

    # ------------------------------------------------------------------
    # Generic backend access
    # ------------------------------------------------------------------

    def get_backend(self, backend_type: str):
        return {"sql": self.sql, "nosql": self.nosql,
                "blob": self.blob, "cache": self.cache}.get(backend_type)

    def list_collections(self) -> list:
        return self.nosql.list_collections()

    # ------------------------------------------------------------------
    # SSM Parameter Store helpers (nosql backend)
    # ------------------------------------------------------------------

    def ssm_put(self, name: str, value: str, param_type: str = "String",
                version: int = 1) -> None:
        self.nosql.put("ssm_params", name,
                       {"Name": name, "Value": value, "Type": param_type,
                        "Version": version})
        self.audit("ssm", "PutParameter", name=name, type=param_type)

    def ssm_get(self, name: str) -> Optional[Dict[str, Any]]:
        item = self.nosql.get("ssm_params", name)
        self.audit("ssm", "GetParameter", name=name, found=item is not None)
        return item

    def ssm_get_by_path(self, path: str) -> List[Dict[str, Any]]:
        all_records = self.nosql.list("ssm_params")
        result = [r["value"] for r in all_records
                  if r["value"].get("Name", "").startswith(path)]
        self.audit("ssm", "GetParametersByPath", path=path, count=len(result))
        return result

    # ------------------------------------------------------------------
    # SecretsManager helpers (nosql backend)
    # ------------------------------------------------------------------

    def secret_put(self, name: str, secret_string: str,
                   arn: Optional[str] = None) -> Dict[str, Any]:
        if arn is None:
            arn = f"arn:aws:secretsmanager:us-east-1:123456789012:secret:{name}"
        entry = {"Name": name, "ARN": arn, "SecretString": secret_string,
                 "VersionId": "a1b2c3d4-e5f6-7890-abcd-ef1234567890"}
        self.nosql.put("secretsmanager", name, entry)
        self.audit("secretsmanager", "PutSecret", name=name)
        return entry

    def secret_get(self, secret_id: str) -> Optional[Dict[str, Any]]:
        item = self.nosql.get("secretsmanager", secret_id)
        self.audit("secretsmanager", "GetSecretValue", name=secret_id,
                   found=item is not None)
        return item

    # ------------------------------------------------------------------
    # ElastiCache helpers (cache backend)
    # ------------------------------------------------------------------

    def cache_cluster_put(self, cluster_id: str, attrs: Dict[str, Any]) -> None:
        self.cache.put("clusters", cluster_id, attrs)
        self.audit("elasticache", "CreateCacheCluster", cluster_id=cluster_id)

    def cache_cluster_get(self, cluster_id: str) -> Optional[Dict[str, Any]]:
        item = self.cache.get("clusters", cluster_id)
        self.audit("elasticache", "DescribeCacheCluster", cluster_id=cluster_id,
                   found=item is not None)
        return item

    def cache_cluster_list(self) -> List[Dict[str, Any]]:
        records = self.cache.list("clusters")
        return [r["value"] for r in records]

    def cache_rg_put(self, rg_id: str, attrs: Dict[str, Any]) -> None:
        self.cache.put("replication_groups", rg_id, attrs)
        self.audit("elasticache", "CreateReplicationGroup", rg_id=rg_id)

    def cache_rg_get(self, rg_id: str) -> Optional[Dict[str, Any]]:
        item = self.cache.get("replication_groups", rg_id)
        self.audit("elasticache", "DescribeReplicationGroup", rg_id=rg_id,
                   found=item is not None)
        return item

    def cache_rg_list(self) -> List[Dict[str, Any]]:
        records = self.cache.list("replication_groups")
        return [r["value"] for r in records]

    # ------------------------------------------------------------------
    # RDS helpers (nosql backend, generated via factory)
    # ------------------------------------------------------------------


for _rtype, (_collection, _audit_svc) in _RDS_RESOURCE_TYPES.items():
    _audit_op = f"put_{_rtype}"
    setattr(StorageManager, f"rds_{_rtype}_put",
            _make_rds_put(_collection, _audit_svc, _audit_op))
    setattr(StorageManager, f"rds_{_rtype}_get",
            _make_rds_get(_collection))
    setattr(StorageManager, f"rds_{_rtype}_delete",
            _make_rds_delete(_collection))
    setattr(StorageManager, f"rds_{_rtype}_list",
            _make_rds_list(_collection))


# ---------------------------------------------------------------------------
# DynamoDB attribute format helpers
# ---------------------------------------------------------------------------

def _flatten_dynamo(item: Dict[str, Any]) -> Dict[str, Any]:
    result = {}
    for k, v in item.items():
        if isinstance(v, dict) and len(v) == 1:
            type_key = next(iter(v))
            if type_key in ("S", "N", "BOOL", "NULL", "B"):
                raw = v[type_key]
                result[k] = float(raw) if type_key == "N" else raw
                continue
        result[k] = v
    return result


def _unflatten_dynamo(item: Dict[str, Any]) -> Dict[str, Any]:
    result = {}
    for k, v in item.items():
        if isinstance(v, bool):
            result[k] = {"BOOL": v}
        elif isinstance(v, (int, float)):
            result[k] = {"N": str(v)}
        elif v is None:
            result[k] = {"NULL": True}
        else:
            result[k] = {"S": str(v)}
    return result


def _dynamo_key(flat: Dict[str, Any]) -> str:
    for candidate in ("id", "pk", "key", "Id", "PK"):
        if candidate in flat:
            return str(flat[candidate])
    return str(flat.get(sorted(flat.keys())[0], "unknown")) if flat else "unknown"


# ---------------------------------------------------------------------------
# DynamoDB UpdateExpression interpreter
# ---------------------------------------------------------------------------

def apply_update_expression(
    existing_flat: Dict[str, Any],
    expression: str,
    attr_values: Dict[str, Any],
    attr_names: Dict[str, str],
) -> Dict[str, Any]:
    if not expression:
        return dict(existing_flat)

    flat_vals = _flatten_dynamo(attr_values) if attr_values else {}
    result    = dict(existing_flat)

    set_m = _re.search(
        r'\bSET\b(.*?)(?:\bREMOVE\b|\bADD\b|\bDELETE\b|$)',
        expression, _re.IGNORECASE | _re.DOTALL,
    )
    if set_m:
        for assignment in set_m.group(1).split(','):
            parts = assignment.split('=', 1)
            if len(parts) != 2:
                continue
            lhs = parts[0].strip()
            rhs = parts[1].strip()
            field = attr_names.get(lhs, lhs)
            if rhs in flat_vals:
                result[field] = flat_vals[rhs]

    rem_m = _re.search(
        r'\bREMOVE\b(.*?)(?:\bSET\b|\bADD\b|\bDELETE\b|$)',
        expression, _re.IGNORECASE | _re.DOTALL,
    )
    if rem_m:
        for token in rem_m.group(1).split(','):
            field = attr_names.get(token.strip(), token.strip())
            result.pop(field, None)

    return result
