"""
Gmail MCP Client - Integrates Gmail functionality with MCP system.
Uses aiogoogle for native async Gmail API calls.
"""

import asyncio
import base64
import json
import logging
import mimetypes
import os
from email import encoders
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from typing import Any, Dict, List, Tuple

from aiogoogle import Aiogoogle
from aiogoogle.auth.creds import ClientCreds, UserCreds

from .base_client import BaseMCPClient
from .types import MCPTool

GMAIL_SCOPES = [
    "https://www.googleapis.com/auth/gmail.readonly",
    "https://www.googleapis.com/auth/gmail.send",
    "https://www.googleapis.com/auth/gmail.modify",
]


def authenticate_gmail_local(
    client_id: str, client_secret: str, token_file: str = None
) -> dict:
    """
    Authenticate with Gmail using OAuth2 with a local server to catch the callback.
    This is an interactive browser-based flow; inherently sync.
    Called via asyncio.to_thread since it blocks on user interaction.
    """
    import os

    from google.auth.transport.requests import Request
    from google.oauth2.credentials import Credentials
    from google_auth_oauthlib.flow import InstalledAppFlow

    try:
        creds = None
        token_path = token_file or os.path.expanduser("~/.mirix/gmail_token.json")

        if os.path.exists(token_path):
            try:
                creds = Credentials.from_authorized_user_file(token_path, GMAIL_SCOPES)
            except Exception:
                pass

        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                try:
                    creds.refresh(Request())
                except Exception:
                    creds = None

            if not creds:
                client_config = {
                    "installed": {
                        "client_id": client_id,
                        "client_secret": client_secret,
                        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                        "token_uri": "https://oauth2.googleapis.com/token",
                        "auth_provider_x509_cert_url": (
                            "https://www.googleapis.com/oauth2/v1/certs"
                        ),
                        "redirect_uris": [
                            "http://localhost:8080/",
                            "http://localhost:8081/",
                            "http://localhost:8082/",
                        ],
                    }
                }

                flow = InstalledAppFlow.from_client_config(client_config, GMAIL_SCOPES)

                for port in [8080, 8081, 8082]:
                    try:
                        creds = flow.run_local_server(
                            port=port,
                            open_browser=True,
                            access_type="offline",
                            prompt="consent",
                        )
                        break
                    except OSError:
                        if port == 8082:
                            creds = flow.run_local_server(
                                port=0,
                                open_browser=True,
                                access_type="offline",
                                prompt="consent",
                            )

            os.makedirs(os.path.dirname(token_path), exist_ok=True)
            with open(token_path, "w") as token:
                token.write(creds.to_json())

        return {
            "success": True,
            "message": "Gmail authenticated successfully",
            "token_path": token_path,
        }

    except Exception as e:
        return {"success": False, "error": str(e)}


logger = logging.getLogger(__name__)


class GmailMCPClient(BaseMCPClient):
    """Gmail MCP Client using aiogoogle for native async Gmail API calls."""

    def __init__(self, server_config):
        super().__init__(server_config)
        self._gmail_api = None
        self._user_creds: UserCreds = None
        self._client_creds: ClientCreds = None
        self._token_file: str = None
        self._client_id: str = None
        self._client_secret: str = None

    def _load_credentials(self) -> bool:
        """Load credentials from the google-auth token file into aiogoogle creds."""
        try:
            if not os.path.exists(self._token_file):
                return False
            with open(self._token_file) as f:
                token_data = json.load(f)
            if not token_data.get("refresh_token"):
                logger.warning(
                    "Token file missing refresh_token, removing invalid token."
                )
                os.remove(self._token_file)
                return False

            self._user_creds = UserCreds(
                access_token=token_data.get("token"),
                refresh_token=token_data.get("refresh_token"),
                expires_at=token_data.get("expiry"),
                scopes=GMAIL_SCOPES,
                token_uri=token_data.get(
                    "token_uri", "https://oauth2.googleapis.com/token"
                ),
            )
            self._client_creds = ClientCreds(
                client_id=self._client_id,
                client_secret=self._client_secret,
            )
            return True
        except Exception as e:
            logger.error("Failed to load credentials: %s", e)
            return False

    async def _initialize_connection(self, server_config, timeout: float) -> bool:
        """Initialize Gmail connection using OAuth and discover the API."""
        try:
            self._token_file = server_config.token_file or os.path.expanduser(
                "~/.mirix/gmail_token.json"
            )
            self._client_id = server_config.client_id
            self._client_secret = server_config.client_secret

            if not self._load_credentials():
                auth_result = await asyncio.to_thread(
                    authenticate_gmail_local,
                    self._client_id,
                    self._client_secret,
                    self._token_file,
                )
                if not auth_result["success"]:
                    logger.error(
                        "Gmail authentication failed: %s",
                        auth_result.get("error", "Unknown error"),
                    )
                    return False

                if not self._load_credentials():
                    logger.error(
                        "Failed to load credentials after authentication"
                    )
                    return False

            async with Aiogoogle(
                user_creds=self._user_creds, client_creds=self._client_creds
            ) as aiogoogle:
                self._gmail_api = await aiogoogle.discover("gmail", "v1")

            logger.info("Gmail service initialized successfully")
            return True

        except Exception as e:
            logger.error("Failed to initialize Gmail connection: %s", str(e))
            return False

    async def _execute_gmail_request(self, request):
        """Execute a single Gmail API request with auto token refresh."""
        async with Aiogoogle(
            user_creds=self._user_creds, client_creds=self._client_creds
        ) as aiogoogle:
            return await aiogoogle.as_user(request)

    async def list_tools(self) -> List[MCPTool]:
        """Return list of available Gmail tools."""
        return [
            MCPTool(
                name="gmail_send_email",
                description="Send an email using Gmail",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "to": {
                            "type": "string",
                            "description": "Recipient email address",
                        },
                        "subject": {
                            "type": "string",
                            "description": "Email subject",
                        },
                        "body": {
                            "type": "string",
                            "description": "Email body (plain text)",
                        },
                        "cc": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "CC recipients (optional)",
                        },
                        "bcc": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "BCC recipients (optional)",
                        },
                        "html_body": {
                            "type": "string",
                            "description": "HTML version of email body (optional)",
                        },
                        "attachments": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "List of file paths to attach (optional)",
                        },
                    },
                    "required": ["to", "subject", "body"],
                },
            ),
            MCPTool(
                name="gmail_read_emails",
                description="Read emails from Gmail inbox",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": (
                                "Gmail search query (optional, e.g., 'is:unread')"
                            ),
                        },
                        "max_results": {
                            "type": "integer",
                            "description": (
                                "Maximum number of emails to retrieve (default: 10)"
                            ),
                            "default": 10,
                        },
                    },
                },
            ),
            MCPTool(
                name="gmail_get_email",
                description="Get a specific email by ID",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "email_id": {
                            "type": "string",
                            "description": "Gmail message ID",
                        }
                    },
                    "required": ["email_id"],
                },
            ),
        ]

    async def execute_tool(
        self, tool_name: str, tool_args: Dict[str, Any]
    ) -> Tuple[str, bool]:
        """Execute a Gmail tool."""
        self._check_initialized()

        if not await self._ensure_gmail_service():
            return (
                "Gmail authentication required. "
                "Please run the Gmail connection process.",
                True,
            )

        try:
            if tool_name == "gmail_send_email":
                return await self._send_email(tool_args)
            elif tool_name == "gmail_read_emails":
                return await self._read_emails(tool_args)
            elif tool_name == "gmail_get_email":
                return await self._get_email(tool_args)
            else:
                return f"Unknown tool: {tool_name}", True
        except Exception as e:
            logger.error("Error executing Gmail tool %s: %s", tool_name, str(e))
            return f"Error executing tool: {str(e)}", True

    async def _ensure_gmail_service(self) -> bool:
        """Ensure Gmail API is discovered and credentials are available."""
        if self._gmail_api is not None and self._user_creds is not None:
            return True

        try:
            success = await self._initialize_connection(
                self.server_config, timeout=30.0
            )
            if success and self._gmail_api is not None:
                logger.info("Gmail service established successfully")
                return True
            else:
                logger.warning("Gmail service initialization failed")
                return False
        except Exception as e:
            logger.error("Failed to ensure Gmail service: %s", str(e))
            return False

    async def _send_email(self, args: Dict[str, Any]) -> Tuple[str, bool]:
        """Send an email using Gmail API (native async via aiogoogle)."""
        try:
            to = args["to"]
            subject = args["subject"]
            body = args["body"]
            cc = args.get("cc", [])
            bcc = args.get("bcc", [])
            html_body = args.get("html_body")
            attachments = args.get("attachments", [])

            message = self._create_message(
                to, subject, body, cc, bcc, attachments, html_body
            )

            result = await self._execute_gmail_request(
                self._gmail_api.users.messages.send(
                    userId="me", json=message
                )
            )

            return f"Email sent successfully! Message ID: {result['id']}", False

        except Exception as e:
            return f"Failed to send email: {str(e)}", True

    async def _read_emails(self, args: Dict[str, Any]) -> Tuple[str, bool]:
        """Read emails from Gmail (native async via aiogoogle)."""
        try:
            query = args.get("query", "")
            max_results = args.get("max_results", 10)

            async with Aiogoogle(
                user_creds=self._user_creds,
                client_creds=self._client_creds,
            ) as aiogoogle:
                results = await aiogoogle.as_user(
                    self._gmail_api.users.messages.list(
                        userId="me", q=query, maxResults=max_results
                    )
                )

                messages = results.get("messages", [])
                if not messages:
                    return "No emails found", False

                email_details = []
                for msg_entry in messages:
                    msg = await aiogoogle.as_user(
                        self._gmail_api.users.messages.get(
                            userId="me", id=msg_entry["id"]
                        )
                    )

                    headers = msg["payload"].get("headers", [])
                    subject = next(
                        (
                            h["value"]
                            for h in headers
                            if h["name"] == "Subject"
                        ),
                        "No Subject",
                    )
                    sender = next(
                        (
                            h["value"]
                            for h in headers
                            if h["name"] == "From"
                        ),
                        "Unknown Sender",
                    )
                    date = next(
                        (
                            h["value"]
                            for h in headers
                            if h["name"] == "Date"
                        ),
                        "Unknown Date",
                    )

                    email_details.append(
                        {
                            "id": msg_entry["id"],
                            "subject": subject,
                            "from": sender,
                            "date": date,
                        }
                    )

            return json.dumps(email_details, indent=2), False

        except Exception as e:
            return f"Failed to read emails: {str(e)}", True

    async def _get_email(self, args: Dict[str, Any]) -> Tuple[str, bool]:
        """Get a specific email by ID (native async via aiogoogle)."""
        try:
            email_id = args["email_id"]

            message = await self._execute_gmail_request(
                self._gmail_api.users.messages.get(userId="me", id=email_id)
            )

            headers = message["payload"].get("headers", [])
            subject = next(
                (h["value"] for h in headers if h["name"] == "Subject"),
                "No Subject",
            )
            sender = next(
                (h["value"] for h in headers if h["name"] == "From"),
                "Unknown Sender",
            )
            date = next(
                (h["value"] for h in headers if h["name"] == "Date"),
                "Unknown Date",
            )

            body = self._extract_message_body(message["payload"])

            email_data = {
                "id": email_id,
                "subject": subject,
                "from": sender,
                "date": date,
                "body": body,
            }

            return json.dumps(email_data, indent=2), False

        except Exception as e:
            return f"Failed to get email: {str(e)}", True

    def _create_message(
        self,
        to: str,
        subject: str,
        body: str,
        cc: List[str] = None,
        bcc: List[str] = None,
        attachments: List[str] = None,
        html_body: str = None,
    ) -> dict:
        """Create an email message dict for the Gmail API."""
        if html_body or attachments:
            message = MIMEMultipart(
                "alternative" if html_body else "mixed"
            )
        else:
            message = MIMEText(body)
            message["to"] = to
            message["subject"] = subject
            return {
                "raw": base64.urlsafe_b64encode(
                    message.as_bytes()
                ).decode()
            }

        message["to"] = to
        message["subject"] = subject

        if cc:
            message["cc"] = ", ".join(cc)
        if bcc:
            message["bcc"] = ", ".join(bcc)

        text_part = MIMEText(body, "plain")
        message.attach(text_part)

        if html_body:
            html_part = MIMEText(html_body, "html")
            message.attach(html_part)

        if attachments:
            for file_path in attachments:
                if os.path.isfile(file_path):
                    content_type, encoding = mimetypes.guess_type(file_path)
                    if content_type is None or encoding is not None:
                        content_type = "application/octet-stream"

                    main_type, sub_type = content_type.split("/", 1)

                    with open(file_path, "rb") as fp:
                        attachment = MIMEBase(main_type, sub_type)
                        attachment.set_payload(fp.read())

                    encoders.encode_base64(attachment)
                    attachment.add_header(
                        "Content-Disposition",
                        f"attachment; filename={os.path.basename(file_path)}",
                    )
                    message.attach(attachment)
                else:
                    logger.warning(
                        "Attachment file '%s' not found", file_path
                    )

        return {
            "raw": base64.urlsafe_b64encode(message.as_bytes()).decode()
        }

    def _extract_message_body(self, payload):
        """Extract message body from Gmail API payload."""
        body = ""

        if "parts" in payload:
            for part in payload["parts"]:
                if part["mimeType"] == "text/plain":
                    data = part["body"]["data"]
                    body = base64.urlsafe_b64decode(data).decode("utf-8")
                    break
        elif payload["body"].get("data"):
            body = base64.urlsafe_b64decode(
                payload["body"]["data"]
            ).decode("utf-8")

        return body
