import asyncio
import math
import os
import tempfile

import modal
from modal.cls import ClsMixin
import tweepy


from suno_utils.worker.loader import S3Loader
from suno_utils.worker.modal_base import get_modal_base_image, MODAL_MOUNTS
from suno_utils.worker.chirp_worker import ChirpV0Worker
from suno_utils.worker.schema import QueueItem

aws_secret = modal.Secret.from_name("studio-aws")
twitter_secret = modal.Secret.from_name("studio-twitter")


def download_model_wrapper_4():
    ChirpV0Worker.download_models()


image = get_modal_base_image()
STUB_NAME = "twitter-botv2"
stub = modal.Stub(STUB_NAME, image=image)


class DummyLoader(S3Loader):
    pass


@stub.cls(
    cpu=1.0,
    secrets=[
        aws_secret,
        twitter_secret,
    ],
    mounts=MODAL_MOUNTS,
)
class TwitterStub(ClsMixin):
    def __enter__(self):
        self.worker = DummyLoader(bg_image="/suno/models/assets/wave-bg-2.png")

    @modal.method()
    async def upload_to_twitter(
        self,
        tmp_file,
        tweet_text="https://discord.com/invite/QRrnYufqZV",
    ):
        # Set up api (twitter api v1) and client (twitter api v2)
        consumer_key = os.environ["TWITTER_CONSUMER_KEY"]
        consumer_secret = os.environ["TWITTER_CONSUMER_SECRET"]
        access_token = os.environ["TWITTER_ACCESS_TOKEN"]
        access_token_secret = os.environ["TWITTER_ACCESS_TOKEN_SECRET"]
        auth = tweepy.OAuth1UserHandler(consumer_key, consumer_secret)
        auth.set_access_token(
            access_token,
            access_token_secret,
        )
        api = tweepy.API(auth)
        client = tweepy.Client(
            consumer_key=consumer_key,
            consumer_secret=consumer_secret,
            access_token=access_token,
            access_token_secret=access_token_secret,
        )

        # Get the file info
        file_info = os.stat(tmp_file.name)
        # Calculate real file size
        tmp_file.seek(0, 0)
        tmp_file.read(file_info.st_size + 50000)
        size = tmp_file.tell()
        print("file size = ", size)

        # Initiate the chunked upload
        media = api.chunked_upload_init(
            total_bytes=size,
            media_type="video/mp4",
            media_category="tweet_video",
        )
        print("initialized chunk upload: ", str(media))

        # Upload the chunks one by one
        chunksize = max(int(math.ceil(file_info.st_size / 1000)), 50000)
        print("chunksize = ", chunksize)
        tmp_file.seek(0, 0)
        num_chunks = int(math.ceil(file_info.st_size / chunksize))
        print("num_chunks =", num_chunks)
        for i in range(0, num_chunks):
            data_chunk = tmp_file.read(chunksize)
            upload_response = api.chunked_upload_append(
                media_id=media.media_id, media=data_chunk, segment_index=i
            )
        # Finalize the upload
        upload_response = api.chunked_upload_finalize(media.media_id)
        print("Finalized upload: ", str(upload_response))
        print(api.get_media_upload_status(media.media_id).processing_info["state"])

        # Poll the upload status until it's done (wait for up to 10 seconds)
        for i in range(20):
            await asyncio.sleep(1)
            state = api.get_media_upload_status(media.media_id).processing_info["state"]
            print(i, state)
            if state == "succeeded":
                break

        # The upload should now be complete and you can post a status with this media
        await asyncio.sleep(0.5)
        response = client.create_tweet(text=tweet_text, media_ids=[media.media_id])
        url = f"https://twitter.com/intent/tweet?text=Wanna%20hear%20me%20%23chirp?%20https://twitter.com/i/status/{response.data['id']}"
        # url = f"https://twitter.com/i/status/{response.data['id']}"
        return url

    @modal.method()
    async def upload_from_clip_id(
        self,
        clip_id,
        gen_message_id,
        batch_n,
        tweet_text="https://discord.com/invite/QRrnYufqZV",
        callback_url: str | None = None,
    ):
        from suno_utils.worker.settings import s3_client

        with tempfile.NamedTemporaryFile(mode="r+b", suffix=".mp4") as tmp_file:
            s3_client.download_fileobj(
                "suno-data-uploads",
                f"studio/uploads/{clip_id}.mp4",
                tmp_file,
            )
            print("Start upload for clip id: ", clip_id)
            url = await self.upload_to_twitter(tmp_file, tweet_text=tweet_text)
            print(url)
            self.worker.notify_finish(
                QueueItem(id=clip_id, metadata={}, callback_url=callback_url),
                {
                    "id": clip_id,
                    "model": "chirp_v0",
                    "type": "make_tweet",
                    "tweet": url,
                    "gen_message_id": gen_message_id,
                    "batch_n": batch_n,
                },
                queue_name="results:q",
            )


@stub.local_entrypoint()
async def main():
    model = TwitterStub()
    gen_request_id = "fc22f413-34de-4867-845c-049dbf53f266"
    clip_id = "3d478f8a-dbff-4ebd-8fd2-237ea5b6de33_0"
    # print("clip_id", clip_id)
    model.upload_from_clip_id.call(
        clip_id,
        gen_request_id,
        1,
        "https://discord.com/invite/QRrnYufqZV",
        "https://troll-selected-presumably.ngrok-free.app/computed/",
    )