import aiohttp
import asyncio
import json
import logging
import re
from urllib import parse

from pytubezero.exceptions import RegexMatchError, MaxRetriesExceeded
from pytubezero.helpers import regex_search

logger = logging.getLogger(__name__)
default_range_size = 9 * 1024 * 1024  # 9MB

class AsyncHTTPClient:
    """Singleton Async HTTP Client with persistent session and handy methods."""

    _instance = None
    _session = None

    def __new__(cls, *args, **kwargs):
        # Singleton pattern: one instance per process
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    async def _get_session(self):
        """Ensure session is alive (recreate if closed)."""
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession()
        return self._session

    async def close(self):
        """Close the internal session (call at end of program or via async with)."""
        if self._session and not self._session.closed:
            await self._session.close()

    async def __aenter__(self):
        await self._get_session()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

    async def _execute_request(
        self, url, method="GET", headers=None, data=None, timeout=None
    ):
        base_headers = {"User-Agent": "Mozilla/5.0", "accept-language": "en-US,en"}
        if headers:
            base_headers.update(headers)
        send_data = None
        if data is not None:
            send_data = (
                json.dumps(data).encode('utf-8') if not isinstance(data, (bytes, str)) else data
            )
            base_headers.setdefault("Content-Type", "application/json")
        if not url.lower().startswith("http"):
            raise ValueError("Invalid URL")
        session = await self._get_session()
        try:
            resp = await session.request(
                method=method,
                url=url,
                headers=base_headers,
                data=send_data,
                timeout=timeout,
            )
            return resp
        except Exception as e:
            logger.error(f"HTTP error: {e}")
            raise

    async def get(self, url, headers=None, timeout=None):
        """GET request, returns response text."""
        resp = await self._execute_request(url, method="GET", headers=headers, timeout=timeout)
        async with resp:
            return await resp.text()

    async def post(self, url, headers=None, data=None, timeout=None):
        """POST request, returns response text."""
        headers = headers or {}
        headers.setdefault("Content-Type", "application/json")
        resp = await self._execute_request(
            url, method="POST", headers=headers, data=data, timeout=timeout
        )
        async with resp:
            return await resp.text()

    async def head(self, url, headers=None, timeout=None):
        """HEAD request, returns headers as dict."""
        resp = await self._execute_request(url, method="HEAD", headers=headers, timeout=timeout)
        async with resp:
            return {k.lower(): v for k, v in resp.headers.items()}

    async def stream(self, url, timeout=None, max_retries=0):
        """Async generator: stream file in chunks with retries and range support."""
        file_size = default_range_size
        downloaded = 0

        while downloaded < file_size:
            stop_pos = min(downloaded + default_range_size, file_size) - 1
            tries = 0
            response = None
            # retry loop
            while True:
                if tries >= 1 + max_retries:
                    raise MaxRetriesExceeded()
                try:
                    response = await self._execute_request(
                        f"{url}&range={downloaded}-{stop_pos}", timeout=timeout
                    )
                    break
                except aiohttp.ClientError as e:
                    logger.error(e)
                    tries += 1
            # get real filesize if first chunk
            if file_size == default_range_size:
                try:
                    resp = await self._execute_request(
                        f"{url}&range=0-99999999999", timeout=timeout
                    )
                    async with resp:
                        content_length = resp.headers.get("Content-Length")
                        if content_length:
                            file_size = int(content_length)
                except Exception as e:
                    logger.error(e)
            async with response:
                while True:
                    chunk = await response.content.readany()
                    if not chunk:
                        break
                    downloaded += len(chunk)
                    yield chunk

    async def seq_stream(self, url, timeout=None, max_retries=0):
        """Async generator: read sequential video segments in order."""
        split_url = parse.urlsplit(url)
        base_url = f"{split_url.scheme}://{split_url.netloc}/{split_url.path}?"
        qs = dict(parse.parse_qsl(split_url.query))
        qs["sq"] = 0
        url_0 = base_url + parse.urlencode(qs)
        buffer = bytearray()
        async for chunk in self.stream(url_0, timeout, max_retries):
            yield chunk
            buffer.extend(chunk)
        # Find segment count
        segment_regex = re.compile(b"Segment-Count: (\\d+)")
        m = segment_regex.search(buffer)
        if not m:
            raise RegexMatchError("seq_stream", segment_regex.pattern)
        segment_count = 0
        for line in buffer.split(b"\r\n"):
            m = re.search(b"Segment-Count: (\\d+)", line)
            if m:
                segment_count = int(m.group(1))
        # Download all segments
        for sq in range(1, segment_count + 1):
            qs["sq"] = sq
            seg_url = base_url + parse.urlencode(qs)
            async for chunk in self.stream(seg_url, timeout, max_retries):
                yield chunk

    async def filesize(self, url):
        """Get file size via HEAD (Content-Length)."""
        head_info = await self.head(url)
        return int(head_info["content-length"])

    async def seq_filesize(self, url):
        """Get total file size for sequential segments."""
        total = 0
        split_url = parse.urlsplit(url)
        base_url = f"{split_url.scheme}://{split_url.netloc}/{split_url.path}?"
        qs = dict(parse.parse_qsl(split_url.query))
        qs["sq"] = 0
        url_0 = base_url + parse.urlencode(qs)
        resp = await self._execute_request(url_0, method="GET")
        async with resp:
            data = await resp.read()
        total += len(data)
        # Find segment count
        segment_count = 0
        for line in data.split(b"\r\n"):
            try:
                segment_count = int(regex_search(b"Segment-Count: (\\d+)", line, 1))
            except RegexMatchError:
                pass
        if segment_count == 0:
            raise RegexMatchError("seq_filesize", b"Segment-Count: (\\d+)")
        for sq in range(1, segment_count + 1):
            qs["sq"] = sq
            seg_url = base_url + parse.urlencode(qs)
            head_info = await self.head(seg_url)
            total += int(head_info["content-length"])
        return total
