from __future__ import annotations

from dataclasses import dataclass, field
from io import BytesIO

from quart.datastructures import FileStorage
from werkzeug.datastructures import Headers
from werkzeug.http import parse_options_header
from werkzeug.sansio.multipart import Data, Epilogue, Field, File, NeedData, MultipartDecoder

@dataclass
class MultipartUploadParts:
  fields: dict[str, str] = field(default_factory=dict)
  file: FileStorage | None = None

def _content_type(headers: Headers) -> str:
  return headers.get('content-type') or 'application/octet-stream'

def _fallback_filename(headers: Headers, default: str = 'file') -> str:
  disposition = headers.get('content-disposition', '')
  _, options = parse_options_header(disposition)
  filename = options.get('filename')
  return filename if filename is not None else default

def parse_multipart_upload_body(body: bytes, boundary: str | bytes | None, file_field_name: str = 'file') -> MultipartUploadParts:
  parts = MultipartUploadParts()
  if not body or not boundary:
    return parts

  boundary_bytes = boundary.encode() if isinstance(boundary, str) else boundary
  parser = MultipartDecoder(boundary_bytes)
  parser.receive_data(body)
  parser.receive_data(None)

  current_part = None
  current_chunks: list[bytes] = []

  while True:
    event = parser.next_event()
    if isinstance(event, (Epilogue, NeedData)):
      break
    if isinstance(event, (Field, File)):
      current_part = event
      current_chunks = []
      continue
    if not isinstance(event, Data) or current_part is None:
      continue

    current_chunks.append(event.data)
    if event.more_data:
      continue

    data = b''.join(current_chunks)
    if current_part.name == file_field_name and data:
      headers = current_part.headers
      parts.file = FileStorage(
        stream=BytesIO(data),
        filename=getattr(current_part, 'filename', None) or _fallback_filename(headers),
        name=file_field_name,
        content_type=_content_type(headers),
        headers=headers,
      )
    elif isinstance(current_part, Field):
      _, options = parse_options_header(current_part.headers.get('content-type', ''))
      charset = options.get('charset') or 'utf-8'
      parts.fields[current_part.name] = data.decode(charset, 'replace')

  return parts

def recover_file_from_multipart_body(body: bytes, boundary: str | bytes | None, field_name: str = 'file') -> FileStorage | None:
  return parse_multipart_upload_body(body, boundary, field_name).file