from sqlalchemy.orm import Session
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import func
from backend.app.db.models.addresses import PostcodeSearch
from utils.logger import setup_logger

logger = setup_logger()


def _get_associated_records(results, uprn, uprn_key="UPRN"):
    matched_record = []
    for x in results:
        if "DPA" in x:
            if x["DPA"].get(uprn_key) == str(uprn):
                matched_record.append(x["DPA"])
        else:
            if x["LPI"].get(uprn_key) == str(uprn):
                matched_record.append(x["LPI"])

    return matched_record


def get_associated_uprns(session: Session, postcode: str, uprn: str):
    """
    Given a postcode and UPRN, for a remote assessment, fetch all associated UPRNs, based
    on parent UPRN. This will be properties in the same building

    Parent UPRN is referenced in the following docs:
    https://static.geoplace.co.uk/downloads/GeoPlace-Data-Entry-Conventions-Best-Practice-for-Addresses.pdf

    :param session: The database session
    :param postcode: The postcode string to search for
    :param uprn: The UPRN string to match
    :return: The matching PostcodeSearch record, or None if not found
    """
    try:

        record = (
            session.query(PostcodeSearch)
            .filter(func.upper(PostcodeSearch.postcode) == postcode)
            .first()
        )
        if not record:
            # No record found for this postcode
            return []

        matched_record = _get_associated_records(results=record.result_data["results"], uprn=uprn)

        if len(matched_record) != 1:
            logger.error("Something went wrong, about to return nothing")
            return []

        if not matched_record[0].get("PARENT_UPRN"):
            logger.info("No parent UPRN found, cannot get associated records")
            return []

        associated_records = _get_associated_records(
            results=record.result_data["results"], uprn=matched_record[0]["PARENT_UPRN"], uprn_key="PARENT_UPRN"
        )
        # We now fetch all UPRNS with the same parent UPRN
        associated_uprns = [int(x["UPRN"]) for x in associated_records if x["UPRN"] != str(uprn)]

        return associated_uprns

    except SQLAlchemyError as e:
        session.rollback()
        raise e