import asyncio
import aiohttp
from os import environ, getenv
import ujson
from fastapi.responses import PlainTextResponse, StreamingResponse
from requests import post, Session, adapters
from loguru import logger
from SpectrumCommon.ServerUtils.ExceptionUtils import SpectrumException
from SpectrumCommon.Const.GB import GB_SW_UNKNOWN
import time

def getProcessTime(start):
    return time.strftime("%M:%S", time.gmtime(time.time()-start))

def log_error(url, data, response):
    logger.error(f'POST {url} status={response.status_code}')
    logger.error(f'POST {url} payload={data}')
    logger.error(f'POST {url} body={response.content}')
    try:
        json_response = ujson.loads(response.content)
        error_code = json_response.get('SWErrorCode', GB_SW_UNKNOWN)
        return json_response 
    except:
        error_code = GB_SW_UNKNOWN
        message = f'POST {url} status={response.status_code} payload={data} body={response.content}'
        raise SpectrumException(error_code,
                                status_code = response.status_code, 
                                message=f'POST {url} status={response.status_code} payload={data} body={response.content}')

# import io
def make_post(url, data=None, isJSON=True, headers={}, files = None, stream=False, backoff_factor=0.3):
    DEBUG_MODE = False

    if DEBUG_MODE == True:
        start = time.time()
        print(time.strftime("%I:%M %p") + ': _REQ ' + url)
    else:
        start = 0

    for key in ["AVENIR_SW_API_KEY"]:
        if key in environ.keys():
            headers[key] = getenv(key, "")

    session = Session()
    retries = adapters.Retry(total=3, backoff_factor=backoff_factor, status_forcelist=[ 500, 502, 503, 504 ])
    session.mount('http://', adapters.HTTPAdapter(max_retries=retries))
    session.mount('https://', adapters.HTTPAdapter(max_retries=retries))

    if isJSON:
        headers["Content-Type"] = "application/json"
        if stream:
            response = session.post(url, data=ujson.dumps(data), headers=headers, stream=stream)
            if response.status_code not in [200, 201]:
                log_error(url, data, response)
            # Pass the error straight back if POST failed for any reason
            if stream and response.ok:
                headers = {"Content-Disposition": response.headers["Content-Disposition"],
                        "Content-Type": response.headers["Content-Type"]}
                
                def iter_chunks(r):
                    for chunk in r.iter_content(1024 * 1024):
                        yield chunk

                return StreamingResponse(iter_chunks(response), headers=headers,  media_type='.zip')
        else:
            response = session.post(url, data=ujson.dumps(data), headers=headers)
            if response.status_code not in [200, 201]:
                log_error(url, data, response)
    else:
        response = session.post(url, data=data, files=files, headers=headers, stream=stream)
        if response.status_code not in [200, 201]:
            log_error(url, data, response)
        # Pass the error straight back if POST failed for any reason
        if stream and response.ok:
            headers = {"Content-Disposition": response.headers["Content-Disposition"],
                       "Content-Type": response.headers["Content-Type"]}
            
            def iter_chunks(r):
                for chunk in r.iter_content(1024 * 1024):
                    yield chunk

            return StreamingResponse(iter_chunks(response), headers=headers,  media_type='.zip')

    if DEBUG_MODE == True:
        print(time.strftime("%I:%M %p") + ': _RES ' + url + " time: " + getProcessTime(start))

    response = PlainTextResponse(response.text, status_code=response.status_code)
    return response

def make_get(url, isJSON=True, headers={}, files = None, stream=False):
    for key in ["AVENIR_SW_API_KEY"]:
        if key in environ.keys():
            headers[key] = getenv(key, "")

    session = Session()
    retries = adapters.Retry(total=3, backoff_factor=0.1, status_forcelist=[ 500, 502, 503, 504 ])
    session.mount('http://', adapters.HTTPAdapter(max_retries=retries))
    session.mount('https://', adapters.HTTPAdapter(max_retries=retries))

    if isJSON:
        headers["Content-Type"] = "application/json"
        if stream:
            response = session.get(url, headers=headers, stream=stream)
            # Pass the error straight back if GET failed for any reason
            if stream and response.ok:
                headers = {"Content-Disposition": response.headers["Content-Disposition"],
                        "Content-Type": response.headers["Content-Type"]}
                
                def iter_chunks(r):
                    for chunk in r.iter_content(1024 * 1024):
                        yield chunk

                return StreamingResponse(iter_chunks(response), headers=headers,  media_type='.zip')
        else:
            response = session.get(url, headers=headers)
            if response.status_code not in [200, 201]:
                logger.error(f'GET {url} status={response.status_code}')
    else:
        response = session.get(url, files=files, headers=headers, stream=stream)
        if response.status_code not in [200, 201]:
            logger.error(f'GET {url} status={response.status_code}')
        # Pass the error straight back if GET failed for any reason
        if stream and response.ok:
            headers = {"Content-Disposition": response.headers["Content-Disposition"],
                       "Content-Type": response.headers["Content-Type"]}
            
            def iter_chunks(r):
                for chunk in r.iter_content(1024 * 1024):
                    yield chunk

            return StreamingResponse(iter_chunks(response), headers=headers,  media_type='.zip')
    return PlainTextResponse(response.text, status_code=response.status_code)

#currently in development
def make_streaming_post(url, data=None, isJSON=True, headers={}, files = None):
    for key in ["AVENIR_SW_API_KEY"]:
        if key in environ.keys():
            headers[key] = getenv(key, "")

    if isJSON:
        headers["Content-Type"] = "application/json"
        response = post(url, data=ujson.dumps(data), headers=headers)
    else:
        # headers["Transfer-Encoding"] = 'chunked'
        response = post(url, data=data, files = files, headers=headers)#, stream=True)
    if response.status_code not in [200, 201]:
        logger.error(f'streaming POST {url} status={response.status_code}')
        logger.error(f'streaming POST {url} payload={data}')
        # Pass the error straight back if POST failed for any reason
    
    fileName = 'test.xlsx'
    # file_bytes = BytesIO() #response.content
    # for chunk in response.iter_content(chunk_size=10 * 1024):
    #     file_bytes.write(chunk)
    file_bytes = response.content
        
    if '.csv' in fileName:
        stream_response = StreamingResponse(
            iter_bytes(file_bytes),
            headers={"Content-Disposition": f"attachment; filename={fileName}",
                     "Content-Type":"text/csv; charset=utf-8"},
            media_type='.csv',
        )
    else:
        stream_response = StreamingResponse(
            file_bytes, #iter_bytes(file_bytes),
            media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            headers={"Content-Disposition": "attachment; filename=test.xlsx", "Content-Length":str(len(file_bytes))}
        )  
    return stream_response#StreamingResponse(response.text, status_code=response.status_code)

def iter_bytes(file_bytes):
    yield file_bytes


async def async_make_post(url, data=None, isJSON=True, headers={}, files=None, stream=False, backoff_factor=0.3):
    for key in ["AVENIR_SW_API_KEY"]:
        if key in environ.keys():
            headers[key] = getenv(key, "")

    # aiohttp does not support built-in retries; you can implement manual retries if needed
    attempt = 0
    max_attempts = 3
    while attempt < max_attempts:
        try:
            async with aiohttp.ClientSession() as session:
                if isJSON:
                    headers["Content-Type"] = "application/json"
                    if stream:
                        async with session.post(url, data=ujson.dumps(data), headers=headers) as response:
                            if response.status not in [200, 201]:
                                logger.error(f'POST {url} status={response.status}')
                                logger.error(f'POST {url} payload={data}')
                                logger.error(f'POST {url} body={await response.text()}')
                                raise SpectrumException(f'POST {url} status={response.status} payload={data} body={await response.text()}')
                            # Stream response
                            return StreamingResponse(response.content.iter_chunked(1024 * 1024),
                                                    headers=dict(response.headers),
                                                    media_type='.zip')
                    else:
                        async with session.post(url, data=ujson.dumps(data), headers=headers) as response:
                            text = await response.text()
                            if response.status not in [200, 201]:
                                logger.error(f'POST {url} status={response.status}')
                                logger.error(f'POST {url} payload={data}')
                                logger.error(f'POST {url} body={text}')
                                raise SpectrumException(f'POST {url} status={response.status} payload={data} body={text}')
                            return PlainTextResponse(text, status_code=response.status)
                else:
                    # For files/multipart
                    async with session.post(url, data=data, headers=headers) as response:
                        text = await response.text()
                        if response.status not in [200, 201]:
                            logger.error(f'POST {url} status={response.status}')
                            logger.error(f'POST {url} payload={data}')
                            logger.error(f'POST {url} body={text}')
                            raise SpectrumException(f'POST {url} status={response.status} payload={data} body={text}')
                        if stream and response.status in [200, 201]:
                            return StreamingResponse(response.content.iter_chunked(1024 * 1024),
                                                    headers=dict(response.headers),
                                                    media_type='.zip')
                        return PlainTextResponse(text, status_code=response.status)
            break
        except Exception:
            attempt += 1
            if attempt >= max_attempts:
                raise
            await asyncio.sleep(backoff_factor * attempt)

