from __future__ import annotations
import base64

from dataclasses import dataclass
import re
import time
from typing import Optional, Union

from ..allspice import AllSpice
from ..apiobject import Content, Ref, Repository
from ..exceptions import NotYetGeneratedException


@dataclass
class SchematicComponent:
    description: str
    manufacturer: str
    part_number: str
    designator: str


@dataclass
class PcbComponent:
    designator: str
    schematic_link: str


@dataclass
class BomEntry:
    description: str
    manufacturer: str
    part_number: str
    designators: list[str]
    quantity: int


Bom = list[BomEntry]


@dataclass
class AttributesMapping:
    """
    Define how we map from attributes in the Altium SchDoc file to the fields of
    the BOM.

    For each field, we define a list of attribute names that we will check in
    the Altium SchDoc file. The *first* one found defined will be used.

    :param description: A list of attribute names in the Altium SchDoc file
        that contain the description of the part.
    :param manufacturer: A list of attribute names in the Altium SchDoc file
        that contain the manufacturer of the part.
    :param part_number: A list of attribute names in the Altium SchDoc file
        that contain the part number of the part.
    :param designator: A list of attribute names in the Altium SchDoc file
        that contain the designator of the part.
    """

    description: list[str]
    manufacturer: list[str]
    part_number: list[str]
    designator: list[str]

    @staticmethod
    def from_dict(dictionary: dict[str, list[str]]) -> AttributesMapping:
        """
        Create an AttributesMapping from a dictionary.

        Example Input::

            {
                "description": ["Description"],
                "manufacturer": ["Manufacturer", "MANUFACTURER"],
                "part_number": ["Part Number"],
                "designator": ["Designator"],
            }
        """
        return AttributesMapping(
            description=dictionary["description"],
            manufacturer=dictionary["manufacturer"],
            part_number=dictionary["part_number"],
            designator=dictionary["designator"],
        )


def generate_bom_for_altium(
    allspice_client: AllSpice,
    repository: Repository,
    prjpcb_file: Union[Content, str],
    pcbdoc_file: Union[Content, str],
    attributes_mapping: AttributesMapping,
    ref: Ref = "main",
) -> Bom:
    """
    Generate a BOM for an Altium project.

    :param allspice_client: The AllSpice client to use.
    :param repository: The repository to generate the BOM for.
    :param prjpcb_file: The Altium project file. This can be a Content object
        returned by the AllSpice API, or a string containing the path to the
        file in the repo.
    :param pcbdoc_file: The Altium PCB document file. This can be a Content
        object returned by the AllSpice API, or a string containing the path to
        the file in the repo.
    :param attributes_mapping: A mapping of Altium attributes to BOM entry
        attributes. See the documentation for AttributesMapping for more
        information.
    :param ref: The ref, i.e. branch, commit or git ref from which to take the
        project files. Defaults to "main".
    :return: A list of BOM entries.
    """

    allspice_client.logger.info(
        f"Generating BOM for {repository.name=} on {ref=} using {attributes_mapping=}"
    )
    allspice_client.logger.info(f"Fetching {prjpcb_file=} and {pcbdoc_file=}")

    prjpcb_file = _get_file_content(repository, prjpcb_file, ref)
    schdoc_files_in_proj = {x for x in _extract_schdoc_list_from_prjpcb(prjpcb_file)}

    allspice_client.logger.info("Found %d SchDoc files", len(schdoc_files_in_proj))

    schdoc_components = _extract_all_schdoc_components(
        repository,
        ref,
        schdoc_files_in_proj,
        attributes_mapping,
    )
    pcbdoc_components = _extract_all_pcbdoc_components(repository, ref, pcbdoc_file)

    ungrouped_bom_entries = _combine_components_to_ungrouped_bom(
        schdoc_components,
        pcbdoc_components,
    )

    return _group_bom_entries(ungrouped_bom_entries)


def _find_first_matching_key(
    alternatives: Union[list[str], str],
    attributes: dict,
) -> Optional[str]:
    """
    Search for a series of alternative keys in a dictionary, and return the
    value of the first one found.
    """

    if isinstance(alternatives, str):
        alternatives = [alternatives]

    for alternative in alternatives:
        if alternative in attributes:
            return attributes[alternative]["text"]

    return None


def _get_file_content(repo, file_path, ref) -> str:
    """
    Get the content of a file in a repo on a branch.
    """

    files_in_repo = repo.get_git_content(ref=ref)
    file = next((x for x in files_in_repo if x.path == file_path), None)
    if not file:
        raise ValueError(
            f"File {file_path} not found in repo {repo.name} at ref {ref.name}"
        )

    content = repo.get_file_content(file, ref=ref)
    return base64.b64decode(content).decode("utf-8")


def _extract_schdoc_list_from_prjpcb(prjpcb_file_content) -> list[str]:
    """
    Get a list of SchDoc files from a PrjPcb file.
    """

    pattern = re.compile(r"DocumentPath=(.*?SchDoc)\r\n")
    return [match.group(1) for match in pattern.finditer(prjpcb_file_content)]


def _schdoc_component_from_attributes(
    attributes: dict,
    mapper: AttributesMapping,
) -> SchematicComponent:
    """
    Make a SchDoc Component object out of the JSON for it in the SchDoc file.

    :param attributes: The attributes for the component in the SchDoc
    :param mapper: An AttributesMapping object, see the documentation for
        AttributesMapping for more information.
    :returns: A SchodocComponent object representing that component.
    """

    return SchematicComponent(
        description=_find_first_matching_key(mapper.description, attributes),
        designator=_find_first_matching_key(mapper.designator, attributes) or "",
        manufacturer=_find_first_matching_key(mapper.manufacturer, attributes),
        part_number=_find_first_matching_key(mapper.part_number, attributes),
    )


def _extract_components_from_schdoc(
    schdoc_file_content: dict,
    attributes_mapper: AttributesMapping,
) -> list[SchematicComponent]:
    """
    Extract all the components from a schdoc file. To see what attributes are
    available, print the schdoc_file_content variable.

    :param schdoc_file_content: The content of the SchDoc file. This should be a
    dictionary.
    """

    components_list = []

    for value in schdoc_file_content.values():
        if isinstance(value, dict):
            if value.get("type") == "Component":
                attributes = value["attributes"]
                components_list.append(
                    _schdoc_component_from_attributes(
                        attributes,
                        attributes_mapper,
                    )
                )

    return components_list


def _extract_all_schdoc_components(
    repository: Repository,
    ref: Ref,
    schdoc_files_in_proj: set[str],
    attributes_mapping: AttributesMapping,
) -> list[SchematicComponent]:
    """
    Fetch all the components from all the SchDoc files in the repo that are in
    the project.
    """

    files_in_repo = repository.get_git_content(ref=ref)
    all_components = []
    for repo_file in files_in_repo:
        if repo_file.name in schdoc_files_in_proj:
            # If the file is not yet generated, we'll retry a few times.
            retry_count = 0
            while True:
                retry_count += 1

                try:
                    schdoc_file_json = repository.get_generated_json(repo_file, ref=ref)
                    all_components.extend(
                        _extract_components_from_schdoc(
                            schdoc_file_json,
                            attributes_mapping,
                        )
                    )

                    break
                except NotYetGeneratedException:
                    if retry_count > 5:
                        break

                    # Wait a bit before retrying.
                    time.sleep(1)
                    continue

    return all_components


def _extract_all_pcbdoc_components(
    repository: Repository,
    ref: Ref,
    pcbdoc_file: str,
) -> list[PcbComponent]:
    """
    Extract all the components from a PcbDoc file in the repo.
    """

    pcb_json = repository.get_generated_json(pcbdoc_file, ref=ref)
    components = []

    component_instances = pcb_json["component_instances"]

    for component in component_instances.values():
        components.append(
            PcbComponent(
                designator=component["designator"],
                schematic_link=component["schematic_link"],
            )
        )

    return components


def _combine_components_to_ungrouped_bom(
    schdoc_components: list[SchematicComponent],
    pcbdoc_components: list[PcbComponent],
) -> list[BomEntry]:
    """
    Link the designators in the PcbDoc to a designator in the SchDocs, and
    generate a list of BOM entries. These BOM entries are not yet grouped by
    part number, and therefore have multiple rows for each part number.
    """

    schdoc_components_by_designator = {
        component.designator: component for component in schdoc_components
    }

    pcb_designators_for_schdoc_designators = {
        component.designator: [] for component in schdoc_components
    }
    orphan_pcb_components = []

    for pcbdoc_component in pcbdoc_components:
        if pcbdoc_component.schematic_link in schdoc_components_by_designator:
            schdoc_designator = pcbdoc_component.schematic_link
            pcb_designator = pcbdoc_component.designator

            pcb_designators_for_schdoc_designators[schdoc_designator].append(
                pcb_designator
            )
        else:
            orphan_pcb_components.append(pcbdoc_component)

    bom = []
    for (
        schdoc_designator,
        pcb_designators,
    ) in pcb_designators_for_schdoc_designators.items():
        schdoc_component = schdoc_components_by_designator[schdoc_designator]

        bom.append(
            BomEntry(
                description=schdoc_component.description,
                manufacturer=schdoc_component.manufacturer,
                part_number=schdoc_component.part_number,
                designators=pcb_designators,
                quantity=len(pcb_designators),
            )
        )
    for orphan_pcb_component in orphan_pcb_components:
        bom.append(
            BomEntry(
                description="",
                manufacturer="",
                part_number="",
                designators=[orphan_pcb_component.designator],
                quantity=1,
            )
        )

    return bom


def _group_bom_entries(bom_entries: list[BomEntry]) -> list[BomEntry]:
    """
    Group BOM Entries by the part number, combining the designators and
    quantities.
    """

    bom_entries_by_part_number = {}

    for bom_entry in bom_entries:
        if bom_entry.part_number != "":
            if bom_entry.part_number not in bom_entries_by_part_number:
                bom_entries_by_part_number[bom_entry.part_number] = bom_entry
            else:
                bom_entries_by_part_number[bom_entry.part_number].designators.extend(
                    bom_entry.designators
                )
                bom_entries_by_part_number[
                    bom_entry.part_number
                ].quantity += bom_entry.quantity

    return list(bom_entries_by_part_number.values())
