import gzip
import json
import pickle
import re
import sys
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from os import environ
from pathlib import Path
from typing import Mapping, Sequence, Type, TypeVar, cast

from rich.table import Table
from rich.text import Text
from yaralyzer.util.helpers.interaction_helper import ask_to_proceed

from epstein_files.documents.config.doc_cfg import Metadata
from epstein_files.documents.config.manual_config import create_configs
from epstein_files.documents.document import Document, DocType
from epstein_files.documents.documents.doc_types_mixin import DocTypesMixin
from epstein_files.documents.documents.search_result import SearchResult
from epstein_files.documents.doj_file import DojFile
from epstein_files.documents.email import EMAILERS_TO_ALWAYS_TRUNCATE, Email
from epstein_files.documents.emails.constants import UNINTERESTING_EMAILERS
from epstein_files.documents.emails.dropsite_email import DropsiteEmail
from epstein_files.documents.emails.emailers import CONFIGURED_NON_ENTITIES, ENTITIES_DICT
from epstein_files.documents.emails.multi_email_files import split_up_multi_email_files
from epstein_files.documents.json_file import JsonFile
from epstein_files.documents.messenger_log import MSG_REGEX, MessengerLog
from epstein_files.documents.messenger_log_pdf import IMESSAGE_PDF_IDS, MessengerLogPdf
from epstein_files.documents.other_file import OtherFile
from epstein_files.output.rich import TABLE_TITLE_STYLE, console
from epstein_files.people.entity import Entity
from epstein_files.people.person import PEOPLE_BIOS, Person
from epstein_files.util.constant.strings import *
from epstein_files.util.constants import *
from epstein_files.util.env import args, logger
from epstein_files.util.helpers.data_helpers import flatten, json_safe, sort_dict_by_keys, uniquify, uniq_sorted
from epstein_files.util.helpers.file_helper import all_txt_paths, doj_txt_paths, extract_file_id, file_size_str
from epstein_files.util.timer import Timer

# Lists of properties to copy into duplicate documents (will be preceded with 'extracted_')
PROPS_TO_COPY = ['author', 'timestamp']
EMAIL_PROPS_TO_COPY = ['recipients']
SLOW_FILE_SECONDS = 1.0


@dataclass
class EpsteinFiles(DocTypesMixin):
    """
    Attributes:
        file_paths (list[Path]): paths to Epstein related text documents
        documents (list[Document]): all parsed Documents except the emails with was_split_up flag
        _uninteresting_ccs (list[Name]): names of tangential people who were just CCed once or similar
    """
    file_paths: list[Path] = field(init=False)
    # Derived fields
    _empty_file_ids: set[str] = field(default_factory=set)
    _emailers: list[Person] = field(default_factory=list)
    _uninteresting_ccs: list[Name] = field(default_factory=list)

    def __post_init__(self):
        """Iterate through files and build appropriate objects."""
        self.file_paths = sorted(all_txt_paths(), reverse=True)
        self._documents = self._load_file_paths(self.file_paths)
        self._documents += self._split_up_big_emails()
        self._finalize_data_and_write_to_disk()

    @classmethod
    def get_files(cls, timer: Timer | None = None) -> 'EpsteinFiles':
        """Alternate constructor that reads/writes a pickled version of the data ('timer' arg is for logging)."""
        timer = timer or Timer()

        if args.pickle_path.exists() and not args.overwrite_pickle:
            with gzip.open(args.pickle_path, 'rb') as file:
                epstein_files = pickle.load(file)
                timer_msg = f"Loaded {len(epstein_files.documents):,} documents from '{args.pickle_path}'"
                timer.print_at_checkpoint(f"{timer_msg} ({file_size_str(args.pickle_path)})")

            if args.load_new:
                epstein_files.load_new_files()
            elif args.reload_doj:
                epstein_files.reload_doj_files()

            return epstein_files

        logger.warning(f"Building new cache file, this will take a few minutes...")
        epstein_files = EpsteinFiles()
        num_synthetic_cfgs = len([c for c in CONFIGS_BY_ID.values() if c.is_synthetic])
        timer.print_at_checkpoint(f'Processed {len(epstein_files.file_paths):,} files, {num_synthetic_cfgs} synthetic configs')
        return epstein_files

    @property
    def counterparties_dict(self) -> dict[Name, list[Name]]:
        """Keys are names, values are lists of all the people who sent/received communication with that person."""
        return sort_dict_by_keys({p.name: p.counterparties for p in self.emailers})

    @property
    def documents(self) -> Sequence[Document]:
        """Overloads mixing @property to exclude split up big files."""
        return [d for d in self._documents if not (isinstance(d, Email) and d._was_split_up)]

    @property
    def emails_with_attachments(self) -> list[Email]:
        return [e for e in self.emails if e.attached_docs]

    @property
    def emailers(self) -> list[Person]:
        """All the people who sent or received an email."""
        return self._emailers

    @property
    def uninteresting_emailers(self) -> list[Name]:
        """Emailers whom we don't want to print a separate section for because they're just CCed."""
        if '_uninteresting_emailers' not in vars(self):
            self._uninteresting_emailers = uniq_sorted(UNINTERESTING_EMAILERS + self._uninteresting_ccs)

        return self._uninteresting_emailers

    def docs_for(self, name: Name) -> list[Document]:
        """All documents with `name` as the author or a recipient (not just someone who is mentioned)."""
        docs = flatten([fxn(name) for fxn in [self.emails_for, self.imessage_logs_for, self.other_files_for]])
        return Document.sort_by_timestamp(Document.uniquify(docs))

    def email_author_counts(self) -> dict[Name, int]:
        """Returns dict counting up how many emails were written by each person."""
        return {
            person.name: len(person.unique_emails_by)
            for person in self.emailers if len(person.unique_emails_by) > 0
        }

    def email_authors_to_device_signatures(self) -> dict[str, set[str]]:
        """Mapping of authors to all the device signatures identified in their history."""
        signatures = defaultdict(set)

        for email in [e for e in self.unique_emails if e.sent_from_device]:
            signatures[email.author_str].add(email.sent_from_device)

        return signatures

    def email_device_signatures_to_authors(self) -> dict[str, set[str]]:
        """Mapping of device signatures to all the users who ever signed an email with them."""
        signatures = defaultdict(set)

        for email in [e for e in self.unique_emails if e.sent_from_device]:
            signatures[email.sent_from_device].add(email.author_str)

        return signatures

    def email_recipient_counts(self) -> dict[Name, int]:
        """Returns dict with names as keys and number of emails received as values."""
        return {p.name: len(p.unique_emails_to) for p in self.emailers if len(p.unique_emails_to) > 0}

    def email_signature_substitution_counts(self) -> dict[str, int]:
        """Return the number of times an email signature was replaced with "<...snipped...>" for each author."""
        substitution_counts = defaultdict(int)

        for email in self.emails:
            for name, num_replaced in email.signature_substitution_counts.items():
                substitution_counts[name] += num_replaced

        return substitution_counts

    def emails_by(self, author: Name) -> list[Email]:
        """All emails sent by `author` (including dupes!)."""
        return Document.sort_by_timestamp([e for e in self.emails if e.author == author])

    def emails_for(self, name: Name) -> list[Email]:
        """All emails to or from 'name' sorted chronologically (including dupes!)."""
        emails = self.emails_by(name) + self.emails_to(name)
        emails += [e for e in self.emails if name and name == e._config.show_with_name]  # Add show_with_name emails

        if len(emails) == 0:
            logger.warning(f"No emails found for '{name}'")

        return Document.sort_by_timestamp(Document.uniquify(emails))

    def emails_to(self, name: Name) -> list[Email]:
        """All `Email`s sent to `name` (including dupes!)."""
        if name is None:
            emails = [e for e in self.emails if len(e.recipients) == 0 or None in e.recipients]
        else:
            emails = [e for e in self.emails if name in e.recipients]

        return Document.sort_by_timestamp(emails)

    def get_ids(self, ids: list[str], rebuild: bool = False) -> Sequence[Document]:
        """Get `Document` objects for `file_ids`. If `rebuild` is True then rebuild `Document` from .txt file."""
        docs = [self.docs_by_id[id] for id in ids]

        if len(docs) != len(ids):
            logger.error(f"{len(ids)} file IDs provided but only {len(docs)} documents found! ids: {ids}")

        return [doc.reload() if rebuild else doc for doc in docs]

    def get_id(self, file_id: str, rebuild: bool = False, required_type: Type[DocType] = Document) -> DocType:
        """Singular ID version of `get_ids()` but with option to require a type of document subclass."""
        doc = self.get_ids([file_id], rebuild)[0]

        if not isinstance(doc, required_type):
            raise ValueError(f"No {required_type.__name__} found for {file_id} (found {doc})")

        return doc

    def grep_documents(self, pattern: re.Pattern | str, names: list[Name] | None = None) -> list[SearchResult]:
        """Find documents whose text matches `pattern` optionally limited to only docs involving `name`)."""
        documents = [d for d in self._documents if (not names) or d.author in names]
        results: list[SearchResult] = []

        for doc in documents:
            if doc.is_duplicate:
                continue

            lines = doc.lines_matching(pattern)

            if args.min_line_length:
                lines = [line for line in lines if len(line.line) > args.min_line_length]

            if len(lines) > 0:
                results.append(SearchResult(doc, lines))

        return results

    def imessage_logs_for(self, name: Name) -> list[MessengerLog]:
        """Return `MessengerLog` objects where Epstein's counterparty is `name`."""
        return [log for log in self.imessage_logs if name == log.author]

    def json_metadata(self) -> str:
        """Create a JSON string containing metadata for all the files."""
        def _sorted_metadata(docs: Sequence[Document]) -> list[Metadata]:
            return [json_safe(d.metadata) for d in Document.sort_by_id(docs)]

        metadata = {
            'files': {
                Email.__name__: _sorted_metadata(self.emails),
                JsonFile.__name__: _sorted_metadata(self.json_files),
                MessengerLog.__name__: _sorted_metadata(self.imessage_logs),
                OtherFile.__name__: _sorted_metadata(self.non_json_other_files),
            },
            'people': {name: bio.plain for name, bio in PEOPLE_BIOS.items()},
        }

        return json.dumps(metadata, indent=4, sort_keys=True)

    def load_new_files(self) -> None:
        """Load any new files detected in the hierarchy."""
        new_docs = self._load_file_paths(self._new_files())

        if not new_docs:
            logger.warning(f"No new files found, doing nothing.")
            return
        elif (cfgs := create_configs(new_docs)):  # Exit if new configs were created
            console.print(f"\n{len(cfgs)} config objects created, exiting...", style='bold yellow1')
            sys.exit()

        self._finalize_data_and_write_to_disk(new_docs)

    def other_files_for(self, name: Name) -> list[OtherFile]:
        """Get files with author `name` or that are marked `show_with_name`."""
        if name is None:
            return []
        else:
            return [f for f in self.other_files if name in [f.author, f._config.show_with_name]]

    def overview_table(self) -> Table:
        """Table showing file counts by type."""
        title = Text('Files Overview ', TABLE_TITLE_STYLE)
        title.append('(last updated: ', 'dim').append(datetime.now().date().isoformat(), 'cyan').append(')', 'dim')
        table = Document.files_summary_table(title, 'File Type')
        table.add_row('Emails', *Document.file_summary_row(self.emails))
        table.add_row('iMessage Logs', *Document.file_summary_row(self.imessage_logs))
        table.add_row('JSON Data', *Document.file_summary_row(self.json_files, True))
        table.add_row('Other', *Document.file_summary_row(self.non_json_other_files))
        return table

    def person_objs(self, names: list[Name]) -> list[Person]:
        """Construct Person objects for a list of names."""
        return [
            Person(
                name=name,
                _documents=self.docs_for(name),
                is_interesting=False if name in self.uninteresting_emailers else None
            )
            for name in names
        ]

    def reload_doj_files(self) -> None:
        """Reload only the DOJ PDF extracts (keep HOUSE_OVERSIGHT stuff unchanged)."""
        def doj_file_counts_str():
            return f"(have {len(self.all_doj_files)}, {len(self.doj_files)} non-email)"

        timer = Timer()
        logger.warning(f"Reloading all DOJ files {doj_file_counts_str()}...")
        self._finalize_data_and_write_to_disk(self._load_file_paths(doj_txt_paths()))
        timer.print_at_checkpoint(f"Reloaded {len(self.doj_files)} DOJ files {doj_file_counts_str()}")

    def repair_ids(self, ids: list[str]) -> None:
        """Reload the `ids` and save updated pickle file (also loads new files)."""
        ids = uniquify(ids)

        if ids == ['EMAIL']:
            doc_paths = [d.file_path for d in self.emails]
            msg = f"Repairing all {len(doc_paths)} emails..."
        else:
            doc_paths = [d.file_path for d in self._documents if d.file_id in ids]
            msg = f"Repairing {len(ids)} file IDs"

            if len(doc_paths) != len(ids):
                raise RuntimeError(f"{len(ids)} specified but only {len(doc_paths)} Document objects found!")

        if (new_paths := self._new_files()):
            msg +=  f" (also loading {len(new_paths)} new files)"
            doc_paths += new_paths

        logger.warning(msg)
        repaired_docs = self._load_file_paths(doc_paths)

        if environ.get('RESPLIT_BIG_EMAILS'):
            repaired_docs += self._split_up_big_emails()
            logger.warning(f"  (RESPLIT_BIG_EMAILS so now have {len(repaired_docs)} repaired_docs)")

        self._finalize_new_docs_if_approved(repaired_docs)

    def _copy_duplicate_doc_properties(self) -> None:
        """Ensure dupe docs have the properties of the docs they duplicate to capture any repairs, config etc."""
        for doc in self.documents:
            if not doc._config.duplicate_of_id:
                continue

            original = self.get_id(doc._config.duplicate_of_id)
            props_to_copy = (EMAIL_PROPS_TO_COPY + PROPS_TO_COPY) if isinstance(doc, Email) else PROPS_TO_COPY
            props_to_copy = [f"extracted_{prop}" for prop in props_to_copy]

            for field_name in props_to_copy:
                original_prop = getattr(original, field_name)
                duplicate_prop = getattr(doc, field_name)

                if original_prop != duplicate_prop:
                    doc._warn(f"Replacing {field_name} {duplicate_prop} with {original_prop} from duplicated '{original.file_id}'")
                    setattr(doc, field_name, original_prop)

    def _finalize_data_and_write_to_disk(self, new_docs: list[Document] | None = None) -> None:
        """Handle computation of fields related to uninterestingness, relationships between documents, etc."""
        new_docs = new_docs or []

        if new_docs:
            old_num_docs = len(self._documents)
            new_doc_ids = [d.file_info.file_id for d in new_docs]
            self._documents = [d for d in self._documents if d.file_info.file_id not in new_doc_ids]  # Remove existing
            logger.warning(f"Adding {len(new_docs)} Documents (replacing {old_num_docs - len(self._documents)} existing)")
            self._documents += new_docs

        self._set_uninteresting_ccs()
        self._copy_duplicate_doc_properties()
        self._emailers = self.person_objs(flatten([e.participants for e in self.emails]))
        self._find_email_attachments_and_set_is_first_for_user()
        self._documents = Document.sort_by_timestamp(self._documents)
        self._docs_by_id = {doc.file_id: doc for doc in self._documents}
        self._save_to_disk()

    def _find_email_attachments_and_set_is_first_for_user(self) -> None:
        """Add documents with configured `attached_to_email_id` to the `Email`s they're attached to."""
        for email in self.emails:
            email.attached_docs = []  # Remove all attachments before re-finding them in case it's a repair

        email_attachments = [f for f in self.other_files if f._config.attached_to_email_id]
        logger.warning(f"Finding homes for {len(email_attachments)} known email attachments...")

        for attachment in email_attachments:
            email = self.get_id(attachment._config.attached_to_email_id, required_type=Email)
            email.attached_docs.append(attachment)

            if email.is_duplicate:
                raise ValueError(f"Cannot attach {attachment.file_id} to duplicate email {email}")

            # Set the attachment timestamp to that of the email if the attachment has no configured timestamp
            if attachment.timestamp and (attachment.timestamp != email.timestamp) and not attachment._config.timestamp:
                attachment._warn(f"Overwriting '{attachment.timestamp}' with {email}'s timestamp {email.timestamp}")
                attachment.extracted_timestamp = email.timestamp

        # Set the is_persons_first_email flag on the earliest Email we have for each person (sent or received)
        for emailer in self.emailers:
            if emailer.name not in EMAILERS_TO_ALWAYS_TRUNCATE and emailer.unique_emails:
                emailer.unique_emails[0].is_persons_first_email = True

    def _finalize_new_docs_if_approved(self, new_docs: list[Document]) -> None:
        """Same as _finalize_data_and_write_to_disk() but prints new docs and asks for permission."""
        if len(new_docs) < 100:
            console.print(*new_docs)  # Print for user review
            logger.warning(f"Finalizing {len(new_docs)} files...")
        else:
            logger.warning(f"Too many new documents ({len(new_docs)}) to show previews...")

        Document._print_ids(new_docs, 'newly loaded or repaired by EpsteinFiles')
        ask_to_proceed("Looks good?")
        self._finalize_data_and_write_to_disk(new_docs)

    def _load_file_paths(self, file_paths: list[Path]) -> list[Document]:
        """Load a list of file paths into a list of `Document` object subclasses."""
        docs: list[Document] = []

        for file_path in file_paths:
            doc_timer = Timer(decimals=2)
            document = Document(file_path)

            if document.length == 0:
                if document.file_id not in self._empty_file_ids:
                    document._warn(f"Skipping empty file...")
                    self._empty_file_ids.add(document.file_id)

                continue

            cls = document_cls(document)
            # docs.append(cls(file_path, lines=document.lines, text=document.text))
            docs.append(cls(file_path))  # TODO: needs to reload DropsiteEmail
            logger.info(str(docs[-1]))

            if doc_timer.seconds_since_start() > SLOW_FILE_SECONDS:
                doc_timer.print_at_checkpoint(f"Slow file: {docs[-1]} processed")

        return docs

    def _new_files(self) -> list[Path]:
        """Find any files that don't exist in the collection. Has side effect of setting `self.file_paths`."""
        self.file_paths = sorted(all_txt_paths(), reverse=True)
        current_doc_ids = [id for id in self.docs_by_id.keys()] + list(self._empty_file_ids)

        return [
            p for p in self.file_paths
            if extract_file_id(p) not in current_doc_ids
        ]

    def _save_to_disk(self) -> None:
        """Write a pickled version of this `EpsteinFiles` object with all documents etc."""
        with gzip.open(args.pickle_path, 'wb') as file:
            pickle.dump(self, file)
            logger.warning(f"Pickled data to '{args.pickle_path}' ({file_size_str(args.pickle_path)})...")

    def _set_uninteresting_ccs(self) -> None:
        """Extract the recipients of emails configured has having uninteresting CCs or BCCs."""
        for email in self.emails:
            if email._config.has_uninteresting_ccs:
                self._uninteresting_ccs += email.recipients

            if email._config.has_uninteresting_bccs:
                self._uninteresting_ccs += [bcc.lower() for bcc in cast(list[str], email.header.bcc)]

        self._uninteresting_ccs = uniq_sorted(self._uninteresting_ccs)
        logger.info(f"Extracted uninteresting_ccs: {self._uninteresting_ccs}")

    def _split_up_big_emails(self) -> list[Email]:
        """Find the big emails that we want to split up into smaller emails."""
        big_emails = [e for e in self._documents if isinstance(e, Email) and e._was_split_up]

        big_emails = big_emails or [
            *self.emails_by(CHRISTOPHER_DILORIO),
            *[self.get_id(LEON_BLACK_EMAIL_ID, required_type=Email)]
        ]

        new_emails = split_up_multi_email_files(big_emails)
        logger.warning(f"Split up {len(big_emails)} into {len(new_emails)} smaller emails...")
        return new_emails


def document_cls(doc: Document) -> Type[Document]:
    """Find the appropriate `Document` subclass for this file based on the contents."""
    if doc.length == 0:
        return Document
    elif doc.file_info.is_eml_file:
        return DropsiteEmail
    elif doc.text[0] == '{':
        return JsonFile
    elif doc.file_id in IMESSAGE_PDF_IDS:
        return MessengerLogPdf
    elif doc.is_email:
        return Email
    elif doc.file_info.is_doj_file:
        return DojFile
    elif MSG_REGEX.search(doc.text[0:5000]):  # Limit search area to avoid pointless scans of huge files
        return MessengerLog
    else:
        return OtherFile


# Find author/recipients that we configured but that have no Entity so we can suppress warnings about them
# TODO: this populates the var in emailers.py which sucks but there's circular import issues
for cfg in CONFIGS_BY_ID.values():
    for name in cfg.names:
        if name in ENTITIES_DICT or name in CONFIGURED_NON_ENTITIES:
            continue

        CONFIGURED_NON_ENTITIES[name] = Entity(name, is_emailer=False, is_interesting=False)
        log_msg = f"Configured name has no Entity object: {quote(name)}"

        if ',' in name or isinstance(cfg, EmailCfg) and (cfg.has_uninteresting_ccs or cfg.has_uninteresting_bccs):
            cfg._debug_log(log_msg)
        else:
            cfg._log(log_msg)
