"""SQS service emulator — supports both XML (query) and JSON (amz-json) protocols."""
import hashlib, json, logging, time, uuid
from dataclasses import dataclass, field
from typing import Optional
from urllib.parse import parse_qs
from flask import Request, Response
from localrun.config import get_config
from localrun.utils import error_response, md5_hex, new_message_id, new_request_id

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

@dataclass
class SQSMessage:
    message_id: str; body: str; receipt_handle: Optional[str] = None
    md5_of_body: str = ""; sent_timestamp: int = 0; visible_after: float = 0; receive_count: int = 0
    def __post_init__(self):
        if not self.md5_of_body: self.md5_of_body = md5_hex(self.body.encode())
        if not self.sent_timestamp: self.sent_timestamp = int(time.time() * 1000)

@dataclass
class SQSQueue:
    name: str; url: str; arn: str
    attributes: dict = field(default_factory=dict)
    messages: list = field(default_factory=list)
    tags: dict = field(default_factory=dict)
    def __post_init__(self):
        ts = str(int(time.time()))
        defaults = {"VisibilityTimeout": "30", "DelaySeconds": "0", "MaximumMessageSize": "262144",
                     "MessageRetentionPeriod": "345600", "ReceiveMessageWaitTimeSeconds": "0",
                     "ApproximateNumberOfMessages": "0", "ApproximateNumberOfMessagesNotVisible": "0",
                     "CreatedTimestamp": ts, "LastModifiedTimestamp": ts, "QueueArn": self.arn}
        for k, v in defaults.items(): self.attributes.setdefault(k, v)

class SQSService:
    def __init__(self):
        self.queues: dict[str, SQSQueue] = {}

    def _is_json_protocol(self, req):
        ct = req.content_type or ""
        return "json" in ct

    def handle(self, req: Request, path: str) -> Response:
        amz_target = req.headers.get("X-Amz-Target", "")
        if amz_target and "." in amz_target:
            action = amz_target.split(".")[-1]
        else:
            action = req.args.get("Action") or req.form.get("Action", "")
            if not action:
                params = parse_qs(req.get_data(as_text=True))
                action = params.get("Action", [""])[0]
        h = {"CreateQueue": self._create_queue, "DeleteQueue": self._delete_queue,
             "ListQueues": self._list_queues, "GetQueueUrl": self._get_queue_url,
             "GetQueueAttributes": self._get_queue_attributes, "SetQueueAttributes": self._set_queue_attributes,
             "SendMessage": self._send_message, "ReceiveMessage": self._receive_message,
             "DeleteMessage": self._delete_message, "PurgeQueue": self._purge_queue,
             "ChangeMessageVisibility": self._change_visibility, "TagQueue": self._tag_queue,
             "UntagQueue": self._untag_queue, "ListQueueTags": self._list_queue_tags}
        handler = h.get(action)
        if not handler:
            if self._is_json_protocol(req):
                return Response(json.dumps({"__type": "InvalidAction", "message": f"Action {action} not valid"}), 400, content_type="application/x-amz-json-1.0")
            return error_response("InvalidAction", f"The action {action} is not valid.", 400)
        return handler(req, path)

    def _p(self, req):
        params = dict(req.args)
        if req.content_type and "form" in req.content_type: params.update(req.form.to_dict())
        if self._is_json_protocol(req):
            try:
                jb = json.loads(req.get_data(as_text=True))
                if isinstance(jb, dict): params.update(jb)
            except Exception: pass
        elif not params.get("Action"):
            for k, v in parse_qs(req.get_data(as_text=True)).items(): params[k] = v[0] if len(v)==1 else v
        return params

    def _url(self, name):
        c = get_config(); return f"http://localhost:{c.port}/{c.account_id}/{name}"
    def _arn(self, name):
        c = get_config(); return f"arn:aws:sqs:{c.region}:{c.account_id}:{name}"
    def _find(self, path):
        for url, q in self.queues.items():
            if url.endswith(f"/{path}") or path.endswith(q.name): return q
        return None

    def _resp(self, req, action, xml_content="", json_data=None):
        if self._is_json_protocol(req):
            return Response(json.dumps(json_data or {}, default=str), 200, content_type="application/x-amz-json-1.0")
        body = f'<?xml version="1.0"?>\n<{action}Response xmlns="http://queue.amazonaws.com/doc/2012-11-05/">\n  <{action}Result>\n{xml_content}\n  </{action}Result>\n  <ResponseMetadata><RequestId>{new_request_id()}</RequestId></ResponseMetadata>\n</{action}Response>'
        return Response(body, 200, content_type="application/xml")

    def _err(self, req, code, msg, status=400):
        if self._is_json_protocol(req):
            return Response(json.dumps({"__type": code, "message": msg}), status, content_type="application/x-amz-json-1.0")
        return error_response(code, msg, status)

    def _resolve_queue(self, params, path):
        url = params.get("QueueUrl", "")
        if not url:
            q = self._find(path); 
            if q: url = q.url
        return self.queues.get(url)

    def _create_queue(self, req, path):
        p = self._p(req); name = p.get("QueueName", "")
        if not name: return self._err(req, "MissingParameter", "QueueName required")
        url = self._url(name)
        if url not in self.queues:
            attrs = {}
            if "Attributes" in p and isinstance(p["Attributes"], dict):
                attrs = p["Attributes"]
            else:
                i = 1
                while f"Attribute.{i}.Name" in p: attrs[p[f"Attribute.{i}.Name"]] = p.get(f"Attribute.{i}.Value", ""); i+=1
            self.queues[url] = SQSQueue(name=name, url=url, arn=self._arn(name), attributes=attrs)
            logger.info("Created queue: %s", name)
        return self._resp(req, "CreateQueue", f"    <QueueUrl>{url}</QueueUrl>", {"QueueUrl": url})

    def _delete_queue(self, req, path):
        p = self._p(req); url = p.get("QueueUrl", "")
        if not url:
            q = self._find(path)
            if q: url = q.url
        self.queues.pop(url, None)
        return self._resp(req, "DeleteQueue", "", {})

    def _list_queues(self, req, path):
        p = self._p(req); prefix = p.get("QueueNamePrefix", "")
        urls = [q.url for q in self.queues.values() if not prefix or q.name.startswith(prefix)]
        xml = "".join(f"    <QueueUrl>{u}</QueueUrl>\n" for u in urls)
        return self._resp(req, "ListQueues", xml, {"QueueUrls": urls})

    def _get_queue_url(self, req, path):
        p = self._p(req); url = self._url(p.get("QueueName", ""))
        if url not in self.queues: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        return self._resp(req, "GetQueueUrl", f"    <QueueUrl>{url}</QueueUrl>", {"QueueUrl": url})

    def _get_queue_attributes(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        now = time.time()
        q.attributes["ApproximateNumberOfMessages"] = str(sum(1 for m in q.messages if m.visible_after <= now))
        q.attributes["ApproximateNumberOfMessagesNotVisible"] = str(sum(1 for m in q.messages if m.visible_after > now))
        xml = "".join(f"    <Attribute><Name>{k}</Name><Value>{v}</Value></Attribute>\n" for k,v in q.attributes.items())
        return self._resp(req, "GetQueueAttributes", xml, {"Attributes": q.attributes})

    def _set_queue_attributes(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        if "Attributes" in p and isinstance(p["Attributes"], dict):
            q.attributes.update(p["Attributes"])
        else:
            i = 1
            while f"Attribute.{i}.Name" in p: q.attributes[p[f"Attribute.{i}.Name"]] = p.get(f"Attribute.{i}.Value", ""); i+=1
        return self._resp(req, "SetQueueAttributes", "", {})

    def _send_message(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        body = p.get("MessageBody", ""); delay = int(p.get("DelaySeconds", q.attributes.get("DelaySeconds", "0")))
        msg = SQSMessage(message_id=new_message_id(), body=body, visible_after=time.time()+delay)
        q.messages.append(msg); logger.info("Sent message to %s", q.name)
        j = {"MessageId": msg.message_id, "MD5OfMessageBody": msg.md5_of_body}
        return self._resp(req, "SendMessage", f"    <MessageId>{msg.message_id}</MessageId>\n    <MD5OfMessageBody>{msg.md5_of_body}</MD5OfMessageBody>", j)

    def _receive_message(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        max_msgs = int(p.get("MaxNumberOfMessages", "1"))
        vt = int(p.get("VisibilityTimeout", q.attributes.get("VisibilityTimeout", "30")))
        now = time.time(); xml = ""; msgs_json = []; count = 0
        for msg in q.messages:
            if count >= max_msgs: break
            if msg.visible_after > now: continue
            msg.visible_after = now + vt; msg.receive_count += 1; msg.receipt_handle = str(uuid.uuid4()); count += 1
            xml += f"    <Message><MessageId>{msg.message_id}</MessageId><ReceiptHandle>{msg.receipt_handle}</ReceiptHandle><MD5OfBody>{msg.md5_of_body}</MD5OfBody><Body>{msg.body}</Body></Message>\n"
            msgs_json.append({"MessageId": msg.message_id, "ReceiptHandle": msg.receipt_handle, "MD5OfBody": msg.md5_of_body, "Body": msg.body})
        return self._resp(req, "ReceiveMessage", xml, {"Messages": msgs_json} if msgs_json else {})

    def _delete_message(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        rh = p.get("ReceiptHandle", ""); q.messages = [m for m in q.messages if m.receipt_handle != rh]
        return self._resp(req, "DeleteMessage", "", {})

    def _purge_queue(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        q.messages.clear()
        return self._resp(req, "PurgeQueue", "", {})

    def _change_visibility(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        rh = p.get("ReceiptHandle", ""); vt = int(p.get("VisibilityTimeout", "30"))
        for m in q.messages:
            if m.receipt_handle == rh: m.visible_after = time.time() + vt; break
        return self._resp(req, "ChangeMessageVisibility", "", {})

    def _tag_queue(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        if "Tags" in p and isinstance(p["Tags"], dict):
            q.tags.update(p["Tags"])
        else:
            i = 1
            while f"Tag.{i}.Key" in p: q.tags[p[f"Tag.{i}.Key"]] = p.get(f"Tag.{i}.Value", ""); i+=1
        return self._resp(req, "TagQueue", "", {})

    def _untag_queue(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        keys = p.get("TagKeys", [])
        if isinstance(keys, list):
            for k in keys: q.tags.pop(k, None)
        else:
            i = 1
            while f"TagKey.{i}" in p: q.tags.pop(p[f"TagKey.{i}"], None); i+=1
        return self._resp(req, "UntagQueue", "", {})

    def _list_queue_tags(self, req, path):
        p = self._p(req); q = self._resolve_queue(p, path)
        if not q: return self._err(req, "AWS.SimpleQueueService.NonExistentQueue", "Queue not found")
        xml = "".join(f"    <Tag><Key>{k}</Key><Value>{v}</Value></Tag>\n" for k,v in q.tags.items())
        return self._resp(req, "ListQueueTags", xml, {"Tags": q.tags})
