"""Browser-bound uploads and safe audio playback for the Gradio application.

Gradio's cache is not a private file store by itself. Register paths from real
upload responses, require the uploading browser on input, and permit only owned
recording downloads. Never expose primary/supporting documents or workspaces.
"""
from __future__ import annotations

from http.cookies import SimpleCookie
import hashlib
import json
import os
from pathlib import Path
import re
import secrets
import threading
import time
from urllib.parse import unquote

from starlette.responses import PlainTextResponse

COOKIE_NAME = 'det_browser'
AUDIO_EXTENSIONS = {'.wav', '.mp3', '.m4a', '.ogg', '.oga', '.flac', '.webm', '.opus', '.aac', '.mp4'}


def cookie_owner(headers) -> str:
    """Read a capability, not an identity; do not accept a session hash alone."""
    try:
        cookie = SimpleCookie()
        cookie.load(headers.get('cookie', ''))
        value = cookie[COOKIE_NAME].value if COOKIE_NAME in cookie else ''
        if not re.fullmatch(r'[0-9a-f]{64}', value):
            return ''
        # On local password installations, bind upload capabilities to the
        # current Gradio login too, so a second login in the same browser does
        # not inherit the previous user's cached uploads.
        if os.getenv('APP_EDITION') == 'local':
            auth = sorted((k, cookie[k].value) for k in cookie if k.startswith('access-token-'))
            if auth:
                return hashlib.sha256((value + repr(auth)).encode()).hexdigest()
        return value
    except (TypeError, AttributeError, ValueError):
        return ''


class UploadRegistry:
    def __init__(self, root: Path, ttl: float = 3600, max_entries: int = 10000):
        self.root = root.resolve()
        self.ttl = ttl
        self.max_entries = max_entries
        self._paths: dict[tuple[str, str], float] = {}
        self._lock = threading.Lock()

    def register(self, browser: str, paths: list[str]) -> None:
        now = time.monotonic()
        with self._lock:
            self._paths = {k: t for k, t in self._paths.items() if now - t <= self.ttl}
            for raw in paths:
                path = Path(raw).resolve()
                if path.is_relative_to(self.root):
                    self._paths[(browser, str(path))] = now
            if len(self._paths) > self.max_entries:
                # Bounded metadata even when a public demo is used heavily.
                ordered = sorted(self._paths, key=self._paths.get)
                for key in ordered[:len(self._paths) - self.max_entries]:
                    self._paths.pop(key, None)

    def owns(self, browser: str, path: Path) -> bool:
        if not browser or not path.resolve().is_relative_to(self.root):
            return False
        with self._lock:
            created = self._paths.get((browser, str(path.resolve())))
            return created is not None and time.monotonic() - created <= self.ttl

    def revoke(self, browser: str) -> None:
        with self._lock:
            self._paths = {key: t for key, t in self._paths.items() if key[0] != browser}


UPLOADS = UploadRegistry(Path(os.environ.get('GRADIO_TEMP_DIR', '/tmp/document-exam-trainer-uploads')))


class PrivateUploads:
    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope['type'] != 'http':
            return await self.app(scope, receive, send)
        headers = {key.decode('latin1').lower(): value.decode('latin1') for key, value in scope.get('headers', [])}
        browser = cookie_owner(headers)
        fresh = not browser
        if fresh:
            browser = secrets.token_hex(32)
        path = scope.get('path', '')
        method = scope.get('method', 'GET')

        if method in ('GET', 'HEAD'):
            match = re.search(r'/(?:gradio_api/)?file(?:=|/)(.*)$', path)
            if match:
                raw_path = unquote(match.group(1))
                candidate = Path(raw_path).resolve()
                permitted = (candidate.suffix.lower() in AUDIO_EXTENSIONS and
                             UPLOADS.owns(browser, candidate) and candidate.is_file())
                if not permitted:
                    return await PlainTextResponse('This file is not available to this browser.', status_code=403)(scope, receive, send)
            if '/gradio_api/stream/' in path:
                return await PlainTextResponse('Streaming file downloads are disabled.', status_code=403)(scope, receive, send)

        if path.rstrip('/') == '/logout':
            UPLOADS.revoke(browser)

        is_upload = method == 'POST' and path.rstrip('/').endswith('/gradio_api/upload')
        response_status = 0
        response_body = bytearray()
        oversized = False

        async def send_guarded(message):
            nonlocal response_status, oversized
            if message['type'] == 'http.response.start':
                response_status = message['status']
                if fresh:
                    # HF embeds the app on another site. CHIPS-capable browsers
                    # can use this partitioned cookie there; standalone is fallback.
                    secure = scope.get('scheme') == 'https' or headers.get('x-forwarded-proto', '').split(',')[0].strip() == 'https'
                    attributes = 'SameSite=None; Secure; Partitioned' if secure else 'SameSite=Lax'
                    value = f'{COOKIE_NAME}={browser}; Path=/; Max-Age=86400; HttpOnly; {attributes}'
                    message = {**message, 'headers': [*message.get('headers', []), (b'set-cookie', value.encode('ascii'))]}
            elif is_upload and response_status == 200 and message['type'] == 'http.response.body':
                if not oversized:
                    response_body.extend(message.get('body', b''))
                    if len(response_body) > 1024 * 1024:
                        oversized = True
                        response_body.clear()
                if not oversized and not message.get('more_body', False):
                    try:
                        paths = json.loads(response_body)
                        if isinstance(paths, list) and all(isinstance(p, str) for p in paths):
                            UPLOADS.register(browser, paths)
                    except (ValueError, TypeError):
                        pass  # Fail closed: unregistered paths cannot reach the engine.
            await send(message)

        await self.app(scope, receive, send_guarded)