from __future__ import annotations

from typing import TYPE_CHECKING, Optional

from mockmesh.interceptors.common import parse_json_body

if TYPE_CHECKING:
    from mockmesh.storage.manager import StorageManager


def handle_dynamodb(sm: "StorageManager", req: dict, operation: str) -> Optional[dict]:
    import json
    from mockmesh.storage.manager import apply_update_expression  # noqa: F401

    bd    = parse_json_body(req)
    table = bd.get("TableName", "unknown")

    if operation == "PutItem":
        item = bd.get("Item", {})
        if item:
            sm.dynamo_put(table, item)
        return {}

    if operation == "GetItem":
        item = sm.dynamo_get(table, bd.get("Key", {}))
        return {"Item": item} if item else {}

    if operation == "DeleteItem":
        sm.dynamo_delete(table, bd.get("Key", {}))
        return {}

    if operation == "UpdateItem":
        sm.dynamo_update(
            table=table,
            key_attrs=bd.get("Key", {}),
            expression=bd.get("UpdateExpression", ""),
            attr_values=bd.get("ExpressionAttributeValues", {}),
            attr_names=bd.get("ExpressionAttributeNames", {}),
        )
        return {}

    if operation in ("Query", "Scan"):
        items = sm.dynamo_scan(table)
        return {"Items": items, "Count": len(items), "ScannedCount": len(items)}

    if operation == "BatchWriteItem":
        for tbl, requests in bd.get("RequestItems", {}).items():
            for entry in requests:
                if "PutRequest" in entry:
                    sm.dynamo_put(tbl, entry["PutRequest"]["Item"])
                elif "DeleteRequest" in entry:
                    sm.dynamo_delete(tbl, entry["DeleteRequest"]["Key"])
        return {"UnprocessedItems": {}}

    if operation == "BatchGetItem":
        responses: dict = {}
        for tbl, spec in bd.get("RequestItems", {}).items():
            items = [sm.dynamo_get(tbl, k) for k in spec.get("Keys", [])]
            responses[tbl] = [i for i in items if i]
        return {"Responses": responses, "UnprocessedKeys": {}}

    if operation == "TransactWriteItems":
        for tx in bd.get("TransactItems", []):
            if "Put" in tx:
                p = tx["Put"]
                sm.dynamo_put(p["TableName"], p["Item"])
            elif "Delete" in tx:
                d = tx["Delete"]
                sm.dynamo_delete(d["TableName"], d["Key"])
            elif "Update" in tx:
                u = tx["Update"]
                sm.dynamo_update(
                    table=u["TableName"],
                    key_attrs=u.get("Key", {}),
                    expression=u.get("UpdateExpression", ""),
                    attr_values=u.get("ExpressionAttributeValues", {}),
                    attr_names=u.get("ExpressionAttributeNames", {}),
                )
        return {}

    if operation == "TransactGetItems":
        responses = []
        for tx in bd.get("TransactItems", []):
            if "Get" in tx:
                g    = tx["Get"]
                item = sm.dynamo_get(g["TableName"], g["Key"])
                responses.append({"Item": item} if item else {})
        return {"Responses": responses}

    if operation == "DescribeTable":
        items = sm.nosql.list(f"dynamodb--{table}")
        return {
            "Table": {
                "TableName":   table,
                "TableStatus": "ACTIVE",
                "ItemCount":   len(items),
                "TableSizeBytes": sum(
                    len(json.dumps(r.get("value", {}))) for r in items
                ),
                "ProvisionedThroughput": {
                    "ReadCapacityUnits":  5,
                    "WriteCapacityUnits": 5,
                },
            }
        }

    return None
