###############################################################################
# (c) Copyright 2026 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
import enum
import json
import logging
import random
from enum import StrEnum
from pathlib import Path
from typing import Annotated, Any, Literal, Self

import typer
import yaml
from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
from pydantic import BaseModel, Field, StringConstraints, model_validator
from rich.console import Console
from rich.logging import RichHandler

# Setup logging with rich
logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    handlers=[RichHandler(rich_tracebacks=True, show_path=False)],
)
logger = logging.getLogger(__name__)
console = Console()


class InputDataResolution(enum.StrEnum):
    Protocol = "InputDataByProtocol"
    Download = "DownloadInputData"


def dirac_make_replica_map(
    input_lfns, local_site_name=None, resolver=InputDataResolution.Protocol
):
    if resolver == InputDataResolution.Protocol:
        from DIRAC.WorkloadManagementSystem.Client.InputDataByProtocol import (
            InputDataByProtocol,
        )

        logger.info("ℹ️  Using InputDataByProtocol for input data resolution")
        module = InputDataByProtocol
    elif resolver == InputDataResolution.Download:
        from DIRAC.WorkloadManagementSystem.Client.DownloadInputData import (
            DownloadInputData,
        )

        module = DownloadInputData
        logger.info("ℹ️  Using DownloadInputData for input data resolution")
    else:
        raise ValueError(f"Unknown input data resolution method: {resolver}")

    from DIRAC import gConfig
    from DIRAC import siteName as DIRAC_siteName
    from DIRAC.Core.Utilities.SiteSEMapping import getSEsForSite
    from DIRAC.Interfaces.API.Dirac import Dirac as Dirac_Api

    api = Dirac_Api()

    diskSE = gConfig.getValue(
        api.section + "/DiskSE", ["-disk", "-DST", "-USER", "-FREEZER"]
    )
    tapeSE = gConfig.getValue(api.section + "/TapeSE", ["-tape", "-RDST", "-RAW"])

    if not local_site_name:
        siteName = DIRAC_siteName()
    else:
        siteName = local_site_name

    # fetch local SEs for the site
    localSEList = returnValueOrRaise(getSEsForSite(siteName))
    logger.info(f"Local SEs for site {siteName}: {localSEList}")

    # Get replicas for jobs from the LFNs
    logger.info(f"Attempting to resolve data for {siteName}")
    replicaDict = returnValueOrRaise(api.getReplicasForJobs(input_lfns))
    if replicaDict.get("Failed", {}):
        raise ValueError(
            f"Failed to resolve replicas for LFNs: {replicaDict['Failed']}"
        )
    replicaDict = replicaDict["Successful"]

    # get file metadata from the LFNs
    guidDict = returnValueOrRaise(api.getLfnMetadata(input_lfns))
    if guidDict.get("Failed", {}):
        raise ValueError(f"Failed to get metadata for LFNs: {guidDict['Failed']}")
    guidDict = guidDict["Successful"]

    # Save the original metadata before modifying guidDict
    # We need this later for building the catalog with size/checksum info
    import copy

    metadataDict = copy.deepcopy(guidDict)

    # prepare arguments for the module

    # for the Input Data Resolution module
    # merge replica info into guidDict
    for lfn, reps in replicaDict.items():
        guidDict[lfn].update(reps)
    resolvedData = guidDict.copy()

    argumentsDict = {
        "FileCatalog": {"Value": {"Successful": resolvedData}},
        "Configuration": {
            "JobID": None,
            "LocalSEList": localSEList,
            "DiskSEList": diskSE,
            "TapeSEList": tapeSE,
            "SiteName": siteName,
        },
        "InputData": input_lfns,
    }

    logger.debug(argumentsDict)

    # execute the module
    result = returnValueOrRaise(module(argumentsDict).execute())
    logger.debug(result)

    # Debug: Check what's in guidDict
    logger.debug(f"guidDict after module execution: {guidDict}")
    logger.debug(f"metadataDict (preserved): {metadataDict}")

    # Build the replica map based on the resolver type
    from diracx.core.models.replica_map import ReplicaMap

    if resolver == InputDataResolution.Download:
        # For DownloadInputData, use the local downloaded paths
        replica_map = ReplicaMap.model_validate(
            {
                lfn: ReplicaMap.MapEntry(
                    replicas=[
                        ReplicaMap.MapEntry.Replica(
                            url=result["Successful"][lfn]["path"],
                            se="DIRAC.Client.Local",
                        )
                    ],
                    size_bytes=file_info.get("Size", None),
                    checksum=ReplicaMap.MapEntry.Checksum(
                        guid=file_info.get("GUID", None),
                        adler32=(
                            file_info.get("Checksum")
                            if file_info.get("ChecksumType") == "Adler32"
                            else None
                        ),
                    ),
                )
                for lfn, file_info in metadataDict.items()
            }
        )
    else:
        # For InputDataByProtocol, use the remote URLs from replicaDict
        replica_map = ReplicaMap.model_validate(
            {
                lfn: ReplicaMap.MapEntry(
                    replicas=[
                        ReplicaMap.MapEntry.Replica(
                            url=as_url_if_local_path(replica_url),
                            se=replica_se,
                        )
                        for replica_se, replica_url in replicaDict[lfn].items()
                    ],
                    size_bytes=file_info.get("Size", None),
                    checksum=ReplicaMap.MapEntry.Checksum(
                        guid=file_info.get("GUID", None),
                        adler32=(
                            file_info.get("Checksum")
                            if file_info.get("ChecksumType") == "Adler32"
                            else None
                        ),
                    ),
                )
                for lfn, file_info in metadataDict.items()
            }
        )

    return replica_map


def as_url_if_local_path(path: str) -> str:
    """Convert a local file path to a file:// URL if needed.

    Args:
        path (str): The file path to convert.

    Returns:
        str: A file:// URL if the path is a local absolute path, otherwise the original path.
    """
    p = Path(path)
    if p.is_absolute() and p.exists():
        return p.as_uri()
    return path


class DataQualityFlag(StrEnum):
    OK = "OK"
    BAD = "BAD"
    UNCHECKED = "UNCHECKED"


class InputDatasetInner(BaseModel):
    """CWL hint for input dataset description."""

    class BookkeepingQuery(BaseModel):
        configName: str
        configVersion: str
        inFileType: str
        inProPass: str
        # TODO: Accept None and lists on these
        inDataQualityFlag: str | list[str] = "OK"
        inExtendedDQOK: Annotated[list[str], Field(min_length=1)] | None = None
        inProductionID: str = "ALL"
        inTCKs: str = "ALL"
        inSMOG2State: Annotated[list[str], Field(min_length=1)] | None = None

    conditions_dict: BookkeepingQuery
    conditions_description: str
    event_type: Annotated[str, StringConstraints(pattern=r"[0-9]{8}")]

    class LaunchParameters(BaseModel):
        """Extra info which is currently needed when making transformations from a request"""

        run_numbers: list[str] | None = None
        sample_max_md5: (
            Annotated[str, StringConstraints(pattern=r"[A-Z0-9]{32}")] | None
        ) = None
        sample_seed_md5: (
            Annotated[str, StringConstraints(pattern=r"[A-Z0-9]{32}")] | None
        ) = None
        start_run: int | None = None
        end_run: int | None = None

        @model_validator(mode="after")
        def validate_run_numbers(self) -> Self:
            if self.run_numbers and (
                self.start_run is not None or self.end_run is not None
            ):
                raise ValueError(
                    "Cannot specify both run_numbers and start_run/end_run"
                )
            return self

    launch_parameters: Annotated[
        LaunchParameters, Field(default_factory=LaunchParameters)
    ]

    def make_bkquery(self) -> dict[str, Any]:
        """Create BKQuery dict from InputDataset for querying the Bookkeeping.

        This transforms the CWL hint structure into the format expected by
        LHCbDIRAC's BookkeepingClient.getFilesWithMetadata().

        Returns:
            A dictionary suitable for BookkeepingClient BK queries.
        """
        # Required base fields
        bk_query_dict: dict[str, Any] = {
            "FileType": self.conditions_dict.inFileType,
            "EventType": str(self.event_type),
            "ConfigName": self.conditions_dict.configName,
            "ConfigVersion": self.conditions_dict.configVersion,
        }

        # Optional: Data taking conditions (sim condition or beam conditions)
        if self.conditions_description:
            bk_query_dict["DataTakingConditions"] = self.conditions_description

        # Optional: Processing pass
        if self.conditions_dict.inProPass:
            bk_query_dict["ProcessingPass"] = self.conditions_dict.inProPass

        # Optional: Data quality flag(s)
        # Can be a single string or a list of strings
        if self.conditions_dict.inDataQualityFlag:
            dq_flag = self.conditions_dict.inDataQualityFlag
            if isinstance(dq_flag, list):
                # List format: convert to list of strings
                bk_query_dict["DataQualityFlag"] = [str(f) for f in dq_flag]
            else:
                # Single value: convert to list with one element
                bk_query_dict["DataQualityFlag"] = str(dq_flag).split(",")

        # Optional: Extended DQ OK flags
        if self.conditions_dict.inExtendedDQOK:
            bk_query_dict["ExtendedDQOK"] = self.conditions_dict.inExtendedDQOK

        # Optional: Production ID filter
        if (
            self.conditions_dict.inProductionID
            and self.conditions_dict.inProductionID != "ALL"
        ):
            bk_query_dict["ProductionID"] = self.conditions_dict.inProductionID

        # Optional: TCKs filter
        if self.conditions_dict.inTCKs and self.conditions_dict.inTCKs != "ALL":
            bk_query_dict["TCKs"] = self.conditions_dict.inTCKs

        # Optional: SMOG2 state
        if self.conditions_dict.inSMOG2State:
            bk_query_dict["SMOG2"] = self.conditions_dict.inSMOG2State

        # Launch parameters: sample selection
        if (
            self.launch_parameters.sample_max_md5
            and self.launch_parameters.sample_seed_md5
        ):
            bk_query_dict["SampleMax"] = self.launch_parameters.sample_max_md5
            bk_query_dict["SampleSeedMD5"] = self.launch_parameters.sample_seed_md5

        # Launch parameters: run number selection
        # Validate that we don't mix run_numbers with start_run/end_run
        if self.launch_parameters.run_numbers and (
            self.launch_parameters.start_run is not None
            or self.launch_parameters.end_run is not None
        ):
            raise ValueError("Cannot specify both run_numbers and start_run/end_run")

        # Validate start_run <= end_run
        if (
            self.launch_parameters.start_run is not None
            and self.launch_parameters.end_run is not None
            and self.launch_parameters.end_run < self.launch_parameters.start_run
        ):
            raise ValueError(
                f"end_run ({self.launch_parameters.end_run}) must be >= start_run ({self.launch_parameters.start_run})"
            )

        # Add run number filters
        if self.launch_parameters.start_run is not None:
            bk_query_dict["StartRun"] = self.launch_parameters.start_run

        if self.launch_parameters.end_run is not None:
            bk_query_dict["EndRun"] = self.launch_parameters.end_run

        if self.launch_parameters.run_numbers:
            # Convert to list of strings if needed
            bk_query_dict["RunNumbers"] = [
                str(r) if not isinstance(r, str) else r
                for r in self.launch_parameters.run_numbers
            ]

        return bk_query_dict


class InputDataset(BaseModel):
    class_: Literal["dirac:Production"] = Field(alias="class")

    input_dataset_plugin: Literal["LHCbBookkeepingPlugin"]
    input_dataset_config: InputDatasetInner


def do_bkquery(
    input_dataset: InputDatasetInner,
    inputFiles: list[str] | None = None,
    numTestLFNs: int | None = None,
    pick_smallest: bool = False,
):
    """Query the Bookkeeping for LFNs matching the input dataset criteria.

    :param input_dataset: InputDataset model with BK query parameters
    :param inputFiles: Optional pre-selected list of LFNs (if provided, skip BK query)
    :param numTestLFNs: Number of LFNs to retrieve. If None, retrieves all available LFNs.
    :param pick_smallest: If True, pick the smallest file(s) for faster testing
    :return: List of LFNs with available replicas
    """
    from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
    from DIRAC.DataManagementSystem.Client.DataManager import DataManager
    from LHCbDIRAC.BookkeepingSystem.Client.BookkeepingClient import BookkeepingClient

    bkQueryDict = input_dataset.make_bkquery()

    if not inputFiles:
        logger.info("Querying Bookkeeping for LFNs...")
        result = returnValueOrRaise(
            BookkeepingClient().getFilesWithMetadata(
                bkQueryDict | {"OnlyParameters": ["FileName", "FileSize"]}
            )
        )
        if result["TotalRecords"] == 0:
            query_str = json.dumps(bkQueryDict, indent=4)
            raise ValueError(
                f"No input files found in the bookkeeping.\n\n"
                f"Bookkeeping query used:\n"
                f"\n{query_str}\n\n"
                f"Please verify that:\n"
                f"  - The bookkeeping path is correct\n"
                f"  - The requested runs (if specified) actually exist\n"
                f"  - The data quality flags are correct\n"
            )

        logger.info(f"Found {result['TotalRecords']} files in Bookkeeping")

        # Get all LFNs first
        filenameIndex = result["ParameterNames"].index("FileName")
        sizeIndex = result["ParameterNames"].index("FileSize")
        records = result["Records"]

        # If sampling (numTestLFNs specified), filter and shuffle
        if numTestLFNs is not None:
            # Sort by size
            records = sorted(records, key=lambda x: x[sizeIndex])

            if pick_smallest:
                # Pick smallest files for fast testing
                logger.info(
                    f"Picking {numTestLFNs} smallest file(s) for fast testing..."
                )
                records = records[:numTestLFNs]
            else:
                # Remove the smallest 50% of files to avoid unusually small files
                if len(records) // 2 >= numTestLFNs:
                    records = records[len(records) // 2 :]
                    logger.debug(
                        f"Filtered out smallest 50% of files, {len(records)} remaining"
                    )
                # Shuffle the LFNs so we pick random ones
                random.shuffle(records)

            logger.info(f"Checking for replicas on disk (need {numTestLFNs} files)...")
        else:
            # Getting all LFNs - just extract them without filtering/shuffling
            logger.info("Retrieving all files from Bookkeeping...")
            inputFiles = [record[filenameIndex] for record in records]
            logger.info(f"Retrieved {len(inputFiles)} LFN(s) from Bookkeeping")
            return inputFiles

        # Only run tests with files which have available replicas (sampling mode)
        inputFiles = []
        skipped_files = []
        for record in records:
            lfn = record[filenameIndex]
            replica_result = returnValueOrRaise(
                DataManager().getReplicasForJobs([lfn], diskOnly=True)
            )
            inputFiles.extend(replica_result["Successful"])
            if len(inputFiles) == numTestLFNs:
                break
            if replica_result["Failed"]:
                skipped_files.extend(replica_result["Failed"])
                logger.warning(
                    f"Skipping LFN (no disk replicas): {replica_result['Failed']}"
                )
        else:
            error_msg = (
                f"Insufficient input files with available (disk) replicas for jobs found.\n\n"
                f"Summary:\n"
                f"  - Files requested for test: {numTestLFNs}\n"
                f"  - Files found with replicas: {len(inputFiles)}\n"
                f"  - Files skipped (no disk replicas): {len(skipped_files)}\n\n"
                f"This usually means the files need to be staged from tape.\n\n"
                f"Solutions:\n"
                f"  1. Request staging of the required files by contacting LHCb Data Management\n\n"
                f"Contact: lhcb-datamanagement@cern.ch"
            )
            if skipped_files:
                error_msg += "\n\nExample skipped files:\n  " + "\n  ".join(
                    skipped_files[:3]
                )
                if len(skipped_files) > 3:
                    error_msg += f"\n  ... and {len(skipped_files) - 3} more"
            raise ValueError(error_msg)

        logger.info(f"Found {len(inputFiles)} files with available replicas")

    if len(inputFiles) < numTestLFNs:
        raise ValueError(
            f"Insufficient input files available.\n\n"
            f"Summary:\n"
            f"  - Files requested for test: {numTestLFNs}\n\n"
            f"  - Files available: {len(inputFiles)}\n"
            f"This could indicate that some files need to be staged from tape storage.\n\n"
            f"Solutions:\n"
            f"  1. Use --num-test-lfns={len(inputFiles)} to work with the available files\n"
            f"  2. Contact LHCb Data Management to request staging of additional files\n"
            f"Contact: lhcb-datamanagement@cern.ch"
        )

    return inputFiles


app = typer.Typer(
    help="Generate replica map and input YAML from CWL workflow with inputDataset hint",
    no_args_is_help=True,
)


@app.command()
def main(
    cwl_file: Path = typer.Argument(
        ...,
        help="Path to CWL file with dirac:Production hint",
        exists=True,
        dir_okay=False,
        readable=True,
    ),
    n_lfns: int | None = typer.Option(
        None,
        "--n-lfns",
        "-n",
        help="Number of LFNs to retrieve from Bookkeeping (default: all available LFNs)",
    ),
    pick_smallest: bool = typer.Option(
        False,
        "--pick-smallest-lfn",
        help="Pick the smallest file(s) for faster testing (requires --n-lfns)",
    ),
    download: bool = typer.Option(
        False,
        "--download",
        help="Download input files locally instead of using remote protocol where available",
    ),
    output_yaml: Path = typer.Option(
        None,
        "--output-yaml",
        "-o",
        help="Output path for CWL input YAML file (default: <cwl_file>-inputs.yml)",
    ),
    output_map: Path = typer.Option(
        None,
        "--output-map",
        "-c",
        help="Output path for replica map JSON (default: <cwl_file>-replica-map.json)",
    ),
    verbose: bool = typer.Option(
        False,
        "--verbose",
        "-v",
        help="Enable verbose logging",
    ),
):
    """Generate replica map and input YAML from a CWL workflow.

    Reads the inputDataset hint from a CWL workflow file, queries the LHCb Bookkeeping
    for matching files with available replicas, and generates:

    1. A CWL input YAML file with the input-data parameter
    2. A replica map JSON file for local execution
    """
    if verbose:
        logger.setLevel(logging.DEBUG)

    # Validate flags
    if pick_smallest and n_lfns is None:
        console.print(
            "[red]Error:[/red] --pick-smallest-lfn requires --n-lfns to be specified"
        )
        raise typer.Exit(1)

    # initialise dirac first
    import DIRAC

    DIRAC.initialize()

    # Set default output paths if not provided
    if output_yaml is None:
        output_yaml = cwl_file.parent / f"{cwl_file.stem}-inputs.yml"
    if output_map is None:
        output_map = cwl_file.parent / f"{cwl_file.stem}-replica-map.json"

    # Read the CWL file and extract inputDataset hint
    logger.info(f"Reading CWL file: {cwl_file}")
    with open(cwl_file, "r") as f:
        cwl_content = yaml.safe_load(f)

    # Find the dirac:Production hint
    hints = cwl_content.get("hints", [])
    input_dataset_dict = None
    for hint in hints:
        if hint.get("class") == "dirac:Production":
            input_dataset_dict = hint
            break

    if not input_dataset_dict:
        console.print("[red]Error:[/red] No dirac:Production hint found in CWL file")
        console.print(
            "Expected a hint with class='dirac:Production' in the 'hints' section"
        )
        raise typer.Exit(1)

    # Parse into InputDataset model
    try:
        input_dataset_outer = InputDataset(**input_dataset_dict)
    except Exception as e:
        console.print(f"[red]Error:[/red] Failed to parse dirac:Production hint: {e}")
        raise typer.Exit(1)

    input_dataset = input_dataset_outer.input_dataset_config

    logger.info("Found dirac:Production hint:")
    logger.info(f"  Event Type: {input_dataset.event_type}")
    logger.info(
        f"  Config: {input_dataset.conditions_dict.configName}/{input_dataset.conditions_dict.configVersion}"
    )
    logger.info(f"  Processing Pass: {input_dataset.conditions_dict.inProPass}")

    # Query Bookkeeping for LFNs
    try:
        lfns_list = do_bkquery(
            input_dataset, numTestLFNs=n_lfns, pick_smallest=pick_smallest
        )
    except Exception as e:
        console.print(f"[red]Error:[/red] Bookkeeping query failed: {e}")
        raise typer.Exit(1)

    logger.info(f"Retrieved {len(lfns_list)} LFNs")

    # Generate replica map using DIRAC

    logger.info("Generating replica map with DIRAC...")
    try:
        replica_map = dirac_make_replica_map(
            lfns_list,
            resolver=(
                InputDataResolution.Download
                if download
                else InputDataResolution.Protocol
            ),
        )
    except Exception as e:
        console.print(f"[red]Error:[/red] Failed to generate replica map: {e}")
        raise typer.Exit(1)

    # Generate CWL input YAML file
    cwl_inputs = {
        "input-data": [
            {
                "class": "File",
                "location": f"LFN:{lfn}",
            }
            for lfn in lfns_list
        ]
    }

    logger.info(f"Writing CWL input YAML to: {output_yaml}")
    with open(output_yaml, "w") as f:
        yaml.dump(cwl_inputs, f, default_flow_style=False, sort_keys=False)

    logger.info(f"Writing replica map to: {output_map}")
    with open(output_map, "w") as f:
        # Convert Pydantic model to dict for JSON serialization
        f.write(replica_map.model_dump_json(indent=2))

    console.print("\n[green]✓[/green] Successfully generated:")
    console.print(f"  • Input YAML: {output_yaml}")
    console.print(f"  • Replica map: {output_map}")
    console.print(f"\nRetrieved {len(lfns_list)} LFN(s)")


if __name__ == "__main__":
    app()
