import math, asyncio, subprocess
from telethon import TelegramClient
from fastapi.responses import StreamingResponse

import logging

logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class Download:
    client: TelegramClient
    route: str
    offset: int
    handler: None
    file: None
    limit: int
    file_size: float

    def __init__(self, handler):
        self.handler = handler
        self.file = handler.message.media
        self.file_size = handler.message.file.size
        self.limit = handler.sanity.limit
        self.offset = handler.sanity.offset
        self.client = handler.client
        self.mime_type = handler.message.file.mime_type

    async def download(self):
        part_size = int(512 * 1024) * 2
        first_part_cut = self.offset % part_size
        first_part = math.floor(self.offset / part_size)
        last_part_cut = part_size - (self.limit % part_size)
        last_part = math.ceil(self.limit / part_size)
        part_count = math.ceil(self.file_size / part_size)
        part = first_part
        try:
            async for chunk in self.client.iter_download(
                self.file, offset=first_part * part_size, request_size=part_size
            ):
                if part == first_part:
                    yield bytes(chunk[first_part_cut:])
                elif part == last_part:
                    yield bytes(chunk[:last_part_cut])
                else:
                    yield bytes(chunk)
                logging.debug(f"Part {part}/{last_part} (total {part_count}) served!")
                part += 1
            logging.debug("serving finished")
        except (GeneratorExit, StopAsyncIteration, asyncio.CancelledError):
            logging.debug("file serve interrupted")

            raise
        except Exception as e:
            print(e)
            logging.debug("file serve errored", exc_info=True)

    async def handle_request(self):
        headers = {
            "content-type": self.mime_type,
            "content-range": f"bytes {self.offset}-{self.limit-1}/{self.file_size}",
            "content-length": str(self.limit - self.offset),
            "accept-ranges": "bytes",
            "content-transfer-encoding": "Binary",
            "content-disposition": f'{self.handler.route}; filename="{self.handler.message.file.name}"',
        }
        logging.info(
            f"Serving file in {self.handler.message.file.name}) ; Range: {self.offset} - {self.limit}"
        )
        if self.handler.head:
            body = None
        else:
            body = self.download()
        return StreamingResponse(
            media_type=self.mime_type,
            content=body,
            headers=headers,
            status_code=206 if self.offset else 200,
        )