from functools import partial
import json
import os
import uuid
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field
from urllib.request import Request, urlopen
import urllib.error

from fastapi import Request as FastAPIRequest
from fal.toolkit.file.providers.fal import LIFECYCLE_PREFERENCE, FalV3TokenManager
from fal.toolkit.file import request_lifecycle_preference
from fal.toolkit.utils.retry import retry

from serverless_comfy.utils import (
    debug_print,
    get_random_filename_for_mime_type,
    async_retry,
)

import asyncio
import datetime
import io
import math

FAL_CDN = "https://v3.fal.media"
UPLOAD_TIMEOUT = 90

# Multipart upload constants
MULTIPART_THRESHOLD = 500 * 1024 * 1024  # 500 MB — avoid multipart 409 bug
MULTIPART_CHUNK_SIZE = 10 * 1024 * 1024  # 10 MB
MULTIPART_MAX_CONCURRENCY = 10

fal_v3_token_manager = FalV3TokenManager()


@retry(max_delay=4)
def get_auth_headers() -> dict[str, str]:
    token = fal_v3_token_manager.get_token()

    return {
        "Authorization": f"{token.token_type} {token.token}",
        "User-Agent": "fal/0.1.0",
    }


def get_async_http_client():
    import httpx

    return httpx.AsyncClient()


def get_lifecycle(
    object_lifecycle_preference: dict[str, str] | None = None,
) -> dict[str, str]:
    olp = object_lifecycle_preference or LIFECYCLE_PREFERENCE.get()
    return olp


@dataclass
class CDNManager:
    thread_pool: ThreadPoolExecutor = field(
        default_factory=partial(ThreadPoolExecutor, max_workers=36)
    )
    uploads: list[Future] = field(default_factory=list)
    upload_headers: dict[str, dict[str, str] | None] = field(default_factory=dict)

    def prepare_access_url(
        self,
        content_type: str,
        object_lifecycle_preference: dict[str, str] | None = None,
        file_name: str | None = None,
    ) -> Future[tuple[str, str]]:
        lifecycle = get_lifecycle(object_lifecycle_preference)

        @retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
        def _get_access_url(request: Request) -> str | None:
            with urlopen(request) as response:
                return json.load(response)["access_url"]

        def _prepare_access_url() -> tuple[str, str]:
            id = str(uuid.uuid4())

            token = fal_v3_token_manager.get_token()
            auth_headers = {
                "Authorization": f"{token.token_type} {token.token}",
                "User-Agent": "fal/serverless-comfy-cdn",
            }

            headers = {
                "Content-Type": content_type,
                "Accept": "application/json",
                "X-Fal-Object-Lifecycle": json.dumps(lifecycle),
                **auth_headers,
            }

            if file_name:
                headers["X-Fal-File-Name"] = file_name

            request = Request(
                FAL_CDN + "/files/upload?async=1",
                method="POST",
                headers=headers,
            )

            self.upload_headers[id] = headers
            return id, _get_access_url(request)

        return self.thread_pool.submit(_prepare_access_url)

    def upload_files(
        self,
        content_type: str,
        file_paths: list[str],
        wait_for_upload: bool = True,
    ) -> list[str]:
        access_url_futures = [
            self.prepare_access_url(
                content_type=content_type, file_name=os.path.basename(file_path)
            )
            for file_path in file_paths
        ]
        access_urls = [
            access_url_future.result() for access_url_future in access_url_futures
        ]

        upload_futures: list[Future[None]] = []
        for (id, access_url), file_path in zip(access_urls, file_paths):
            headers = self.upload_headers.pop(id)
            upload_future = self.upload_file(access_url, file_path, headers=headers)
            upload_futures.append(upload_future)

        if wait_for_upload:
            for future in upload_futures:
                future.result()

        return [access_url for _, access_url in access_urls]

    def upload_file(
        self,
        access_url: str,
        file_path: str,
        headers: dict[str, str] | None = None,
    ) -> Future[None]:
        headers = headers or {}

        with open(file_path, "rb") as fp:
            file_data = fp.read()

        @retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
        def _upload_data() -> None:
            request = Request(
                access_url,
                method="PUT",
                headers=headers,
                data=file_data,
            )

            try:
                with urlopen(request, timeout=UPLOAD_TIMEOUT):
                    pass
            except urllib.error.HTTPError as e:
                print(e.read().decode("utf-8"))
                raise e

        return self.thread_pool.submit(_upload_data)


@dataclass
class AsyncCDNManager:
    thread_pool: ThreadPoolExecutor = field(
        default_factory=lambda: ThreadPoolExecutor(max_workers=10)
    )
    auth_headers: dict[str, str] = field(default_factory=get_auth_headers)
    upload_headers: dict[str, dict[str, str] | None] = field(default_factory=dict)
    async_client: object = field(default_factory=get_async_http_client)
    _client_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    _auth_headers_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    _upload_headers_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    header_last_updated: datetime.datetime = field(
        default_factory=datetime.datetime.now
    )

    async def _ensure_valid_auth_headers(self) -> None:
        """
        Ensure that the authentication headers are valid.
        """
        if datetime.datetime.now() - self.header_last_updated > datetime.timedelta(
            days=3
        ):
            async with self._auth_headers_lock:
                if (
                    datetime.datetime.now() - self.header_last_updated
                    > datetime.timedelta(days=3)
                ):
                    self.auth_headers = get_auth_headers()
                    self.header_last_updated = datetime.datetime.now()

    async def _reset_client(self) -> None:
        """
        Reset the client by closing the old one and creating a new one.
        """
        import httpx

        async with self._client_lock:
            old_client = self.async_client
            self.async_client: httpx.AsyncClient = get_async_http_client()
            try:
                await old_client.aclose()
            except Exception:
                # Best effort cleanup - don't fail if closing the old client fails
                pass

    @async_retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
    async def _upload_part(
        self,
        url: str,
        data: bytes,
        content_type: str,
    ) -> str:
        """Upload a single part with retry logic. Returns the ETag."""
        headers = {
            **self._get_auth_headers(),
            "Content-Type": content_type,
        }
        try:
            response = await self.async_client.request(
                "PUT",
                url,
                headers=headers,
                content=data,
                timeout=UPLOAD_TIMEOUT,
            )
        except Exception:
            await self._reset_client()
            raise
        response.raise_for_status()
        return response.headers.get("ETag", "").strip('"')

    @async_retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
    async def _complete_multipart(
        self,
        complete_url: str,
        parts: list[dict],
    ) -> None:
        """Complete the multipart upload with retry logic."""
        headers = {
            **self._get_auth_headers(),
            "Accept": "application/json",
            "Content-Type": "application/json",
        }
        try:
            response = await self.async_client.request(
                "POST",
                complete_url,
                headers=headers,
                json={"parts": parts},
                timeout=UPLOAD_TIMEOUT,
            )
        except Exception:
            await self._reset_client()
            raise
        response.raise_for_status()

    async def upload_blob_multipart(
        self,
        access_url: str,
        upload_id: str,
        blob: io.BytesIO,
        content_type: str = "application/octet-stream",
    ) -> None:
        """
        Upload a blob using multipart upload (1 or more parts based on size).
        The access_url and upload_id come from prepare_access_url.
        """
        data = blob.getvalue()
        blob_size = len(data)
        debug_print(f"Uploading {blob_size} bytes to {access_url} via multipart")

        if blob_size > MULTIPART_THRESHOLD:
            # Large file: upload in multiple parts concurrently
            num_parts = math.ceil(blob_size / MULTIPART_CHUNK_SIZE)
            semaphore = asyncio.Semaphore(MULTIPART_MAX_CONCURRENCY)
            parts: list[dict] = []
            parts_lock = asyncio.Lock()

            async def upload_part(part_number: int, part_data: bytes) -> None:
                async with semaphore:
                    url = f"{access_url}/multipart/{upload_id}/{part_number}"
                    etag = await self._upload_part(url, part_data, content_type)
                    async with parts_lock:
                        parts.append({"partNumber": part_number, "etag": etag})

            tasks = []
            for part_number in range(1, num_parts + 1):
                start = (part_number - 1) * MULTIPART_CHUNK_SIZE
                end = start + MULTIPART_CHUNK_SIZE
                part_data = data[start:end]
                tasks.append(upload_part(part_number, part_data))

            await asyncio.gather(*tasks)

            # Sort parts and complete
            sorted_parts = sorted(parts, key=lambda p: p["partNumber"])
        else:
            # Small file: upload as single part
            url = f"{access_url}/multipart/{upload_id}/1"
            etag = await self._upload_part(url, data, content_type)
            sorted_parts = [{"partNumber": 1, "etag": etag}]

        # Complete the multipart upload
        complete_url = f"{access_url}/multipart/{upload_id}/complete"
        await self._complete_multipart(complete_url, sorted_parts)
        debug_print(f"Multipart upload completed: {access_url}")

    def _get_auth_headers(self) -> dict[str, str]:
        """Get authentication headers for multipart uploads."""
        token = fal_v3_token_manager.get_token()
        return {
            "Authorization": f"{token.token_type} {token.token}",
            "User-Agent": "fal/0.1.0",
        }

    @async_retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
    async def _initiate_multipart(
        self,
        content_type: str,
        file_name: str,
        object_lifecycle_preference: dict[str, str] | None = None,
    ) -> tuple[str, str]:
        """Initiate a multipart upload with retry logic. Returns (access_url, upload_id)."""
        token = fal_v3_token_manager.get_token()
        url = f"{token.base_upload_url}/files/upload/multipart"

        # Always get lifecycle with proper fallback to LIFECYCLE_PREFERENCE.get()
        lifecycle = get_lifecycle(object_lifecycle_preference)

        headers = {
            **self._get_auth_headers(),
            "Accept": "application/json",
            "Content-Type": content_type,
            "X-Fal-File-Name": file_name,
            # Used by V3 CDN
            "X-Fal-Object-Lifecycle": json.dumps(lifecycle),
            # Used by V1 CDN
            "X-Fal-Object-Lifecycle-Preference": json.dumps(lifecycle),
        }

        try:
            response = await self.async_client.request(
                "POST",
                url,
                headers=headers,
                timeout=UPLOAD_TIMEOUT,
            )
        except Exception:
            await self._reset_client()
            raise
        response.raise_for_status()

        result = response.json()
        return result["access_url"], result["uploadId"]

    @async_retry(max_retries=5, base_delay=1, backoff_type="exponential", jitter=True)
    async def _prepare_direct_url(
        self,
        content_type: str,
        file_name: str,
        object_lifecycle_preference: dict[str, str] | None = None,
    ) -> str:
        """Get a direct (non-multipart) upload URL."""
        token = fal_v3_token_manager.get_token()
        lifecycle = get_lifecycle(object_lifecycle_preference)
        headers = {
            **self._get_auth_headers(),
            "Accept": "application/json",
            "Content-Type": content_type,
            "X-Fal-File-Name": file_name,
            "X-Fal-Object-Lifecycle": json.dumps(lifecycle),
        }
        try:
            response = await self.async_client.request(
                "POST",
                FAL_CDN + "/files/upload?async=1",
                headers=headers,
                timeout=UPLOAD_TIMEOUT,
            )
        except Exception:
            await self._reset_client()
            raise
        response.raise_for_status()
        return response.json()["access_url"]

    async def prepare_access_url(
        self,
        content_type: str = "image/jpeg",
        object_lifecycle_preference: dict[str, str] | None = None,
        file_name: str | None = None,
        request: FastAPIRequest | None = None,
    ) -> tuple[str, str | None]:
        """
        Prepare an upload URL. Uses direct PUT for small files, multipart for large.
        Returns (id, access_url) where id can be used to retrieve upload metadata.
        """
        if object_lifecycle_preference is None and request is not None:
            object_lifecycle_preference = request_lifecycle_preference(request)

        id = str(uuid.uuid4())
        _file_name = file_name or get_random_filename_for_mime_type(content_type)

        await self._ensure_valid_auth_headers()

        # Try direct upload first, fall back to multipart for large files
        access_url = await self._prepare_direct_url(
            content_type, _file_name, object_lifecycle_preference
        )

        # Store upload metadata — upload_id=None signals direct PUT
        self.upload_headers[id] = {
            "Content-Type": content_type,
            "upload_id": None,
            "access_url": access_url,
        }

        return id, access_url

    async def prepare_sync_access_url(self) -> tuple[str, str | None]:
        """
        Create a data URI for a given file.
        """
        return (str(uuid.uuid4()), None)

    def prepare_urls(
        self,
        count: int,
        content_type: str = "image/jpeg",
        sync_mode: bool = False,
        request: FastAPIRequest | None = None,
        create_file_names: bool = False,
        object_lifecycle_preference: dict[str, str] | None = None,
    ) -> list[asyncio.Task[tuple[str, str | None]]]:
        """
        Prepare access URLs for a given number of images.
        """
        if sync_mode:
            coros = [self.prepare_sync_access_url() for _ in range(count)]
        else:
            file_names: list[str | None] = [None] * count
            if create_file_names:
                file_names = [
                    get_random_filename_for_mime_type(content_type)
                    for _ in range(count)
                ]
            coros = [
                self.prepare_access_url(
                    content_type=content_type,
                    request=request,
                    file_name=file_names[i],
                    object_lifecycle_preference=object_lifecycle_preference,
                )
                for i in range(count)
            ]

        return [asyncio.create_task(coro) for coro in coros]

    async def upload_files(
        self,
        content_type: str,
        file_paths: list[str],
        wait_for_upload: bool = True,
    ) -> list[str]:
        """
        Upload files given their file paths.
        """

        urls = self.prepare_urls(count=len(file_paths), content_type=content_type)
        url_results = await asyncio.gather(*urls)

        assert len(url_results) == len(file_paths)
        assert all(url is not None for _, url in url_results)

        def _read_file(file_path: str) -> io.BytesIO:
            blob = io.BytesIO()
            with open(file_path, "rb") as fp:
                blob.write(fp.read())
            return blob

        def _upload_done_callback(task: asyncio.Task, access_url: str) -> None:
            exc = task.exception()
            if exc is not None:
                debug_print(f"Failed to upload image to {access_url}: {exc}")
            else:
                debug_print(f"Uploaded image to {access_url}")

        async def _direct_put(access_url: str, data: bytes, content_type: str,
                              file_name: str = "output"):
            """Upload via a simple PUT. On 409/400, get a fresh URL and retry.
            Returns the final access_url (may differ from input if retried)."""
            import httpx
            for attempt in range(5):
                headers = {
                    **self._get_auth_headers(),
                    "Content-Type": content_type,
                }
                try:
                    response = await self.async_client.request(
                        "PUT",
                        access_url,
                        headers=headers,
                        content=data,
                        timeout=UPLOAD_TIMEOUT,
                    )
                    response.raise_for_status()
                    return access_url
                except httpx.HTTPStatusError as exc:
                    if exc.response.status_code in (409, 400) and attempt < 4:
                        debug_print(f"Upload got {exc.response.status_code}, getting fresh URL (attempt {attempt+1})")
                        try:
                            access_url = await self._prepare_direct_url(
                                content_type, file_name,
                            )
                        except Exception:
                            pass
                        await asyncio.sleep(min(2 ** attempt * 0.5, 30))
                        continue
                    raise
                except Exception as exc:
                    if attempt == 4:
                        raise
                    await asyncio.sleep(min(2 ** attempt * 0.5, 30))
            return access_url

        async def upload_file(
            file_path: str,
            id: str,
            access_url: str,
            content_type: str,
        ):
            event_loop = asyncio.get_event_loop()
            blob = await event_loop.run_in_executor(
                self.thread_pool, _read_file, file_path
            )

            async with self._upload_headers_lock:
                metadata = self.upload_headers.pop(id, {})

            upload_id = metadata.get("upload_id")  # type: ignore
            content_type = metadata.get("Content-Type", content_type)  # type: ignore
            data = blob.getvalue()

            if upload_id is None:
                # Direct PUT upload (non-multipart) — more reliable
                if len(data) > MULTIPART_THRESHOLD:
                    # File too large for direct PUT, fall back to multipart
                    _access_url, _upload_id = await self._initiate_multipart(
                        content_type,
                        os.path.basename(file_path) if hasattr(file_path, '__len__') else "output",
                    )
                    if wait_for_upload:
                        await self.upload_blob_multipart(
                            access_url=_access_url,
                            upload_id=_upload_id,
                            blob=blob,
                            content_type=content_type,
                        )
                    else:
                        t = asyncio.create_task(
                            self.upload_blob_multipart(
                                access_url=_access_url, upload_id=_upload_id,
                                blob=blob, content_type=content_type,
                            )
                        )
                        t.add_done_callback(partial(_upload_done_callback, access_url=access_url))
                else:
                    _fname = os.path.basename(file_path) if isinstance(file_path, str) else "output"
                    if wait_for_upload:
                        access_url = await _direct_put(access_url, data, content_type, file_name=_fname)
                    else:
                        t = asyncio.create_task(_direct_put(access_url, data, content_type, file_name=_fname))
                        t.add_done_callback(partial(_upload_done_callback, access_url=access_url))
            else:
                # Legacy multipart path
                if wait_for_upload:
                    await self.upload_blob_multipart(
                        access_url=access_url,
                        upload_id=upload_id,
                        blob=blob,
                        content_type=content_type,
                    )
                else:
                    upload_task = asyncio.create_task(
                        self.upload_blob_multipart(
                            access_url=access_url,
                            upload_id=upload_id,
                            blob=blob,
                            content_type=content_type,
                        )
                    )
                    upload_task.add_done_callback(
                        partial(
                            _upload_done_callback,
                            access_url=access_url,
                        )
                    )
            return access_url

        upload_tasks: list[asyncio.Task[str]] = []
        if wait_for_upload:
            async with asyncio.TaskGroup() as tg:
                for file_idx, (id, url) in enumerate(url_results):
                    file_path = file_paths[file_idx]

                    upload_task = tg.create_task(
                        upload_file(file_path, id, url, content_type)  # type: ignore
                    )
                    upload_tasks.append(upload_task)

            return [task.result() for task in upload_tasks]
        else:
            for file_idx, (id, url) in enumerate(url_results):
                file_path = file_paths[file_idx]
                upload_task = asyncio.create_task(
                    upload_file(file_path, id, url, content_type)
                )  # type: ignore
                upload_tasks.append(upload_task)

            return [url for _, url in url_results]
