import base64
import datetime
import logging
import os
import io
from PIL import Image
import requests

from app.appTypes import PathSegment, SiteConfig
from app.clients import async_openai_client, openai_client
from app.services import usage_history
from app.utilities.path_utils import PathUtils

from dotenv import load_dotenv
load_dotenv()

path_utils = PathUtils()

ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}


def get_default_image_format() -> str:
    """Read IMAGE_STORAGE_FORMAT from env. Returns 'png' or 'jpeg', defaults to 'png'."""
    fmt = os.environ.get('IMAGE_STORAGE_FORMAT', 'png').lower()
    if fmt in ('png', 'jpeg'):
        return fmt
    return 'png'


class ImageService:

    def __init__(self):
        self._last_base64_data = None

    async def create_image(self, prompt: str, site: SiteConfig, size: str = "1536x1024", quality: str = "medium", output_format: str = None) -> list:
        """
        Generate image using OpenAI Images API with gpt-image-1.5

        Args:
            prompt: Text description of the image to generate
            site: Site configuration for file storage
            size: Image dimensions (default: 1536x1024 landscape)
            quality: Image quality - low, medium, high (default: medium)
            output_format: Image format - png, jpeg, webp (default: env IMAGE_STORAGE_FORMAT)

        Returns:
            List of image URLs
        """
        resolved_format = output_format or get_default_image_format()

        response = await async_openai_client.images.generate(
            model="gpt-image-1.5",
            prompt=prompt,
            size=size,
            quality=quality,
            output_format=resolved_format,
        )

        image_data = response.data[0].b64_json
        if image_data:
            usage_history.record_image(
                model='gpt-image-1.5',
                size=size,
                quality=quality,
                operation='article_image',
                b64=image_data,
                site_slug=site.slug,
            )
            self._last_base64_data = image_data
            return self.save_b64_and_get_url(image_data, site, resolved_format)

        return []

    def save_b64_and_get_url(self, image_base64: str, site: SiteConfig, output_format: str = "png") -> list:
        """
        Helper method to save base64 image and return URL

        Args:
            image_base64: Base64 encoded image data
            site: Site configuration for file storage
            output_format: Image format for saving

        Returns:
            List containing the image URL
        """
        create_dir = path_utils.get_site_files_dir(site) + PathSegment.CHATBOT_CREATE.value
        file_path = self.save_base64_image(create_dir, image_base64, output_format=output_format)

        if not file_path:
            logging.error(f"Failed to persist image to {create_dir} (see prior error)")
            return []

        myHost = os.environ.get('DRUPAL_HOST')
        if not myHost:
            logging.error("DRUPAL_HOST environment variable not set")
            return []

        my_url = myHost + file_path[len(path_utils.get_drupal_base_dir()):]
        return [my_url]

    def get_last_generated_base64(self):
        """Return and clear the base64 data from the last generated image"""
        data = self._last_base64_data
        self._last_base64_data = None  # Clear after retrieval to free memory
        return data

    def reduce_png_size(self, base64_data, output_width=None, output_height=None, output_quality=80):
        image_data = base64.b64decode(base64_data)
        image = Image.open(io.BytesIO(image_data))
        if output_width and output_height:
            image = image.resize((output_width, output_height))
        output_buffer = io.BytesIO()
        image.save(output_buffer, format='PNG',
                   optimize=True, quality=output_quality)
        compressed_data = output_buffer.getvalue()
        compressed_base64 = base64.b64encode(compressed_data).decode('utf-8')
        return compressed_base64

    def save_base64_image(self, directory, base64_image, file_prefix="", output_format="png"):
        try:
            fmt = output_format.value if hasattr(output_format, 'value') else output_format
            ext = fmt if fmt in ('png', 'jpeg', 'webp') else 'png'

            # Only compress for PNG output; JPEG/WebP are already optimized by the API
            if ext == 'png':
                save_data = self.reduce_png_size(
                    base64_image, output_width=None, output_height=None, output_quality=60)
            else:
                save_data = base64_image

            os.makedirs(directory, exist_ok=True)
            logging.info(
                f"saving b64 image to {directory} prefix {file_prefix} format {ext}")
            current_date = datetime.date.today().strftime("%Y-%m-%d")
            # Each logical image owns one basename across every supported
            # extension, so transcodes group with their source. A fresh save
            # in any format advances the -N suffix if another format already
            # claimed the basename today.
            base_stem = f"{file_prefix}{current_date}"
            stem = base_stem
            suffix = 0
            while self._basename_in_use(directory, stem):
                suffix += 1
                stem = f"{base_stem}-{suffix}"
            file_name = f"{stem}.{ext}"
            file_path = os.path.join(directory, file_name)

            image_data = base64.b64decode(save_data)
            with open(file_path, "wb") as image_file:
                image_file.write(image_data)

            logging.info(f"Image saved successfully as {file_path}")
            try:
                self.get_or_create_thumbnail(file_path)
            except Exception as thumb_err:
                logging.warning(f"Could not pre-generate thumbnail for {file_path}: {thumb_err}")
            return file_path
        except Exception as e:
            logging.error(
                f"Error saving b64 image to {directory}: {e}", exc_info=True)
            return None

    def png_to_jpg(self, png_path: str, jpg_path: str, compression_level: int):
        with Image.open(png_path) as img:
            img = img.convert("RGB")
            img.save(jpg_path, "JPEG", quality=compression_level)

    def png_to_webp(self, png_path: str, webp_path: str, compression_level: int):
        with Image.open(png_path) as img:
            # Keep original mode to preserve transparency if present
            img.save(webp_path, "WEBP", quality=compression_level)

    def get_compressed_jpg_size(self, png_path: str, compression_level: int) -> int:
        with Image.open(png_path) as img:
            img = img.convert("RGB")
            with io.BytesIO() as output:
                img.save(output, "JPEG", quality=compression_level)
                size = output.tell()
        return size

    @staticmethod
    def _basename_in_use(directory: str, stem: str) -> bool:
        for candidate_ext in ALLOWED_EXTENSIONS:
            if os.path.exists(os.path.join(directory, f"{stem}{candidate_ext}")):
                return True
        return False

    def get_unique_webp_path(self, png_path):
        # Variants of the same source share a basename — the webp transcode
        # just swaps the extension. Re-transcoding overwrites the prior webp.
        return png_path.rsplit('.', 1)[0] + '.webp'

    def get_or_create_thumbnail(self, source_path: str, size: int = 200, quality: int = 70) -> str:
        """Generate a thumbnail for the given image, cached to a thumb/ subdirectory.

        Returns the absolute path to the thumbnail file.
        """
        source_dir = os.path.dirname(source_path)
        thumb_dir = os.path.join(source_dir, 'thumb')
        basename = os.path.splitext(os.path.basename(source_path))[0]
        thumb_path = os.path.join(thumb_dir, f"{basename}.jpg")

        # Cache hit: thumbnail exists and is newer than source
        if os.path.exists(thumb_path):
            if os.path.getmtime(thumb_path) >= os.path.getmtime(source_path):
                return thumb_path

        os.makedirs(thumb_dir, exist_ok=True)

        with Image.open(source_path) as img:
            img = img.convert("RGB")
            img.thumbnail((size, size))
            img.save(thumb_path, "JPEG", quality=quality)

        return thumb_path

    # rename this function - it adds the image to the media library
    def compress_and_save(self, site, token, url, compression_level):
        logging.info(f'token is {token}')

        png_path = path_utils.url_to_file_path(url, site)
        webp_path = self.get_unique_webp_path(png_path)

        self.png_to_webp(png_path, webp_path, compression_level)

        try:
            self.get_or_create_thumbnail(webp_path)
        except Exception as thumb_err:
            logging.warning(f"Could not pre-generate thumbnail for {webp_path}: {thumb_err}")

        filename = webp_path.split('/')[-1]
        logging.info(f'filename is {filename}')

        media_url = f"{site.wp_host}/wp-json/wp/v2/media"

        headers = {
            'Content-Disposition': f'attachment; filename={filename}',
            'Authorization': f'Bearer {token}'
        }

        with open(webp_path, 'rb') as img:
            file = {'file': img}
            logging.info(f'file: {file}')
            myResponse = requests.post(media_url, headers=headers, files=file)

        if myResponse.status_code == 201:
            return myResponse.json()
        else:
            logging.error(f'WordPress media upload failed with status {myResponse.status_code}')
            logging.error(f'Response body: {myResponse.text}')
            return {
                'error': myResponse.status_code,
                'message': myResponse.text
            }