import logging
import os
from copy import copy
from pathlib import Path
from typing import Any, Dict

import numpy as np
import pandas as pd
import qlacref_authorities
from pandas import DataFrame
from qlacref_postcodes import Postcodes

logger = logging.getLogger(__name__)

# TODO security point. remove this line.
os.environ["QLACREF_PC_INSECURE"] = "True"

la_df = pd.DataFrame.from_records(qlacref_authorities.records)
logger.info("Loaded authorities")

# A bit of a hack, but will keep things working
if os.getenv("QLACREF_PC_KEY") is None:
    key_file = Path(__file__).parent.parent / ".qlacref/id_rsa.pub"
    if key_file.is_file():
        logger.warning(
            "Using repository key for pickle signature. "
            "To stay safe from tampering, set the key path in 'QLACREF_PC_KEY'"
        )
        os.environ["QLACREF_PC_KEY"] = str(key_file.absolute())

postcodes = Postcodes()
logger.info("Initialised Postcodes")


def create_datastore(data: Dict[str, Any], metadata: Dict[str, Any]):
    """
    Returns a dictionary with keys for
    - Every table name in config.py
    - Every table name again, with the suffix '_last', if last years data was uploaded.
    - A 'metadata' key, which links to a nested dictionary with
      - 'postcodes' - a postcodes csv, with columns "laua" (LA code), "oseast1m", "osnrth1m" (coordinates) and "pcd" (the postcode)
      - 'localAuthority' - the code of the local authority (long form)
      - 'collectionYear' - the collection year string (e.g. '2019/20)

    :param data: Dict of raw DataFrames by name (from config.py) together with the '_last' data.
    :param metadata:
    """
    logger.info(", ".join(data.keys()))
    data = copy(data)
    data["metadata"] = _process_metadata(metadata)
    if "Episodes" in data:
        data["Episodes"] = _add_postcode_derived_fields(
            data["Episodes"],
            metadata["localAuthority"],
        )
    if "Episodes_last" in data:
        data["Episodes_last"] = _add_postcode_derived_fields(
            data["Episodes_last"], metadata["localAuthority"]
        )
    # quick n dirty fix for weird postcode related columns showing up in episode tables
    for table_name in ["Episodes", "Episodes_last"]:
        if table_name in data:
            data[table_name] = data[table_name].drop(
                columns={"_14", "_15", "_16", "_17"} & set(data[table_name].columns)
            )

    names_and_lengths = ", ".join(f"{t}: {len(data[t])} rows" for t in data)
    logger.info(f"Datastore created -- {names_and_lengths}")
    return data


def _process_metadata(metadata):
    collection_year = int(metadata["collectionYear"][:4])
    metadata["collection_start"] = f"01/04/{collection_year - 1}"
    metadata["collection_end"] = f"31/03/{collection_year}"
    return metadata


def merge_postcodes(df: DataFrame, postcode_field: str) -> DataFrame:
    df[postcode_field] = df[postcode_field].str.upper()

    key_field = f"__{postcode_field}__first_letter"
    df[key_field] = df[postcode_field].str[0]
    pcs = df[key_field].dropna().unique()

    logger.info(f"Loading the following postcode letters: {pcs}")
    postcodes.load_postcodes(pcs)

    logger.info(f"Adding postcode abbreviations for {postcode_field}")
    pc_abbr_field = f"__{postcode_field}__abbr"
    df[pc_abbr_field] = df[postcode_field].str.replace(" ", "")

    df_merged = df[[pc_abbr_field]].merge(
        postcodes.dataframe, how="left", left_on=pc_abbr_field, right_on="pcd_abbr"
    )
    del df[key_field]
    del df[pc_abbr_field]

    return df_merged


def _add_postcode_derived_fields(episodes_df, local_authority):
    episodes_df = episodes_df.copy()

    home_details = merge_postcodes(episodes_df, "HOME_POST")
    pl_details = merge_postcodes(episodes_df, "PL_POST")

    # The indices remain the same post merge as the length of the dataframes doesn't change, so we can set directly.
    pl_details = pl_details.merge(
        la_df, how="left", left_on="laua", right_on="LTLA21CD"
    )
    episodes_df["PL_LA"] = pl_details["UTLA21CD"]

    logger.info(f"Adding IN/OUT")
    episodes_df["PL_LOCATION"] = "IN"
    episodes_df.loc[episodes_df["PL_LA"].ne(local_authority), "PL_LOCATION"] = "OUT"
    episodes_df.loc[pl_details["laua"].isna(), "PL_LOCATION"] = pd.NA

    # Add country codes for placements outside england
    for letter, code in {"S": "SCO", "N": "NIR", "W": "WAL"}.items():
        mask = (
            pl_details["laua"].str.upper().str.startswith(letter)
            & episodes_df["PL_POST"].notnull()
        )
        episodes_df.loc[mask, ["PL_LA", "PL_LOCATION"]] = [code, "OUT"]

    logger.info(f"Calculating distances")
    # This formula is taken straight from the guidance, to get miles between two postcodes
    episodes_df["PL_DISTANCE"] = (
        np.sqrt(
            (home_details["oseast1m"] - pl_details["oseast1m"]) ** 2
            + (home_details["osnrth1m"] - pl_details["osnrth1m"]) ** 2
        )
        / 1000
        / 1.6093
    )
    episodes_df["PL_DISTANCE"] = episodes_df["PL_DISTANCE"].round(decimals=1)

    logger.info(f"Completed postcode calculations")
    return episodes_df


def copy_datastore(data_store):
    """
    This is used for getting a shallow copy that is passed to the validation rules.
    """
    return {
        k: v.copy(deep=False) if k != "metadata" else v for k, v in data_store.items()
    }
