import base64
import logging
import httpx
from typing import Optional
from datetime import datetime, timedelta
from app.retrieval_engine.base import RetrievalHandler
from app.schemas import RetrievalPayload, RetrievalResponse, StatusEnum, DocumentDownloadResponse, ListedDocument
from app.tenant_config import TenantSettings

logger = logging.getLogger(__name__)

class MicrosoftEntraRetrievalHandler(RetrievalHandler):
    """
    Retrieval handler for fetching emails from Microsoft Exchange Online via Microsoft Graph API.
    Uses OAuth2 client credentials flow for authentication.
    """

    def __init__(self, settings: TenantSettings, mailbox: str, folder_path: Optional[list[str]] = None):
        super().__init__(settings)

        # Store mailbox passed at instantiation (allows multiple instances with different mailboxes)
        self.mailbox = mailbox

        # Validate required Entra settings
        if not all([
            self.settings.ENTRA_TENANT_ID,
            self.settings.ENTRA_CLIENT_ID,
            self.settings.ENTRA_CLIENT_SECRET,
            self.mailbox
        ]):
            raise ValueError("ENTRA_TENANT_ID, ENTRA_CLIENT_ID, ENTRA_CLIENT_SECRET must be configured, and mailbox must be provided.")

        # Folder path - defaults to Inbox if not specified
        self.folder_path = folder_path if folder_path else ["Postvak IN"]

        logger.info(f"Initialized MicrosoftEntraRetrievalHandler for tenant: {self.settings.TENANT_NAME}")
        logger.info(f"Entra Mailbox: {self.mailbox}")
        logger.info(f"Target folder path: {' / '.join(self.folder_path)}")

        self._access_token = None
        self._token_expiry = None
        self._target_folder_id = None

    async def _get_access_token(self) -> str:
        """Acquire an OAuth2 access token from Microsoft Entra."""
        # Check if we have a valid cached token
        if self._access_token and self._token_expiry and datetime.now() < self._token_expiry:
            return self._access_token

        logger.info("Acquiring new access token from Microsoft Entra")

        token_url = f"https://login.microsoftonline.com/{self.settings.ENTRA_TENANT_ID}/oauth2/v2.0/token"
        payload = {
            "grant_type": "client_credentials",
            "client_id": self.settings.ENTRA_CLIENT_ID,
            "client_secret": self.settings.ENTRA_CLIENT_SECRET,
            "scope": "https://graph.microsoft.com/.default"
        }

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(token_url, data=payload)
                response.raise_for_status()

                token_data = response.json()
                self._access_token = token_data["access_token"]

                # Cache token with expiry (subtract 5 minutes for safety)
                expires_in = token_data.get("expires_in", 3600)
                self._token_expiry = datetime.now() + timedelta(seconds=expires_in - 300)

                logger.info("Successfully acquired access token")
                return self._access_token

        except httpx.HTTPStatusError as e:
            logger.error(f"Failed to acquire access token: {e}")
            logger.error(f"Response: {e.response.text}")
            raise
        except httpx.RequestError as e:
            logger.error(f"Failed to acquire access token: {e}")
            raise

    async def _find_target_folder(self, access_token: str) -> Optional[str]:
        """Traverse the folder path to find the target folder ID."""
        if self._target_folder_id:
            return self._target_folder_id

        logger.info(f"Searching for folder path: {' / '.join(self.folder_path)}")

        headers = {"Authorization": f"Bearer {access_token}"}

        # Start with root folders
        url = f"https://graph.microsoft.com/v1.0/users/{self.mailbox}/mailFolders"
        parent_folder_id = None

        async with httpx.AsyncClient() as client:
            for i, folder_name in enumerate(self.folder_path):
                logger.info(f"  -> Searching for '{folder_name}'...")

                if i > 0 and parent_folder_id:
                    # Search in child folders
                    url = f"https://graph.microsoft.com/v1.0/users/{self.mailbox}/mailFolders/{parent_folder_id}/childFolders"

                try:
                    response = await client.get(url, headers=headers)
                    response.raise_for_status()
                    folders = response.json().get("value", [])

                    # Find the folder (case-insensitive)
                    found_folder = None
                    for folder in folders:
                        if folder.get("displayName", "").lower() == folder_name.lower():
                            found_folder = folder
                            break

                    if not found_folder:
                        logger.error(f"Could not find folder '{folder_name}' in path")
                        logger.error(f"Available folders: {[f.get('displayName') for f in folders]}")
                        return None

                    parent_folder_id = found_folder["id"]
                    logger.info(f"  Found '{found_folder['displayName']}' with ID: {parent_folder_id}")

                except httpx.HTTPStatusError as e:
                    logger.error(f"Error searching for folder: {e}")
                    logger.error(f"Response: {e.response.text}")
                    return None
                except httpx.RequestError as e:
                    logger.error(f"Error searching for folder: {e}")
                    return None

        self._target_folder_id = parent_folder_id
        logger.info(f"Successfully found target folder with ID: {self._target_folder_id}")
        return self._target_folder_id

    def _generate_filename_from_message(self, message: dict) -> str:
        """Generate a filename from a Graph API message object using just the subject."""
        try:
            import re
            subject = message.get("subject", "")
            if subject:
                # Clean subject for filename use (remove invalid characters)
                subject_clean = re.sub(r'[<>:"/\\|?*]', '_', subject).strip()
                return f"{subject_clean}.eml"
            else:
                # Fallback if no subject
                message_id = message.get("id", "unknown")[:20]
                return f"{message_id}_email.eml"

        except Exception as e:
            logger.warning(f"Error generating filename: {e}")
            return f"{message.get('id', 'unknown')[:20]}_email.eml"

    async def list_documents(self, payload: RetrievalPayload) -> RetrievalResponse:
        """List all emails in the target folder since the specified date."""
        logger.info(f"MicrosoftEntraRetrievalHandler: list_documents called for source: {payload.source_identifier}")

        try:
            # Get access token
            access_token = await self._get_access_token()

            # Find target folder
            folder_id = await self._find_target_folder(access_token)
            if not folder_id:
                return RetrievalResponse(
                    status=StatusEnum.FAILED,
                    message=f"Could not find folder: {' / '.join(self.folder_path)}"
                )

            # Build query for messages since date_from
            date_str = payload.date_from.strftime("%Y-%m-%dT00:00:00Z")

            # Fetch messages from the folder
            url = (
                f"https://graph.microsoft.com/v1.0/users/{self.mailbox}"
                f"/mailFolders/{folder_id}/messages"
                f"?$filter=receivedDateTime ge {date_str}"
                "&$select=id,subject,from,receivedDateTime"
                "&$orderby=receivedDateTime desc"
                "&$top=999"
            )

            headers = {"Authorization": f"Bearer {access_token}"}

            async with httpx.AsyncClient() as client:
                response = await client.get(url, headers=headers)
                response.raise_for_status()

                messages = response.json().get("value", [])
            logger.info(f"Found {len(messages)} messages since {date_str}")

            # Convert to ListedDocument objects
            listed_documents = []
            for message in messages:
                message_id = message.get("id")
                if not message_id:
                    continue

                filename = self._generate_filename_from_message(message)

                listed_documents.append(
                    ListedDocument(
                        document_identifier=message_id,
                        filename=filename
                    )
                )

            logger.info(f"Successfully listed {len(listed_documents)} email documents")
            return RetrievalResponse(
                status=StatusEnum.SUCCESS,
                documents=listed_documents
            )

        except Exception as e:
            logger.error(f"Error in list_documents: {str(e)}", exc_info=True)
            return RetrievalResponse(
                status=StatusEnum.FAILED,
                message=f"Error listing documents: {str(e)}"
            )

    async def get_document(self, document_identifier: str) -> DocumentDownloadResponse:
        """Retrieve a specific email by its message ID and return as .eml file."""
        logger.info(f"MicrosoftEntraRetrievalHandler: get_document called for ID: {document_identifier}")

        try:
            # Get access token
            access_token = await self._get_access_token()

            # Fetch the message with MIME content
            url = (
                f"https://graph.microsoft.com/v1.0/users/{self.mailbox}"
                f"/messages/{document_identifier}/$value"
            )

            headers = {"Authorization": f"Bearer {access_token}"}

            async with httpx.AsyncClient() as client:
                response = await client.get(url, headers=headers)
                response.raise_for_status()

                # The response is the raw MIME content of the email
                email_data = response.content

            logger.info(f"Successfully fetched email data for ID {document_identifier}, size: {len(email_data)} bytes")

            # Parse email to generate filename using just the subject
            import email
            msg = email.message_from_bytes(email_data)

            # Try to get a better filename using message headers
            subject = msg.get('Subject', '')
            if subject:
                from email.header import decode_header
                decoded_parts = []
                for bytes_part, charset in decode_header(subject):
                    if isinstance(bytes_part, bytes):
                        try:
                            decoded_parts.append(bytes_part.decode(charset or 'utf-8', errors='replace'))
                        except:
                            decoded_parts.append(bytes_part.decode('utf-8', errors='replace'))
                    else:
                        decoded_parts.append(str(bytes_part))
                subject = "".join(decoded_parts)

            # Generate filename - use just the subject
            import re
            if subject:
                subject_clean = re.sub(r'[<>:"/\\|?*]', '_', subject).strip()
                filename = f"{subject_clean}.eml"
            else:
                filename = f"{document_identifier[:20]}_email.eml"

            # Encode as base64
            content_b64 = base64.b64encode(email_data).decode('utf-8')

            logger.info(f"Successfully retrieved email: {document_identifier} as {filename}")
            return DocumentDownloadResponse(
                document_identifier=filename,
                adapter_document_id=document_identifier,
                content_b64=content_b64
            )

        except Exception as e:
            logger.error(f"Error in get_document for ID {document_identifier}: {str(e)}", exc_info=True)
            raise

    async def retrieve_documents(self, payload: RetrievalPayload) -> RetrievalResponse:
        """
        Retrieve documents from Microsoft Exchange via Graph API.

        This method orchestrates the full retrieval process:
        1. Lists all available documents (emails) since the specified date
        2. Downloads each document as a complete email (.eml file)
        """
        logger.info(f"MicrosoftEntraRetrievalHandler: retrieve_documents called for source: {payload.source_identifier}")

        try:
            # First, list all available documents
            list_response = await self.list_documents(payload)

            if list_response.status != StatusEnum.SUCCESS:
                logger.error(f"Failed to list documents: {list_response.message}")
                return list_response

            if not list_response.documents:
                logger.info("No documents found to retrieve")
                return RetrievalResponse(
                    status=StatusEnum.SUCCESS,
                    message="No documents found for the specified criteria"
                )

            logger.info(f"Found {len(list_response.documents)} documents to retrieve")

            # Track retrieval statistics
            successful_retrievals = 0
            failed_retrievals = 0

            # Retrieve each document
            for document in list_response.documents:
                try:
                    download_response = await self.get_document(document.document_identifier)
                    if download_response:
                        successful_retrievals += 1
                        logger.debug(f"Successfully retrieved document: {document.filename}")
                    else:
                        failed_retrievals += 1
                        logger.warning(f"Failed to retrieve document: {document.filename}")

                except Exception as e:
                    failed_retrievals += 1
                    logger.warning(f"Error retrieving document {document.filename}: {str(e)}")
                    continue

            # Prepare response message
            total_documents = len(list_response.documents)
            message = f"Retrieved {successful_retrievals}/{total_documents} documents successfully"

            if failed_retrievals > 0:
                message += f", {failed_retrievals} failed"
                logger.warning(message)
            else:
                logger.info(message)

            # Return success if at least some documents were retrieved
            status = StatusEnum.SUCCESS if successful_retrievals > 0 else StatusEnum.FAILED

            return RetrievalResponse(
                status=status,
                message=message,
                documents=list_response.documents
            )

        except Exception as e:
            error_message = f"Error in retrieve_documents: {str(e)}"
            logger.error(error_message, exc_info=True)
            return RetrievalResponse(
                status=StatusEnum.FAILED,
                message=error_message
            )