import datetime
from typing import Annotated

from pydantic import BaseModel, EmailStr, computed_field, model_validator
from typing_extensions import Self

from agentdojo.default_suites.v1.tools.types import CalendarEvent, CloudDriveFileID, Email, EmailID, EmailStatus
from agentdojo.functions_runtime import Depends


class EmailContact(BaseModel):
    email: EmailStr
    name: str


class Inbox(BaseModel):
    account_email: EmailStr
    initial_emails: list[Email]
    emails: dict[EmailID, Email] = {}
    contact_list: list[EmailContact] = []
    trash: dict[EmailID, Email] = {}

    @model_validator(mode="after")
    def _create_emails(self) -> Self:
        all_ids = [email.id_ for email in self.initial_emails]
        if len(all_ids) != len(set(all_ids)):
            duplicate_ids = [id_ for id_ in all_ids if all_ids.count(id_) > 1]
            raise ValueError(f"Email IDs must be unique. Duplicates: {duplicate_ids}")
        self.emails = {email.id_: email for email in self.initial_emails}
        return self

    @model_validator(mode="after")
    def _create_contact_list(self) -> Self:
        existing_contacts = set()
        for email in self.initial_emails:
            email_contacts = [*email.recipients, email.sender, *email.cc, *email.bcc]
            for recipient in email_contacts:
                if recipient in existing_contacts:
                    continue
                if recipient not in [contact.email for contact in self.contact_list]:
                    name = recipient.split("@")[0]
                    name = name.replace(".", " ").title()
                    new_contact = EmailContact(email=recipient, name=name)
                    self.contact_list.append(new_contact)
                    existing_contacts.add(recipient)
        return self

    def from_address(self, address: str) -> list[Email]:
        return [email for email in self.emails.values() if email.sender == address]

    def _get_next_id(self) -> EmailID:
        largest_id = max((int(key) for key in self.emails.keys()), default=0)
        return EmailID(largest_id + 1)

    def _get_emails_by_status(self, status: EmailStatus) -> list[Email]:
        return [email for email in self.emails.values() if email.status == status]

    @computed_field
    @property
    def received(self) -> list[Email]:
        return self._get_emails_by_status(EmailStatus.received)

    @computed_field
    @property
    def sent(self) -> list[Email]:
        return self._get_emails_by_status(EmailStatus.sent)

    @computed_field
    @property
    def drafts(self) -> list[Email]:
        return self._get_emails_by_status(EmailStatus.draft)

    def send_email(
        self,
        recipients: list[EmailStr],
        subject: str,
        body: str,
        attachments: list[CloudDriveFileID | CalendarEvent] | None = None,
        cc: list[EmailStr] | None = None,
        bcc: list[EmailStr] | None = None,
    ) -> Email:
        if cc is None:
            cc = []
        if bcc is None:
            bcc = []
        if attachments is None:
            attachments = []
        new_email = Email(
            id_=self._get_next_id(),
            sender=self.account_email,
            body=body,
            subject=subject,
            recipients=recipients,
            status=EmailStatus.sent,
            timestamp=datetime.datetime.now(),
            cc=cc,
            bcc=bcc,
            attachments=attachments,
            read=True,
        )
        self.emails[new_email.id_] = new_email
        return new_email

    def search_emails(self, query: str, sender: EmailStr | None = None) -> list[Email]:
        if sender is not None:
            emails = self.from_address(sender)
        else:
            emails = self.emails.values()
        return [
            email for email in emails if query.lower() in email.subject.lower() or query.lower() in email.body.lower()
        ]

    def delete_email(self, email_id: EmailID) -> Email:
        if email_id not in self.emails:
            raise ValueError(f"Email with ID '{email_id}' not found.")
        email = self.emails.pop(email_id)
        self.trash[email_id] = email
        return email

    def get_unread(self) -> list[Email]:
        unread_emails = [email for email in self.emails.values() if not email.read]
        for email in unread_emails:
            email.read = True
        return unread_emails

    def find_contacts_by_name(self, name: str) -> list[EmailContact]:
        matches = []
        for contact in self.contact_list:
            if name.lower() in contact.name.lower():
                matches.append(contact)
        if len(matches) == 0:
            raise ValueError(f"Contact with name '{name}' not found.")
        return matches

    def find_contacts_by_email(self, email: EmailStr) -> list[EmailContact]:
        matches = []
        for contact in self.contact_list:
            if email.lower() in contact.email.lower():
                matches.append(contact)
        if not matches:
            raise ValueError(f"Contact with email '{email}' not found.")
        return matches

    def get_most_recently_received_email(self) -> Email:
        return max(self.received, key=lambda email: email.timestamp)


def get_unread_emails(inbox: Annotated[Inbox, Depends("inbox")]) -> list[Email]:
    """Returns all the unread emails in the inbox. Each email has a sender, a subject, and a body.
    The emails are marked as read after this function is called."""
    return inbox.get_unread()


def send_email(
    inbox: Annotated[Inbox, Depends("inbox")],
    recipients: list[str],
    subject: str,
    body: str,
    attachments: list[dict] | None = None,
    cc: list[str] | None = None,
    bcc: list[str] | None = None,
) -> Email:
    """Sends an email with the given `body` to the given `address`. Returns a dictionary with the email details.

    :param recipients: The list with the email addresses of the recipients.
    :param subject: The subject of the email.
    :param body: The body of the email.
    :param attachments: The list of attachments to include in the email. If `null`, no attachments are included.
    If the attachment has as "type" "file", then it is expected to have a field "file_id", with the ID of the file in the
    cloud drive. If the attachment has as "type" "event", then it is expected to be a calendar event in the field "event_details".
    A calendar event has the following fields: `title`, `description`, `start_time` (in ISO format), `end_time`
    (in ISO format), `location`, and participants (a list of emails).
    :param cc: The list of email addresses to include in the CC field. If `null`, no email addresses are included.
    :param bcc: The list of email addresses to include in the BCC field. If `null`, no email addresses are included.
    """
    if attachments is not None:
        parsed_attachments = []
        for attachment in attachments:
            if not isinstance(attachment, dict):
                raise ValueError("Attachments must be dictionaries.")
            if "type" not in attachment:
                raise ValueError("Attachment must have a 'type' field.")
            if attachment["type"] == "file" or "file_id" in attachment:
                if "file_id" not in attachment:
                    raise ValueError("Attachment of type 'file' must have a 'file_id' field.")
                parsed_attachments.append(CloudDriveFileID(attachment["file_id"]))
            else:
                if "event_details" not in attachment:
                    raise ValueError("Attachment of type 'event' must have an 'event_details' field.")
                parsed_attachments.append(CalendarEvent(**attachment["event_details"]))
    else:
        parsed_attachments = None
    return inbox.send_email(recipients, subject, body, parsed_attachments, cc, bcc)


def search_emails(
    inbox: Annotated[Inbox, Depends("inbox")],
    query: str,
    sender: str | None = None,
) -> list[Email]:
    """Searches for emails in the inbox that contain the given query in the subject or body. If `address` is provided,
    only emails from that address are searched.

    :param query: The query to search for in the email subject or body. If empty, all emails are returned.
    :param sender: The email address of the sender. If `null`, all emails are searched. Use `search_contacts_by_name`
    if the email address is unknown.
    """
    matches = inbox.search_emails(query, sender)
    if len(matches) == 0:
        raise ValueError("No emails found. Try with a different query.")
    return matches


def delete_email(inbox: Annotated[Inbox, Depends("inbox")], email_id: str) -> str:
    """Deletes the email with the given `email_id` from the inbox.

    :param email_id: The id of the email to delete.
    """
    inbox.delete_email(email_id)
    return f"Email with id {email_id} deleted successfully."


def get_sent_emails(inbox: Annotated[Inbox, Depends("inbox")]) -> list[Email]:
    """Returns all the sent emails in the inbox. Each email has a recipient, a subject, and a body."""
    return inbox.sent


def get_received_emails(inbox: Annotated[Inbox, Depends("inbox")]) -> list[Email]:
    """Returns all the received emails in the inbox. Each email has a sender, a subject, and a body."""
    return inbox.received


def get_draft_emails(inbox: Annotated[Inbox, Depends("inbox")]) -> list[Email]:
    """Returns all the draft emails in the inbox. Each email has a recipient, a subject, and a body."""
    return inbox.drafts


def search_contacts_by_name(inbox: Annotated[Inbox, Depends("inbox")], query: str) -> list[EmailContact]:
    """Finds contacts in the inbox's contact list by name.
    It returns a list of contacts that match the given name.

    :param query: The name of the contacts to search for.
    """
    return inbox.find_contacts_by_name(query)


def search_contacts_by_email(inbox: Annotated[Inbox, Depends("inbox")], query: str) -> list[EmailContact]:
    """Finds contacts in the inbox's contact list by email.
    It returns a list of contacts that match the given email.

    :param query: The email of the contacts to search for.
    """
    return inbox.find_contacts_by_email(query)
