"""Script to rename fields and transform values prior to submission to Loculus"""

# Needs to be configurable via yaml file
# Start off with a simple mapping
# Add transformations that can be applied to certain fields
# Like separation of country into country and division

import csv
import hashlib
import json
import logging
from dataclasses import dataclass

import click
import orjsonl
import pandas as pd
import yaml

logger = logging.getLogger(__name__)
logging.basicConfig(
    encoding="utf-8",
    level=logging.DEBUG,
    format="%(asctime)s %(levelname)8s (%(filename)20s:%(lineno)4d) - %(message)s ",
    datefmt="%H:%M:%S",
)

FastaIdField = str


@dataclass
class Config:
    compound_country_field: str
    fasta_id_field: str
    rename: dict[str, str]
    keep: list[str]
    segmented: bool


def resolve_host_information(record: dict[str, str]) -> dict[str, str]:
    """Create a new host field and populate it from hostTaxonId or
    hostNameScientific (falling back to the empty string) to be consistent
    with how direct submissions specify the host organism. Any existing
    hostTaxonId, hostNameScientific, and hostNameCommon fields on the
    record will be removed.

    This should be done after computing the hash for a record to not trigger
    revisions for all INSDC data
    """
    host = record.get("hostTaxonId") or record.get("hostNameScientific")
    record.pop("hostTaxonId", None)
    record.pop("hostNameScientific", None)
    record.pop("hostNameCommon", None)
    record["host"] = host

    return record


@click.command()
@click.option("--config-file", required=True, type=click.Path(exists=True))
@click.option("--input", required=True, type=click.Path(exists=True))
@click.option("--segments", required=False, type=click.Path())
@click.option("--sequence-hashes-file", required=True, type=click.Path(exists=True))
@click.option("--output", required=True, type=click.Path())
@click.option(
    "--log-level",
    default="INFO",
    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]),
)
def main(
    config_file: str,
    input: str,
    segments: str | None,
    sequence_hashes_file: str,
    output: str,
    log_level: str,
) -> None:
    logger.setLevel(log_level)

    with open(config_file, encoding="utf-8") as file:
        full_config = yaml.safe_load(file)
        relevant_config = {key: full_config[key] for key in Config.__annotations__}
        config = Config(**relevant_config)
    logger.debug(config)

    logger.info(f"Reading metadata from {input}")
    df = pd.read_csv(
        input, sep="\t", dtype=str, keep_default_na=False, quoting=csv.QUOTE_NONE, escapechar="\\"
    )
    metadata: list[dict[str, str]] = df.to_dict(orient="records")

    sequence_hashes: dict[FastaIdField, str] = {
        record["id"]: record["hash"] for record in orjsonl.load(sequence_hashes_file)
    }

    if segments:
        segments_dict: dict[FastaIdField, dict[str, str]] = {}
        segment_df = pd.read_csv(segments, sep="\t")
        segmented_fields = list(segment_df.columns)

        rows_as_dicts = segment_df.to_dict(orient="records")

        for row in rows_as_dicts:
            segments_dict[row["seqName"]] = row

    for record in metadata:
        # Transform the metadata
        try:
            record["division"] = record[config.compound_country_field].split(":", 1)[1].strip()
        except IndexError:
            record["division"] = ""
        record["country"] = record[config.compound_country_field].split(":", 1)[0].strip()
        record["id"] = record[config.fasta_id_field]
        record["insdcAccessionBase"] = record[config.fasta_id_field].split(".", 1)[0]
        record["insdcVersion"] = record[config.fasta_id_field].split(".", 1)[1]
        if segments:
            results_dic = segments_dict.get(record[config.fasta_id_field], {})
            for key in segmented_fields:
                record[key] = results_dic.get(key, "")

    # Get rid of all records without segment
    if config.segmented:
        metadata = [record for record in metadata if record["segment"]]
        missing_a_segment = [
            record["insdcAccessionBase"] for record in metadata if not record["segment"]
        ]
        if missing_a_segment:
            logger.info(
                f"Missing segment for {len(missing_a_segment)} records: {', '.join(missing_a_segment)}"
            )

    for record in metadata:
        for from_key, to_key in config.rename.items():
            # segment is a required field for the ingest grouping scripts
            # Keep segment field if config.segmented even if it is specified to be renamed
            val = (
                record.get(from_key)
                if (from_key == "segment" and config.segmented)
                else record.pop(from_key)
            )
            record[to_key] = val

    keys_to_keep = set(config.rename.values()) | set(config.keep)
    if config.segmented:
        keys_to_keep.add("segment")
    for record in metadata:
        for key in list(record.keys()):
            if key not in keys_to_keep:
                record.pop(key)

    # Calculate overall hash of metadata + sequence
    for record in metadata:
        fasta_id_field = config.fasta_id_field
        if config.fasta_id_field in config.rename:
            fasta_id_field = config.rename[config.fasta_id_field]
        sequence_hash = sequence_hashes.get(record[fasta_id_field], "")
        if not sequence_hash:
            msg = f"No hash found for {record[config.fasta_id_field]}"
            raise ValueError(msg)

        # Hash of all metadata fields should be the same if
        # 1. field is not in keys_to_keep and
        # 2. field is in keys_to_keep but is "" or None
        filtered_record = {k: str(v) for k, v in record.items() if v is not None and str(v)}

        # rename "id" to "submissionId" for back-compatibility with old hashes
        filtered_record["submissionId"] = filtered_record.pop("id")

        metadata_dump = json.dumps(filtered_record, sort_keys=True)
        prehash = metadata_dump + sequence_hash

        record["hash"] = hashlib.md5(prehash.encode(), usedforsecurity=False).hexdigest()

        # for segmented organisms, this has to happen in `heuristic_group_segments.py`
        # and `override_group_segments.py`
        if not config.segmented:
            record = resolve_host_information(record)

        orjsonl.append(output, {"id": record[fasta_id_field], "metadata": record})

    logger.info(f"Saved metadata for {len(metadata)} sequences")


if __name__ == "__main__":
    main()
