"""DynamoDB service emulator."""
import copy, json, logging, time, uuid
from dataclasses import dataclass, field
from typing import Any, Optional
from flask import Request, Response
from localrun.config import get_config
from localrun.utils import error_response, iso_timestamp, new_request_id

logger = logging.getLogger("localrun.dynamodb")

def _json_resp(data, status=200):
    return Response(json.dumps(data, default=str), status=status, content_type="application/x-amz-json-1.0")

def _json_err(code, msg, status=400):
    return Response(json.dumps({"__type": code, "message": msg}), status=status, content_type="application/x-amz-json-1.0")

class DynamoDBService:
    def __init__(self):
        self.tables: dict[str, dict] = {}  # table_name -> table_meta
        self.table_items: dict[str, list[dict]] = {}  # table_name -> list of items

    def handle(self, req: Request, path: str) -> Response:
        target = req.headers.get("X-Amz-Target", "")
        action = target.split(".")[-1] if "." in target else ""
        body = {}
        try:
            body = json.loads(req.get_data(as_text=True) or "{}")
        except (json.JSONDecodeError, Exception):
            pass
        actions = {
            "CreateTable": self._create_table, "DeleteTable": self._delete_table,
            "ListTables": self._list_tables, "DescribeTable": self._describe_table,
            "PutItem": self._put_item, "GetItem": self._get_item,
            "DeleteItem": self._delete_item, "UpdateItem": self._update_item,
            "Query": self._query, "Scan": self._scan,
            "BatchWriteItem": self._batch_write, "BatchGetItem": self._batch_get,
        }
        handler = actions.get(action)
        if not handler:
            return _json_err("UnknownOperationException", f"Unknown operation: {action}", 400)
        return handler(body)

    def _table_meta(self, name):
        t = self.tables.get(name)
        if not t: return None
        count = len(self.table_items.get(name, []))
        return {
            "TableName": t["TableName"], "TableStatus": "ACTIVE",
            "KeySchema": t["KeySchema"],
            "AttributeDefinitions": t["AttributeDefinitions"],
            "TableArn": t["TableArn"], "ItemCount": count,
            "TableSizeBytes": count * 100,
            "CreationDateTime": t["CreationDateTime"],
            "ProvisionedThroughput": t.get("ProvisionedThroughput", {"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}),
        }

    def _create_table(self, body):
        name = body.get("TableName", "")
        if not name: return _json_err("ValidationException", "TableName required")
        if name in self.tables: return _json_err("ResourceInUseException", f"Table {name} already exists")
        c = get_config()
        meta = {
            "TableName": name, "KeySchema": body.get("KeySchema", []),
            "AttributeDefinitions": body.get("AttributeDefinitions", []),
            "TableArn": f"arn:aws:dynamodb:{c.region}:{c.account_id}:table/{name}",
            "CreationDateTime": time.time(),
            "ProvisionedThroughput": body.get("ProvisionedThroughput", {"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}),
        }
        self.tables[name] = meta; self.table_items[name] = []
        logger.info("Created table: %s", name)
        return _json_resp({"TableDescription": self._table_meta(name)})

    def _delete_table(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        meta = self._table_meta(name)
        del self.tables[name]; self.table_items.pop(name, None)
        logger.info("Deleted table: %s", name)
        return _json_resp({"TableDescription": meta})

    def _list_tables(self, body):
        names = sorted(self.tables.keys())
        limit = body.get("Limit", 100)
        return _json_resp({"TableNames": names[:limit]})

    def _describe_table(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        return _json_resp({"Table": self._table_meta(name)})

    def _get_key_attrs(self, table_name):
        t = self.tables[table_name]
        return [ks["AttributeName"] for ks in t["KeySchema"]]

    def _match_key(self, item, key):
        for k, v in key.items():
            if k not in item: return False
            if item[k] != v: return False
        return True

    def _find_item_idx(self, table_name, key):
        items = self.table_items.get(table_name, [])
        for i, item in enumerate(items):
            if self._match_key(item, key): return i
        return -1

    def _put_item(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        item = body.get("Item", {})
        key_attrs = self._get_key_attrs(name)
        key = {k: item[k] for k in key_attrs if k in item}
        idx = self._find_item_idx(name, key)
        old = None
        if idx >= 0:
            old = self.table_items[name][idx]
            self.table_items[name][idx] = item
        else:
            self.table_items[name].append(item)
        result = {}
        if body.get("ReturnValues") == "ALL_OLD" and old:
            result["Attributes"] = old
        return _json_resp(result)

    def _get_item(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        key = body.get("Key", {})
        idx = self._find_item_idx(name, key)
        if idx < 0: return _json_resp({})
        item = self.table_items[name][idx]
        proj = body.get("ProjectionExpression")
        if proj:
            attrs = [a.strip() for a in proj.split(",")]
            item = {k: v for k, v in item.items() if k in attrs}
        return _json_resp({"Item": item})

    def _delete_item(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        key = body.get("Key", {})
        idx = self._find_item_idx(name, key)
        old = None
        if idx >= 0: old = self.table_items[name].pop(idx)
        result = {}
        if body.get("ReturnValues") == "ALL_OLD" and old:
            result["Attributes"] = old
        return _json_resp(result)

    def _update_item(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        key = body.get("Key", {})
        idx = self._find_item_idx(name, key)
        if idx < 0:
            item = dict(key); self.table_items[name].append(item); idx = len(self.table_items[name]) - 1
        else:
            item = self.table_items[name][idx]
        # Simple SET expression support
        expr = body.get("UpdateExpression", "")
        values = body.get("ExpressionAttributeValues", {})
        names = body.get("ExpressionAttributeNames", {})
        if expr.upper().startswith("SET "):
            parts = expr[4:].split(",")
            for part in parts:
                part = part.strip()
                if "=" in part:
                    lhs, rhs = part.split("=", 1)
                    lhs = lhs.strip(); rhs = rhs.strip()
                    attr = names.get(lhs, lhs)
                    val = values.get(rhs, {"S": rhs})
                    item[attr] = val
        elif expr.upper().startswith("REMOVE "):
            parts = expr[7:].split(",")
            for part in parts:
                attr = names.get(part.strip(), part.strip())
                item.pop(attr, None)
        self.table_items[name][idx] = item
        result = {}
        rv = body.get("ReturnValues", "NONE")
        if rv == "ALL_NEW": result["Attributes"] = item
        elif rv == "ALL_OLD": result["Attributes"] = item
        return _json_resp(result)

    def _eval_condition(self, item, cond_expr, cond_val, attr_names=None):
        """Very basic condition evaluation for KeyConditionExpression."""
        if not cond_expr: return True
        # Replace attribute names
        expr = cond_expr
        if attr_names:
            for placeholder, real_name in attr_names.items():
                expr = expr.replace(placeholder, real_name)
        # Simple: attr = :val
        if " = " in expr:
            parts = expr.split(" AND ")
            for part in parts:
                part = part.strip()
                if " = " in part:
                    lhs, rhs = part.split(" = ", 1)
                    lhs = lhs.strip(); rhs = rhs.strip()
                    expected = cond_val.get(rhs)
                    actual = item.get(lhs)
                    if actual != expected: return False
                elif " begins_with" in part.lower():
                    # begins_with(attr, :val)
                    import re
                    m = re.match(r'begins_with\((\w+),\s*(\S+)\)', part, re.IGNORECASE)
                    if m:
                        attr, valref = m.group(1), m.group(2)
                        expected = cond_val.get(valref, {})
                        actual = item.get(attr, {})
                        exp_s = expected.get("S", "")
                        act_s = actual.get("S", "")
                        if not act_s.startswith(exp_s): return False
            return True
        return True

    def _query(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        items = self.table_items.get(name, [])
        expr = body.get("KeyConditionExpression", "")
        values = body.get("ExpressionAttributeValues", {})
        attr_names = body.get("ExpressionAttributeNames", {})
        result = [i for i in items if self._eval_condition(i, expr, values, attr_names)]
        limit = body.get("Limit")
        if limit: result = result[:limit]
        if not body.get("ScanIndexForward", True): result.reverse()
        return _json_resp({"Items": result, "Count": len(result), "ScannedCount": len(result)})

    def _scan(self, body):
        name = body.get("TableName", "")
        if name not in self.tables: return _json_err("ResourceNotFoundException", f"Table {name} not found")
        items = list(self.table_items.get(name, []))
        limit = body.get("Limit")
        if limit: items = items[:limit]
        return _json_resp({"Items": items, "Count": len(items), "ScannedCount": len(items)})

    def _batch_write(self, body):
        requests = body.get("RequestItems", {})
        for table_name, ops in requests.items():
            if table_name not in self.tables: continue
            for op in ops:
                if "PutRequest" in op:
                    self._put_item({"TableName": table_name, "Item": op["PutRequest"]["Item"]})
                elif "DeleteRequest" in op:
                    self._delete_item({"TableName": table_name, "Key": op["DeleteRequest"]["Key"]})
        return _json_resp({"UnprocessedItems": {}})

    def _batch_get(self, body):
        requests = body.get("RequestItems", {})
        responses = {}
        for table_name, spec in requests.items():
            if table_name not in self.tables: continue
            keys = spec.get("Keys", [])
            items = []
            for key in keys:
                idx = self._find_item_idx(table_name, key)
                if idx >= 0: items.append(self.table_items[table_name][idx])
            responses[table_name] = items
        return _json_resp({"Responses": responses, "UnprocessedKeys": {}})
