from __future__ import annotations

from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Tuple


class ObsidianInjector:
    script_id = "aura-obsidian"
    max_body_bytes = 2 * 1024 * 1024

    def wrap_send(self, send: Callable[[Dict[str, Any]], Awaitable[None]], rules: Mapping[str, Any]):
        injector = self

        state: Dict[str, Any] = {
            "start": None,
            "headers": [],
            "body": bytearray(),
            "more": False,
        }

        async def send_wrapped(message: Dict[str, Any]):
            t = message.get("type")
            if t == "http.response.start":
                hdrs = list(message.get("headers") or [])
                # headers hardening
                if rules.get("obsidian_headers"):
                    hdrs = _set_header(hdrs, b"x-content-type-options", b"nosniff")
                    hdrs = _set_header(hdrs, b"x-frame-options", b"SAMEORIGIN")
                    hdrs = _set_header(hdrs, b"x-xss-protection", b"1; mode=block")

                state["start"] = {**message, "headers": hdrs}
                state["headers"] = hdrs
                return

            if t == "http.response.body":
                body = message.get("body") or b""
                more = bool(message.get("more_body"))

                # if no start yet, forward
                if state["start"] is None:
                    await send(message)
                    return

                # buffer (best-effort)
                if len(state["body"]) + len(body) <= injector.max_body_bytes:
                    state["body"].extend(body)
                else:
                    # too big, flush unmodified
                    await send(state["start"])
                    await send({"type": "http.response.body", "body": bytes(state["body"]) + body, "more_body": more})
                    state["start"] = None
                    state["body"] = bytearray()
                    return

                if more:
                    state["more"] = True
                    return

                # final chunk
                start = state["start"]
                state["start"] = None
                html = bytes(state["body"])
                state["body"] = bytearray()

                injected = injector.inject_into_html(html, rules, headers=start.get("headers") or [])
                await send({**start, "headers": _set_content_length(start.get("headers") or [], len(injected))})
                await send({"type": "http.response.body", "body": injected, "more_body": False})
                return

            await send(message)

        return send_wrapped

    def inject_into_html(self, body: bytes, rules: Mapping[str, Any], *, headers: List[Tuple[bytes, bytes]]):
        if not body:
            return body

        ct = _get_header(headers, b"content-type")
        is_html = (ct and b"text/html" in ct.lower()) or (b"<html" in body.lower())
        if not is_html:
            return body

        script = build_client_script(rules)
        if not script:
            return body

        injection = f"<script id=\"{self.script_id}\">{script}</script>".encode("utf-8")
        lower = body.lower()
        idx = lower.rfind(b"</body>")
        if idx != -1:
            return body[:idx] + injection + body[idx:]
        return body + injection


def build_client_script(rules: Mapping[str, Any]) -> str:
    mode = "aggressive" if rules.get("obsidian_mode") == "aggressive" else "compat"

    js = ""
    if rules.get("guard_rightclick"):
        js += "document.addEventListener('contextmenu',e=>e.preventDefault());"
    if rules.get("guard_devtools"):
        js += "document.onkeydown=function(e){if(e.keyCode==123||(e.ctrlKey&&e.shiftKey&&e.keyCode=='I'.charCodeAt(0))||(e.ctrlKey&&e.shiftKey&&e.keyCode=='J'.charCodeAt(0))||(e.ctrlKey&&e.keyCode=='U'.charCodeAt(0)))return false;};"
    if rules.get("guard_copy"):
        js += "document.addEventListener('copy',function(e){try{var t=(e.clipboardData||window.clipboardData);if(!t)return;var s=window.getSelection?String(window.getSelection()):'';if(!s)return;t.setData('text/plain',s+'\\n\\n[Protected by Aura]');e.preventDefault();}catch(_e){}});"

    if rules.get("obsidian_text") and mode == "aggressive":
        js += "(function(){function s(n){if(n.nodeType===3){n.nodeValue=n.nodeValue.split('').join('\\u200C');}else if(n.nodeType===1&&n.tagName!=='SCRIPT'&&n.tagName!=='STYLE'){for(var i=0;i<n.childNodes.length;i++){s(n.childNodes[i]);}}}}window.addEventListener('load',function(){try{s(document.body);}catch(_e){}},{once:true});})();"

    if rules.get("obsidian_css"):
        if mode == "aggressive":
            js += "(function(){function run(){try{var n=['ax-99','bz-22','c-al'];var all=document.querySelectorAll('*');for(var j=0;j<all.length;j++){all[j].classList.add(n[Math.floor(Math.random()*n.length)]);}}catch(_e){}}window.addEventListener('load',function(){setTimeout(run,50);},{once:true});})();"
        else:
            js += "(function(){function run(){try{var roots=['#app','#__next','[data-reactroot]'];for(var i=0;i<roots.length;i++){if(document.querySelector(roots[i]))return;}var n=['ax-99','bz-22','c-al'];var all=document.querySelectorAll('*');for(var j=0;j<all.length;j++){all[j].classList.add(n[Math.floor(Math.random()*n.length)]);}}catch(_e){}}if('requestIdleCallback' in window){requestIdleCallback(run,{timeout:1200});}else{window.addEventListener('load',function(){setTimeout(run,250);});}})();"

    if not js:
        return ""
    return f"(function(){{try{{{js}}}catch(_e){{}}}})();"


def _get_header(headers: List[Tuple[bytes, bytes]], key: bytes) -> Optional[bytes]:
    k = key.lower()
    for hk, hv in headers:
        if hk.lower() == k:
            return hv
    return None


def _set_header(headers: List[Tuple[bytes, bytes]], key: bytes, value: bytes) -> List[Tuple[bytes, bytes]]:
    out: List[Tuple[bytes, bytes]] = []
    k = key.lower()
    found = False
    for hk, hv in headers:
        if hk.lower() == k:
            if not found:
                out.append((key, value))
                found = True
            continue
        out.append((hk, hv))
    if not found:
        out.append((key, value))
    return out


def _set_content_length(headers: List[Tuple[bytes, bytes]], n: int) -> List[Tuple[bytes, bytes]]:
    # Remove any existing content-length; then set.
    out = [(k, v) for (k, v) in headers if k.lower() != b"content-length"]
    out.append((b"content-length", str(int(n)).encode("ascii")))
    return out
