"""
Job Merging System for Multi-Config Pipeline

Merges jobs from multiple validated pipeline configuration files into a single
execution plan while preserving source attribution for error handling.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Tuple,
)

from .config_discovery import ConfigFile


@dataclass
class JobSource:
    """Tracks the source of a job for error attribution."""

    config_file: Path
    config_name: str
    job_index: int
    original_job: Dict[str, Any]

    def __str__(self) -> str:
        return f"{self.config_file}[job:{self.job_index}]"


@dataclass
class MergedJobConfig:
    """Represents merged pipeline configuration from multiple config files."""

    stage: str
    jobs: List[Dict[str, Any]]
    job_sources: List[JobSource]
    total_jobs: int
    source_configs: List[str]

    def get_job_source(self, merged_job_index: int) -> Optional[JobSource]:
        """Get source information for a merged job by its index."""
        if 0 <= merged_job_index < len(self.job_sources):
            return self.job_sources[merged_job_index]
        return None

    def get_config_summary(self) -> Dict[str, int]:
        """Get summary of jobs per config file."""
        summary = {}
        for source in self.job_sources:
            config_name = str(source.config_file)
            summary[config_name] = summary.get(config_name, 0) + 1
        return summary


class JobMergingError(Exception):
    """Raised when job merging encounters an error."""

    pass


def extract_jobs_from_config(
    config_file: ConfigFile, stage_name: str
) -> List[Dict[str, Any]]:
    """
    Extract jobs for a specific stage from a config file.

    Args:
        config_file: Config file to extract jobs from
        stage_name: Name of the stage to extract jobs for

    Returns:
        List of job configurations for the specified stage
    """
    content = config_file.content or config_file.load_content()
    stages = content.get("stages", {})

    if stage_name not in stages:
        return []

    stage_config = stages[stage_name]

    # Handle different stage config formats
    if isinstance(stage_config, list):
        # Direct list of jobs
        return stage_config
    elif isinstance(stage_config, dict):
        # Object with jobs key
        return stage_config.get("jobs", [])
    else:
        # Invalid format (should be caught by validation)
        return []


def create_job_with_source_tracking(
    job: Dict[str, Any], config_file: ConfigFile, job_index: int
) -> Tuple[Dict[str, Any], JobSource]:
    """
    Create a job with source tracking metadata and job ID for data lineage.

    Args:
        job: Original job configuration
        config_file: Source config file
        job_index: Index of job within the config file

    Returns:
        Tuple of (enhanced_job, job_source)
    """
    # Generate unique job ID for data lineage tracking
    job_id = f"{config_file.name}_{job_index}"

    # Extract metadata config from source file
    config_content = config_file.content or config_file.load_content()
    metadata_config = config_content.get("stages", {}).get("metadata", {})

    # Create source tracking info
    job_source = JobSource(
        config_file=config_file.relative_path,
        config_name=config_file.name,
        job_index=job_index,
        original_job=job.copy(),
    )

    # Add source tracking and job lineage to job
    enhanced_job = job.copy()
    enhanced_job["_config_source"] = {
        "file": str(config_file.relative_path),
        "name": config_file.name,
        "job_index": job_index,
    }
    enhanced_job["_job_id"] = job_id
    enhanced_job["_metadata_config"] = metadata_config

    return enhanced_job, job_source


def merge_pipeline_jobs(
    validated_configs: List[ConfigFile], stage_name: str
) -> MergedJobConfig:
    """
    Merge jobs from all validated configs into single execution plan.

    Args:
        validated_configs: List of validated config files
        stage_name: Target stage name

    Returns:
        MergedJobConfig with unified job list and source tracking

    Raises:
        JobMergingError: If merging encounters an error
    """
    if not validated_configs:
        raise JobMergingError("No config files provided for job merging")

    merged_jobs = []
    job_sources = []
    configs_with_jobs = []

    for config_file in validated_configs:
        try:
            # Extract jobs for this stage
            jobs = extract_jobs_from_config(config_file, stage_name)

            if jobs:
                configs_with_jobs.append(str(config_file.relative_path))

                # Process each job with source tracking
                for job_index, job in enumerate(jobs):
                    if not isinstance(job, dict):
                        raise JobMergingError(
                            f"Invalid job format in {config_file.relative_path}[{job_index}]: "
                            f"expected dict, got {type(job).__name__}"
                        )

                    enhanced_job, job_source = create_job_with_source_tracking(
                        job, config_file, job_index
                    )

                    merged_jobs.append(enhanced_job)
                    job_sources.append(job_source)

        except Exception as e:
            raise JobMergingError(
                f"Error merging jobs from {config_file.relative_path}: {e}"
            )

    return MergedJobConfig(
        stage=stage_name,
        jobs=merged_jobs,
        job_sources=job_sources,
        total_jobs=len(merged_jobs),
        source_configs=configs_with_jobs,
    )


def validate_merged_job_consistency(merged_config: MergedJobConfig) -> List[str]:
    """
    Validate consistency across merged jobs and identify potential issues.

    Args:
        merged_config: Merged job configuration

    Returns:
        List of warning messages about potential issues
    """
    warnings = []

    if not merged_config.jobs:
        warnings.append("No jobs found across all config files for this stage")
        return warnings

    # Check for date range consistency (for generate stage)
    if merged_config.stage == "generate":
        date_ranges = []
        for job in merged_config.jobs:
            start_date = job.get("start_date")
            end_date = job.get("end_date")
            if start_date and end_date:
                date_ranges.append((start_date, end_date))

        if len(set(date_ranges)) > 1:
            warnings.append(
                f"Found {len(set(date_ranges))} different date ranges across jobs. "
                f"This may result in inconsistent time periods in the output dataset."
            )

    # Check for reference data file overlaps
    reference_files = []
    for i, job in enumerate(merged_config.jobs):
        ref_data = job.get("reference_data")
        if ref_data:
            job_source = merged_config.get_job_source(i)
            reference_files.append(
                (ref_data, str(job_source.config_file) if job_source else "unknown")
            )

    ref_file_usage = {}
    for ref_file, config_file in reference_files:
        if ref_file not in ref_file_usage:
            ref_file_usage[ref_file] = []
        ref_file_usage[ref_file].append(config_file)

    for ref_file, config_files in ref_file_usage.items():
        if len(set(config_files)) > 1:
            warnings.append(
                f"Reference file '{ref_file}' is used by multiple configs: {', '.join(set(config_files))}. "
                f"This may result in duplicate or conflicting data."
            )

    return warnings


def get_job_merging_summary(merged_config: MergedJobConfig) -> Dict[str, Any]:
    """
    Generate a summary of the job merging results.

    Args:
        merged_config: Merged job configuration

    Returns:
        Dictionary with merging statistics and details
    """
    config_summary = merged_config.get_config_summary()

    return {
        "stage": merged_config.stage,
        "total_jobs": merged_config.total_jobs,
        "source_configs": len(merged_config.source_configs),
        "config_files": merged_config.source_configs,
        "jobs_per_config": config_summary,
        "has_jobs": merged_config.total_jobs > 0,
    }


def remove_source_metadata_from_jobs(
    jobs: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """
    Remove source tracking metadata from jobs before pipeline execution.

    This ensures that the source tracking metadata doesn't interfere with
    pipeline processing while preserving job ID and metadata config for data lineage.

    Args:
        jobs: List of jobs with source metadata

    Returns:
        List of jobs with non-essential source metadata removed
    """
    cleaned_jobs = []

    for job in jobs:
        cleaned_job = job.copy()
        # Remove the source tracking metadata (keep _job_id and _metadata_config)
        cleaned_job.pop("_config_source", None)
        cleaned_jobs.append(cleaned_job)

    return cleaned_jobs


def build_job_metadata_mapping(jobs: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
    """
    Build mapping from job IDs to their metadata configurations.

    Args:
        jobs: List of enhanced jobs with _job_id and _metadata_config

    Returns:
        Dictionary mapping job_id -> metadata_config
    """
    mapping = {}

    for job in jobs:
        job_id = job.get("_job_id")
        metadata_config = job.get("_metadata_config", {})

        if job_id:
            mapping[job_id] = metadata_config

    return mapping


def merge_and_validate_pipeline_jobs(
    validated_configs: List[ConfigFile], stage_name: str, verbose: bool = False
) -> MergedJobConfig:
    """
    Main entry point for merging and validating pipeline jobs.

    Combines job merging, consistency validation, and optional verbose reporting.

    Args:
        validated_configs: List of validated config files
        stage_name: Target stage name
        verbose: Whether to print merging details

    Returns:
        MergedJobConfig with unified job list

    Raises:
        JobMergingError: If merging fails
    """
    # Phase 1: Merge jobs from all configs
    merged_config = merge_pipeline_jobs(validated_configs, stage_name)

    # Phase 2: Validate consistency
    warnings = validate_merged_job_consistency(merged_config)

    # Phase 3: Summary (optional verbose output)
    if verbose:
        summary = get_job_merging_summary(merged_config)
        print("\n🔄 Job Merging Summary:")
        print(f"   Stage: {summary['stage']}")
        print(f"   Total jobs: {summary['total_jobs']}")
        print(f"   Source configs: {summary['source_configs']}")

        if summary["jobs_per_config"]:
            print("\n📋 Jobs per config:")
            for config_file, job_count in summary["jobs_per_config"].items():
                print(f"   • {config_file}: {job_count} job(s)")

        if warnings:
            print("\n⚠️  Warnings:")
            for warning in warnings:
                print(f"   • {warning}")

    return merged_config
