from typing import Optional, Union, Literal
import aiohttp
from pathlib import Path
from pdf2image import convert_from_path
from brother_ql.conversion import convert
from brother_ql.raster import BrotherQLRaster
import logging
import asyncio
import time
from PIL import Image
from contextlib import asynccontextmanager
from app.schemas.file import FileInDB

logger = logging.getLogger(__name__)

class LabelPrinterService:
    def __init__(self, printer_api_url: str = "http://localhost:8000"):
        """Initialize the label printer service.
        
        Args:
            printer_api_url: Base URL of the printer API endpoint
        """
        self.printer_api_url = printer_api_url.rstrip('/')
        self.status_url = f"{self.printer_api_url}/status"
        self.print_url = f"{self.printer_api_url}/print"
        self.cache_dir = Path("app/data/cache/prints")
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self._session = None
        self._lock = asyncio.Lock()
        
    async def cleanup(self):
        """Clean up resources, particularly the aiohttp session."""
        if self._session:
            await self._session.close()
            self._session = None
            logger.info("Label printer service session closed")
        
    @asynccontextmanager
    async def _get_session(self):
        """Context manager for aiohttp session."""
        if self._session is None:
            self._session = aiohttp.ClientSession()
        try:
            yield self._session
        except Exception as e:
            logger.error(f"Session error: {e}")
            if self._session:
                await self._session.close()
                self._session = None
            raise
        
    async def _wait_for_printer_ready(self, max_wait: int = 300) -> bool:
        """Wait for the printer to be ready.
        
        Args:
            max_wait: Maximum time to wait in seconds
            
        Returns:
            bool: True if printer is ready, False if timeout
        """
        start_time = time.time()
        while time.time() - start_time < max_wait:
            try:
                async with self._get_session() as session:
                    async with session.get(self.status_url) as response:
                        if response.status == 200:
                            data = await response.json()
                            if data.get('status') == 'ready':
                                return True
                            elif data.get('status') == 'busy':
                                logger.info("Printer is busy, waiting...")
                        elif response.status == 404:
                            logger.error(f"Printer status endpoint not found at {self.status_url}")
                            return False
                        elif response.status == 500:
                            data = await response.json()
                            error_msg = data.get('message', 'Unknown printer error')
                            logger.error(f"Printer error: {error_msg}")
                            raise Exception(f"Printer error: {error_msg}")
            except aiohttp.ClientError as e:
                logger.warning(f"Error checking printer status: {e}")
                if "Cannot connect to host" in str(e):
                    logger.error("Printer reciever is not available")
                    raise Exception("Printer reciever is not available")
            except Exception as e:
                logger.error(f"Unexpected error in _wait_for_printer_ready: {e}")
                return False
            
            await asyncio.sleep(1)
            
        logger.error("Timeout waiting for printer to be ready")
        return False
        
    async def _send_print_request(self, file_path: Union[str, Path]) -> bool:
        """Send print data to printer API.
        
        Args:
            file_path: Path to the binary data file to send to the printer
            
        Returns:
            bool: True if request was successful, False otherwise
        """
        try:
            # Read the binary data from the cache file
            with open(file_path, "rb") as f:
                print_data = f.read()
            
            # Send the request to the printer API using aiohttp
            async with self._get_session() as session:
                async with session.post(
                    self.print_url,
                    data=print_data,
                    headers={'Content-Type': 'application/octet-stream'},
                    timeout=30
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        if data.get('message') == 'Print request processed successfully':
                            return True
                        logger.error(f"Unexpected success response: {data}")
                        return False
                    elif response.status == 404:
                        logger.error(f"Print endpoint not found at {self.print_url}")
                        return False
                    elif response.status == 429:
                        logger.error("Printer is busy")
                        return False
                    elif response.status == 500:
                        data = await response.json()
                        error_msg = data.get('message', 'Unknown printer error')
                        logger.error(f"Printer error: {error_msg}")
                        raise Exception(f"Printer error: {error_msg}")
                    else:
                        data = await response.json()
                        logger.error(f"Print request failed with status {response.status}: {data.get('message')}")
                        return False
                
        except aiohttp.ClientError as e:
            logger.error(f"Error sending print request: {str(e)}")
            return False
        except Exception as e:
            logger.error(f"Unexpected error in _send_print_request: {e}")
            return False
            
    async def print_file(self, file_path: Union[str, Path, FileInDB], label_size: Literal["dk1201", "dk1241"], label_type: Optional[Literal["address_label", "packing_slip", "set_label", "return_label", "pirate_ship_label"]] = None, copies: Optional[int] = None) -> bool:
        """Print a PDF or PNG file to the label printer.
        
        Args:
            file_path: Path to the PDF or PNG file, or a FileInDB object
            label_size: Size of label to use ("dk1201" or "dk1241")
            label_type: Type of label to use ("address_label" or "packing_slip" or "set_label")
            copies: Optional number of copies to print. If None, prints once.
            
        Returns:
            bool: True if print was successful, False otherwise
        """
        async with self._lock:  # Ensure only one print operation at a time
            try:
                if file_path is None:
                    logger.error("No file path provided")
                    return False
                    
                # Handle FileInDB objects
                if isinstance(file_path, FileInDB):
                    file_path = file_path.path
                    
                file_path = Path(file_path)
                if not file_path.exists():
                    logger.error(f"File not found: {file_path}")
                    return False

                # Wait for printer to be ready
                if not await self._wait_for_printer_ready():
                    logger.error("Printer not ready after waiting")
                    return False

                if file_path.suffix.lower() == '.pdf':
                    # Convert PDF to images
                    logger.info(f"Converting PDF {file_path} to images")
                    images = convert_from_path(file_path)
                    if not images:
                        logger.error(f"No images could be extracted from {file_path}")
                        return False
                    logger.info(f"Successfully converted PDF to {len(images)} images")
                else:
                    # For PNG files, we can use them directly
                    images = [Image.open(file_path)]
                
                # Process each page
                for i, image in enumerate(images):
                    logger.info(f"Processing page {i+1} with dimensions {image.size}")
                    # Resize image based on label size and type
                    resized_image = image.copy()  # Create a copy to work with
                    
                    # Store the original label size before we modify it
                    original_label_size = label_size
                    
                    # Handle resizing based on label size and type
                    if original_label_size == "dk1241":
                        if label_type == "packing_slip":
                            resized_image = resized_image.resize((1164, 1660), Image.Resampling.LANCZOS)
                        elif label_type == "address_label":
                            resized_image = resized_image.resize((1660, 1164), Image.Resampling.LANCZOS)
                        else:
                            resized_image = resized_image.resize((1164, 1660), Image.Resampling.LANCZOS)
                    elif original_label_size == "dk1201":
                        resized_image = resized_image.resize((991, 306), Image.Resampling.LANCZOS)
                    
                    # if file path contains address_label, rotate image 90 degrees
                    if label_type == "address_label" or label_type == "set_label" or label_type == "return_label":
                        rotate = "90"
                        cut = False
                    else:
                        rotate = "0"
                        cut = True

                    # Convert to label format
                    qlr = BrotherQLRaster("QL-1100")
                    qlr.exception_on_warning = True
                    
                    # Get label size based on type
                    brother_label_size = "29x90" if original_label_size == "dk1201" else "102x152"
                    
                    converted_image = convert(
                        qlr=qlr,
                        images=[resized_image],
                        label=brother_label_size,
                        rotate=rotate,
                        threshold=70.0,
                        dither=False,
                        compress=False,
                        red=False,
                        dpi_600=False,
                        #hq=True,
                        hq=False,
                        cut=cut
                    )
                    
                    # Cache the converted binary data
                    cache_path = self.cache_dir / f"{file_path.stem}_{brother_label_size}_page_{i+1}_converted.bin"
                    with open(cache_path, "wb") as f:
                        f.write(converted_image)
                    
                    if copies:
                        # Send to API for each copy
                        for copy in range(copies):
                            logger.info(f"Printing copy {copy + 1} of {copies}")
                            if not await self._send_print_request(cache_path):
                                logger.error(f"Failed to print page {i+1}, copy {copy + 1}")
                                return False
                            
                            # Wait for printer to be ready before next copy or page
                            if copy < copies - 1 or i < len(images) - 1:
                                if not await self._wait_for_printer_ready():
                                    logger.error("Printer not ready for next copy/page")
                                    return False
                    else:
                        # Send to API once (original behavior)
                        if not await self._send_print_request(cache_path):
                            logger.error(f"Failed to print page {i+1}")
                            return False
                        
                        # Wait for printer to be ready before processing next page
                        if i < len(images) - 1:  # Don't wait after the last page
                            if not await self._wait_for_printer_ready():
                                logger.error("Printer not ready for next page")
                                return False
                
                return True
                
            except Exception as e:
                logger.error(f"Error printing file {file_path}: {str(e)}")
                return False 