"""Utility functions for loading data profile artifacts."""

import json
import pickle
from pathlib import Path
from typing import Any, Dict, Optional, Set, Tuple


def find_latest_run(profile_dir: Path, version: str, split: str = "train") -> Optional[Path]:
    """Find the latest run directory for a given version.

    Args:
        profile_dir: Base directory containing all profiles
        version: Version string (e.g., "v0", "v1", "v9")
        split: Data split ("train" or "val")

    Returns:
        Path to the latest run directory, or None if not found
    """
    # Look for directories matching the version pattern
    pattern = f"{version}_full_*"
    matching_dirs = sorted(profile_dir.glob(pattern))

    if not matching_dirs:
        print(f"Warning: No directories found for {version} in {profile_dir}")
        return None

    # Use the latest run (last in sorted order)
    latest_dir = matching_dirs[-1]
    split_dir = latest_dir / split

    if not split_dir.exists():
        print(f"Warning: Split directory {split_dir} does not exist")
        return None

    return split_dir


def load_analysis_results(run_dir: Path) -> Optional[Dict[str, Any]]:
    """Load analysis_results.json from a run directory.

    Args:
        run_dir: Path to the run directory (e.g., v0_full_20250930_150622/train/)

    Returns:
        Dictionary containing analysis results, or None if file not found
    """
    analysis_file = run_dir / "analysis_results.json"

    if not analysis_file.exists():
        print(f"Warning: Analysis file not found: {analysis_file}")
        return None

    try:
        with open(analysis_file, "r") as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading {analysis_file}: {e}")
        return None


def load_id_set(run_dir: Path) -> Optional[Set[str]]:
    """Load ID set from pickle file.

    Args:
        run_dir: Path to the run directory (e.g., v0_full_20250930_150622/train/)

    Returns:
        Set of ID strings, or None if file not found
    """
    id_set_file = run_dir / "id_set_train.pkl"

    # Try alternate filename pattern
    if not id_set_file.exists():
        id_set_file = run_dir / "id_set_val.pkl"

    if not id_set_file.exists():
        print(f"Warning: ID set file not found in {run_dir}")
        return None

    try:
        with open(id_set_file, "rb") as f:
            id_set = pickle.load(f)

        if not isinstance(id_set, set):
            print(f"Warning: ID set is not a Python set, got {type(id_set)}")
            return None

        return id_set
    except Exception as e:
        print(f"Error loading ID set from {id_set_file}: {e}")
        return None


def load_version_data(
    profile_dir: Path, version: str, split: str = "train"
) -> Optional[Tuple[Dict[str, Any], Set[str], Path]]:
    """Load both analysis results and ID set for a version.

    Args:
        profile_dir: Base directory containing all profiles
        version: Version string (e.g., "v0", "v1", "v9")
        split: Data split ("train" or "val")

    Returns:
        Tuple of (analysis_results, id_set, run_dir) or None if loading failed
    """
    run_dir = find_latest_run(profile_dir, version, split)

    if run_dir is None:
        return None

    analysis_results = load_analysis_results(run_dir)
    id_set = load_id_set(run_dir)

    if analysis_results is None or id_set is None:
        print(f"Warning: Failed to load data for {version}")
        return None

    return analysis_results, id_set, run_dir


def get_source_jsonl_path(analysis_results: Dict[str, Any]) -> Optional[Path]:
    """Extract source JSONL file path from analysis results.

    Args:
        analysis_results: Dictionary from analysis_results.json

    Returns:
        Path to source JSONL file, or None if not found
    """
    try:
        metadata = analysis_results.get("_metadata", {})
        file_path = metadata.get("file_path")

        if file_path:
            return Path(file_path)

        print("Warning: No file_path found in _metadata")
        return None
    except Exception as e:
        print(f"Error extracting source JSONL path: {e}")
        return None


def get_version_summary(analysis_results: Dict[str, Any]) -> Dict[str, Any]:
    """Extract key summary statistics from analysis results.

    Args:
        analysis_results: Dictionary from analysis_results.json

    Returns:
        Dictionary with summary statistics
    """
    id_section = analysis_results.get("id", {}).get("summary", {})
    tags_section = analysis_results.get("tags", {}).get("summary", {})
    text_section = analysis_results.get("text", {}).get("summary", {})
    language_section = analysis_results.get("language", {}).get("summary", {})
    weight_section = analysis_results.get("weight", {}).get("summary", {})
    stems_section = analysis_results.get("stems", {}).get("summary", {})
    paths_section = analysis_results.get("paths", {}).get("summary", {})
    metadata_section = analysis_results.get("_metadata", {})

    return {
        "total_records": id_section.get("total_records", 0),
        "unique_ids": id_section.get("unique_ids", 0),
        "duplicate_ids": id_section.get("duplicate_ids", 0),
        "records_with_tags": tags_section.get("records_with_tags", 0),
        "unique_tags": tags_section.get("unique_tags", 0),
        "records_with_text": text_section.get("records_with_text", 0),
        "records_with_lang": language_section.get("records_with_lang", 0),
        "unique_languages": language_section.get("unique_languages", 0),
        "language_mismatch_rate": language_section.get("mismatch_rate", 0.0),
        "records_with_weight": weight_section.get("records_with_weight", 0),
        "records_with_stems": stems_section.get("records_with_stems", 0),
        "unique_stem_names": stems_section.get("unique_stem_names", 0),
        "local_path_count": paths_section.get("records_with_local_path", 0),
        "s3_path_count": paths_section.get("records_with_s3_path", 0),
        "path_anomalies": paths_section.get("total_anomalies", 0),
        "source_file": metadata_section.get("file_path", ""),
        "processing_time": metadata_section.get("timing", {}).get("overall_time", 0),
    }