#! /usr/bin/env python
"""CoLRev pdf_get operation: Get PDF documents."""
from __future__ import annotations

import shutil
from glob import glob
from multiprocessing.pool import ThreadPool as Pool
from pathlib import Path

import colrev.exceptions as colrev_exceptions
import colrev.process.operation
import colrev.record.record_pdf
from colrev.constants import Colors
from colrev.constants import EndpointType
from colrev.constants import Fields
from colrev.constants import OperationsType
from colrev.constants import PDFPathType
from colrev.constants import RecordState
from colrev.writer.write_utils import write_file


class PDFGet(colrev.process.operation.Operation):
    """Get the PDFs"""

    to_retrieve: int
    retrieved: int
    not_retrieved: int

    type = OperationsType.pdf_get

    def __init__(
        self,
        *,
        review_manager: colrev.review_manager.ReviewManager,
        notify_state_transition_operation: bool = True,
    ) -> None:
        super().__init__(
            review_manager=review_manager,
            operations_type=self.type,
            notify_state_transition_operation=notify_state_transition_operation,
        )

        self.package_manager = self.review_manager.get_package_manager()

        pdf_dir = self.review_manager.paths.pdf
        pdf_dir.mkdir(exist_ok=True, parents=True)

        self.pdf_qm = self.review_manager.get_pdf_qm()

        self.filepath_directory_pattern = ""
        pdf_endpoints = [
            s
            for s in self.review_manager.settings.sources
            if s.endpoint == "colrev.files_dir"
        ]
        if pdf_endpoints:
            self.filepath_directory_pattern = (
                pdf_endpoints[0].search_parameters["scope"].get("subdir_pattern", {})
            )

    def copy_pdfs_to_repo(self) -> None:
        """Copy the PDFs to the repository"""
        self.review_manager.logger.info("Copy PDFs to dir")
        records = self.review_manager.dataset.load_records_dict()

        for record_dict in records.values():
            if Fields.FILE not in record_dict:
                continue
            fpath = Path(record_dict[Fields.FILE])
            new_fpath = fpath.absolute()
            if fpath.is_symlink():
                linked_file = fpath.resolve()
                if linked_file.is_file():
                    fpath.unlink()
                    shutil.copyfile(linked_file, new_fpath)
                    self.review_manager.logger.info(
                        f" {Colors.GREEN}copied PDF for {record_dict[Fields.ID]} {Colors.END}"
                    )

            elif new_fpath.is_file() and self.review_manager.verbose_mode:
                self.review_manager.logger.info(
                    f"No need to copy PDF - already exits ({record_dict[Fields.ID]})"
                )

    def link_pdf(
        self, record: colrev.record.record.Record
    ) -> colrev.record.record.Record:
        """Link the PDF in its record (should be {ID}.pdf)"""

        pdf_filepath = self.review_manager.paths.PDF_DIR / Path(
            f"{record.data['ID']}.pdf"
        )
        if (self.review_manager.path / pdf_filepath).is_file() and str(
            pdf_filepath
        ) != record.data.get(Fields.FILE, "NA"):
            record.update_field(
                key=Fields.FILE, value=str(pdf_filepath), source="link_pdf"
            )
            self.import_pdf(record)
            if RecordState.rev_prescreen_included == record.data[Fields.STATUS]:
                record.set_status(RecordState.pdf_imported)

        return record

    def get_target_filepath(self, record: colrev.record.record.Record) -> Path:
        """Get the target filepath for a PDF"""

        target_filepath = self.review_manager.paths.PDF_DIR / Path(
            f"{record.data['ID']}.pdf"
        )
        if self.filepath_directory_pattern == Fields.YEAR:
            target_filepath = self.review_manager.paths.PDF_DIR / Path(
                f"{record.data.get('year', 'no_year')}/{record.data['ID']}.pdf"
            )

        elif self.filepath_directory_pattern == "volume_number":
            if Fields.VOLUME in record.data and Fields.NUMBER in record.data:
                target_filepath = self.review_manager.paths.PDF_DIR / Path(
                    f"{record.data['volume']}/{record.data['number']}/{record.data['ID']}.pdf"
                )

            if Fields.VOLUME in record.data and Fields.NUMBER not in record.data:
                target_filepath = self.review_manager.paths.PDF_DIR / Path(
                    f"{record.data['volume']}/{record.data['ID']}.pdf"
                )

        return target_filepath

    def import_pdf(self, record: colrev.record.record.Record) -> None:
        """Import a file (PDF) and copy/symlink it"""
        # self.review_manager.pdf_dir.mkdir(exist_ok=True)
        # new_fp = self.review_manager.PDF_DIR_RELATIVE / Path(record.data[Fields.ID] + ".pdf").name
        new_fp = self.get_target_filepath(record)
        original_fp = Path(record.data[Fields.FILE])

        if new_fp != original_fp and not new_fp.is_file():
            new_fp.parents[0].mkdir(exist_ok=True, parents=True)
            if (
                PDFPathType.symlink
                == self.review_manager.settings.pdf_get.pdf_path_type
            ):
                new_fp.symlink_to(original_fp)
            elif PDFPathType.copy == self.review_manager.settings.pdf_get.pdf_path_type:
                shutil.copyfile(original_fp, new_fp.resolve())
            # Note : else: leave absolute paths

        record.data[Fields.FILE] = str(new_fp)

    def _log_infos(self, record: colrev.record.record.Record) -> None:
        if Fields.FILE not in record.data:
            if (
                not self.review_manager.settings.pdf_get.pdf_required_for_screen_and_synthesis
            ):
                return
            self.review_manager.logger.info(
                f" {Colors.ORANGE}{record.data['ID']}".ljust(46)
                + f"rev_prescreen_included → pdf_needs_manual_retrieval{Colors.END}"
            )
            return

        if RecordState.pdf_prepared == record.data[Fields.STATUS]:
            self.review_manager.logger.info(
                f" {Colors.GREEN}{record.data['ID']}".ljust(46)
                + f"rev_prescreen_included → pdf_prepared{Colors.END}"
            )
        elif RecordState.pdf_needs_manual_preparation == record.data[Fields.STATUS]:
            self.review_manager.logger.info(
                f" {Colors.ORANGE}{record.data['ID']}".ljust(46)
                + f"rev_prescreen_included → pdf_needs_manual_preparation{Colors.END}"
            )

    # Note : no named arguments (multiprocessing)
    def get_pdf(self, item: dict) -> dict:
        """Get PDFs (based on the package endpoints in the settings)"""

        record_dict = item["record"]

        if record_dict[Fields.STATUS] not in [
            RecordState.rev_prescreen_included,
            RecordState.pdf_needs_manual_retrieval,
        ]:
            if Fields.FILE in record_dict:
                record = colrev.record.record_pdf.PDFRecord(
                    record_dict, path=self.review_manager.path
                )
                record.remove_field(key=Fields.FILE)
                return record.get_data()
            return record_dict

        record = colrev.record.record_pdf.PDFRecord(
            record_dict, path=self.review_manager.path
        )

        for (
            pdf_get_package_endpoint
        ) in self.review_manager.settings.pdf_get.pdf_get_package_endpoints:

            pdf_get_class = self.package_manager.get_package_endpoint_class(
                package_type=EndpointType.pdf_get,
                package_identifier=pdf_get_package_endpoint["endpoint"],
            )
            endpoint = pdf_get_class(
                pdf_get_operation=self, settings=pdf_get_package_endpoint
            )

            endpoint.get_pdf(record)  # type: ignore

            if Fields.FILE in record.data:
                self.review_manager.report_logger.info(
                    f"{endpoint.settings.endpoint}"  # type: ignore
                    f"({record.data[Fields.ID]}): retrieved .../"
                    f"{Path(record.data[Fields.FILE]).name}"
                )
                break

        if Fields.FILE in record.data:
            record.run_pdf_quality_model(self.pdf_qm, set_prepared=True)
        else:
            record.set_status(RecordState.pdf_needs_manual_retrieval)

        self._log_infos(record)

        return record.get_data()

    def _fix_broken_symlinks(self) -> None:

        pdf_dir = self.review_manager.paths.pdf
        broken_symlinks = []
        for pdf_candidate in list(pdf_dir.glob("**/*.pdf")):
            relative_path = pdf_candidate.relative_to(self.review_manager.path)
            if not relative_path.is_file():
                if pdf_candidate.is_symlink():
                    broken_symlinks.append(pdf_candidate)
                continue

        if broken_symlinks:
            print("To fix broken symlinks:")
            old_path = input("Enter the old path: ")
            new_path = input("Enter the new path: ")
            for broken_symlink in broken_symlinks:
                new_file = str(broken_symlink.resolve()).replace(old_path, new_path)
                print(f"Fix {broken_symlink}")
                broken_symlink.unlink()
                broken_symlink.symlink_to(new_file)

    def _relink_pdfs_in_source(self, source: colrev.settings.SearchSource) -> None:

        # pylint: disable=too-many-locals

        self.review_manager.logger.info(
            "Checking PDFs in same directory to reassign when "
            f"the cpid is identical {source.filename}"
        )

        pdf_dir = self.review_manager.paths.pdf
        pdf_candidates = {}
        for pdf_candidate in list(pdf_dir.glob("**/*.pdf")):
            colrev_pdf_id = colrev.record.record_pdf.PDFRecord.get_colrev_pdf_id(
                pdf_candidate
            )
            relative_path = pdf_candidate.relative_to(self.review_manager.path)
            pdf_candidates[relative_path] = colrev_pdf_id

        source_records_dict = colrev.loader.load_utils.load(
            filename=source.filename,
            logger=self.review_manager.logger,
        )
        source_records = list(source_records_dict.values())
        corresponding_origin = str(source.filename)
        records = self.review_manager.dataset.load_records_dict()
        for record in records.values():
            if Fields.FILE not in record:
                continue

            # Note: we check the source_records based on the cpids
            # in the record because cpids are not stored in the source_record
            # (pdf hashes may change after import/preparation)
            source_rec = {}
            if corresponding_origin != "":
                source_origin_l = [
                    o for o in record[Fields.ORIGIN] if corresponding_origin in o
                ]
                if len(source_origin_l) == 1:
                    source_origin = source_origin_l[0]
                    source_origin = source_origin.replace(
                        f"{corresponding_origin}/", ""
                    )
                    source_rec_l = [
                        s for s in source_records if s[Fields.ID] == source_origin
                    ]
                    if len(source_rec_l) == 1:
                        source_rec = source_rec_l[0]

            if source_rec:
                if (
                    self.review_manager.path / Path(record[Fields.FILE])
                ).is_file() and (
                    self.review_manager.path / Path(source_rec[Fields.FILE])
                ).is_file():
                    continue
            else:
                if (self.review_manager.path / Path(record[Fields.FILE])).is_file():
                    continue

            self.review_manager.logger.info(record[Fields.ID])

            for pdf_candidate, cpid in pdf_candidates.items():
                if record.get("colrev_pdf_id", "") == cpid:
                    record[Fields.FILE] = str(pdf_candidate)
                    source_rec[Fields.FILE] = str(pdf_candidate)

                    self.review_manager.logger.info(
                        f"Found and linked PDF: {pdf_candidate}"
                    )
                    break

        if len(source_records) > 0:
            source_records_dict = {r[Fields.ID]: r for r in source_records}
            write_file(records_dict=source_records_dict, filename=source.filename)

        self.review_manager.dataset.save_records_dict(records)
        self.review_manager.dataset.add_changes(source.filename)

    def relink_pdfs(self) -> None:
        """Relink record files to the corresponding PDFs (if available)"""

        self._fix_broken_symlinks()

        sources = [
            s
            for s in self.review_manager.settings.sources
            if s.endpoint == "colrev.files_dir" and s.filename.is_file()
        ]
        for source in sources:
            self._relink_pdfs_in_source(source)

        self.review_manager.dataset.create_commit(msg="Relink PDFs")

    def check_existing_unlinked_pdfs(
        self,
        records: dict,
    ) -> dict:
        """Check for PDFs that are in the pdfs directory but not linked in the record file"""

        linked_pdfs = [
            str(Path(x[Fields.FILE]).resolve())
            for x in records.values()
            if Fields.FILE in x
        ]
        pdf_dir = self.review_manager.paths.pdf
        pdf_files = glob(str(pdf_dir) + "/**.pdf", recursive=True)
        unlinked_pdfs = [
            Path(x)
            for x in pdf_files
            if str(Path(x).resolve()) not in linked_pdfs
            and not any(kw in x for kw in ["_with_lp.pdf", "_with_cp.pdf", "_ocr.pdf"])
        ]

        if len(unlinked_pdfs) == 0:
            return records

        grobid_service = self.review_manager.get_grobid_service()
        grobid_service.start()
        self.review_manager.logger.info("Check unlinked PDFs")
        for file in unlinked_pdfs:
            msg = f"Check unlinked PDF: {file.relative_to(self.review_manager.path)}"
            self.review_manager.logger.info(msg)
            if file.stem not in records.keys():
                tei = self.review_manager.get_tei(pdf_path=file)
                pdf_record = tei.get_metadata()

                if "error" in pdf_record:
                    continue

                max_similarity = 0.0
                max_sim_record = None
                for record in records.values():
                    sim = colrev.record.record_pdf.PDFRecord.get_record_similarity(
                        colrev.record.record_pdf.PDFRecord(
                            pdf_record, path=self.review_manager.path
                        ),
                        colrev.record.record_pdf.PDFRecord(
                            record.copy(), path=self.review_manager.path
                        ),
                    )
                    if sim > max_similarity:
                        max_similarity = sim
                        max_sim_record = record
                if max_sim_record:
                    if max_similarity > 0.5:
                        if RecordState.pdf_prepared == max_sim_record[Fields.STATUS]:
                            continue

                        record = colrev.record.record_pdf.PDFRecord(
                            max_sim_record, path=self.review_manager.path
                        )
                        record.update_field(
                            key=Fields.FILE,
                            value=str(file),
                            source="linking-available-files",
                        )
                        self.import_pdf(record)
                        if (
                            RecordState.rev_prescreen_included
                            == record.data[Fields.STATUS]
                        ):
                            record.set_status(RecordState.pdf_imported)

                        self.review_manager.report_logger.info(
                            "linked unlinked pdf:" f" {file.name}"
                        )
                        self.review_manager.logger.info(
                            "linked unlinked pdf:" f" {file.name}"
                        )
                        # max_sim_record = \
                        #     pdf_prep.validate_pdf_metadata(max_sim_record)
                        # colrev_status = max_sim_record['colrev_status']
                        # if RecordState.pdf_needs_manual_preparation == colrev_status:
                        #     # revert?
            else:
                record = records[file.stem]
                self.link_pdf(
                    colrev.record.record_pdf.PDFRecord(
                        record, path=self.review_manager.path
                    )
                )

        self.review_manager.dataset.save_records_dict(records)

        return records

    def _rename_pdf(
        self,
        *,
        record_dict: dict,
        file: Path,
        new_filename: Path,
        pdfs_search_file: Path,
    ) -> None:
        record_dict[Fields.FILE] = new_filename
        if pdfs_search_file.is_file():
            colrev.env.utils.inplace_change(
                filename=pdfs_search_file,
                old_string="{" + str(file) + "}",
                new_string="{" + str(new_filename) + "}",
            )

        if not file.is_file():
            corrected_path = Path(str(file).replace("  ", " "))
            if corrected_path.is_file():
                file = corrected_path

        if file.is_file():
            shutil.move(str(file), str(new_filename))
        elif file.is_symlink():
            shutil.move(str(file), str(new_filename))

        record_dict[Fields.FILE] = str(new_filename)
        self.review_manager.logger.info(f"rename {file.name} > {new_filename}")
        if RecordState.rev_prescreen_included == record_dict[Fields.STATUS]:
            record = colrev.record.record_pdf.PDFRecord(
                record_dict, path=self.review_manager.path
            )
            record.set_status(RecordState.pdf_imported)

    def rename_pdfs(self) -> None:
        """Rename the PDFs"""

        self.review_manager.logger.info("Rename PDFs")

        records = self.review_manager.dataset.load_records_dict()

        # We may use other pdfs_search_files from the sources:
        # review_manager.settings.sources
        pdfs_search_file = Path("data/search/pdfs.bib")

        for record_dict in records.values():
            if Fields.FILE not in record_dict:
                continue
            if record_dict[Fields.STATUS] not in RecordState.get_post_x_states(
                state=RecordState.md_processed
            ):
                continue

            file = Path(record_dict[Fields.FILE])
            new_filename = file.parents[0] / Path(f"{record_dict['ID']}{file.suffix}")
            # Possible option: move to top (pdfs) directory:
            # new_filename = self.review_manager.PDF_DIR_RELATIVE / Path(
            #     f"{record['ID']}.pdf"
            # )
            if str(file) == str(new_filename):
                continue

            self._rename_pdf(
                record_dict=record_dict,
                file=file,
                new_filename=new_filename,
                pdfs_search_file=pdfs_search_file,
            )

        self.review_manager.dataset.save_records_dict(records)

        if pdfs_search_file.is_file():
            self.review_manager.dataset.add_changes(pdfs_search_file)

    def _get_data(self) -> dict:
        # pylint: disable=duplicate-code

        records_headers = self.review_manager.dataset.load_records_dict(
            header_only=True
        )
        record_header_list = list(records_headers.values())

        nr_tasks = len(
            [
                x
                for x in record_header_list
                if x[Fields.STATUS]
                in [
                    RecordState.pdf_needs_manual_retrieval,
                    RecordState.rev_prescreen_included,
                ]
            ]
        )

        items = self.review_manager.dataset.read_next_record(
            conditions=[
                {Fields.STATUS: RecordState.rev_prescreen_included},
                {Fields.STATUS: RecordState.pdf_needs_manual_retrieval},
            ],
        )

        self.to_retrieve = nr_tasks

        pdf_get_data = {
            "nr_tasks": nr_tasks,
            "items": [{"record": item} for item in items],
        }

        return pdf_get_data

    def _print_stats(self, retrieved_record_list: list) -> None:
        self.retrieved = len([r for r in retrieved_record_list if Fields.FILE in r])

        self.not_retrieved = self.to_retrieve - self.retrieved

        retrieved_string = "Overall pdf_imported".ljust(34)
        if self.retrieved == 0:
            retrieved_string += f"{self.retrieved}".rjust(6, " ")
            retrieved_string += " PDFs"
        elif self.retrieved == 1:
            retrieved_string += f"{Colors.GREEN}"
            retrieved_string += f"{self.retrieved}".rjust(6, " ")
            retrieved_string += f"{Colors.END} PDF"
        else:
            retrieved_string += f"{Colors.GREEN}"
            retrieved_string += f"{self.retrieved}".rjust(6, " ")
            retrieved_string += f"{Colors.END} PDFs"

        not_retrieved_string = "Overall pdf_needs_manual_retrieval".ljust(34)
        if self.not_retrieved == 0:
            not_retrieved_string += f"{self.not_retrieved}".rjust(6, " ")
            not_retrieved_string += " PDFs"
        elif self.not_retrieved == 1:
            not_retrieved_string += f"{Colors.ORANGE}"
            not_retrieved_string += f"{self.not_retrieved}".rjust(6, " ")
            not_retrieved_string += f"{Colors.END} PDF"
        else:
            not_retrieved_string += f"{Colors.ORANGE}"
            not_retrieved_string += f"{self.not_retrieved}".rjust(6, " ")
            not_retrieved_string += f"{Colors.END} PDFs"

        self.review_manager.logger.info(retrieved_string)
        self.review_manager.logger.info(not_retrieved_string)

    def _set_status_if_pdf_linked(self, records: dict) -> dict:
        for record_dict in records.values():
            if record_dict[Fields.STATUS] in [
                RecordState.rev_prescreen_included,
                RecordState.pdf_needs_manual_retrieval,
            ]:
                record = colrev.record.record_pdf.PDFRecord(
                    record_dict, path=self.review_manager.path
                )
                if Fields.FILE in record_dict:
                    if any(
                        Path(fpath).is_file()
                        for fpath in record.data[Fields.FILE].split(";")
                    ):
                        if (
                            RecordState.rev_prescreen_included
                            == record.data[Fields.STATUS]
                        ):
                            record.set_status(RecordState.pdf_imported)
                    else:
                        self.review_manager.logger.warning(
                            "Remove non-existent file link "
                            f"({record_dict[Fields.ID]}: {record_dict[Fields.FILE]}"
                        )
                        record.remove_field(key=Fields.FILE)
        self.review_manager.dataset.save_records_dict(records)

        return records

    def setup_custom_script(self) -> None:
        """Setup a custom pfd-get script"""

        filedata = colrev.env.utils.get_package_file_content(
            module="colrev.ops",
            filename=Path("custom_scripts/custom_pdf_get_script.py"),
        )
        if filedata:
            with open("custom_pdf_get_script.py", "w", encoding="utf-8") as file:
                file.write(filedata.decode("utf-8"))

        self.review_manager.dataset.add_changes(Path("custom_pdf_get_script.py"))

        self.review_manager.settings.pdf_get.pdf_get_man_package_endpoints.append(
            {"endpoint": "custom_pdf_get_script"}
        )

        self.review_manager.save_settings()

    @colrev.process.operation.Operation.decorate()
    def main(self) -> None:
        """Get PDFs (main entrypoint)"""

        if (
            self.review_manager.in_ci_environment()
            and not self.review_manager.in_test_environment()
        ):
            raise colrev_exceptions.ServiceNotAvailableException(
                dep="colrev pdf-prep",
                detailed_trace="pdf-prep not available in ci environment",
            )

        self.review_manager.logger.info("Get PDFs")
        self.review_manager.logger.info(
            "Get PDFs of prescreen-included records from local and remote sources."
        )
        self.review_manager.logger.info(
            "PDFs are stored in the directory data/pdfs "
            f"({Colors.ORANGE}colrev pdfs --dir{Colors.END})"
        )
        self.review_manager.logger.info(
            "See https://colrev-environment.github.io/colrev/manual/pdf_retrieval/pdf_get.html"
        )

        records = self.review_manager.dataset.load_records_dict()
        records = self._set_status_if_pdf_linked(records)
        records = self.check_existing_unlinked_pdfs(records)

        pdf_get_data = self._get_data()

        if pdf_get_data["nr_tasks"] == 0:
            self.review_manager.logger.info("No additional pdfs to retrieve")
        else:
            self.review_manager.logger.info(
                "PDFs to get".ljust(38) + f'{pdf_get_data["nr_tasks"]} PDFs'
            )

            pool = Pool(4)
            retrieved_record_list = pool.map(self.get_pdf, pdf_get_data["items"])
            pool.close()
            pool.join()

            self.review_manager.dataset.save_records_dict(
                {r[Fields.ID]: r for r in retrieved_record_list}, partial=True
            )

            self._print_stats(retrieved_record_list)

        # Note: rename should be after copy.
        # Note : do not pass records as an argument.
        if self.review_manager.settings.pdf_get.rename_pdfs:
            self.rename_pdfs()

        self.review_manager.dataset.create_commit(msg="PDFs: get and prepare")
        self.review_manager.logger.info(
            f"{Colors.GREEN}Completed pdf-get operation{Colors.END}"
        )
