#!/usr/bin/env python3
"""Main pipeline for data monitoring."""

import sys
import time
import json
import pickle
import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
import argparse
from datetime import datetime

# Add parent directory to path
sys.path.append(str(Path(__file__).parent.parent.parent))

from sunodata.data_monitor.analyzers.id_analyzer import IDAnalyzer
from sunodata.data_monitor.analyzers.tags_analyzer import TagsAnalyzer
from sunodata.data_monitor.analyzers.text_analyzer import TextAnalyzer
from sunodata.data_monitor.analyzers.lang_analyzer import LanguageAnalyzer
from sunodata.data_monitor.analyzers.weight_analyzer import WeightAnalyzer
from sunodata.data_monitor.analyzers.stems_analyzer import StemsAnalyzer
from sunodata.data_monitor.analyzers.paths_analyzer import PathsAnalyzer
from sunodata.data_monitor.analyzers.lists_analyzer import ListsAnalyzer
from sunodata.data_monitor.analyzers.sft_analyzer import SFTAnalyzer
from sunodata.data_monitor.analyzers.sample_collector import SampleCollector

from sunodata.data_monitor.utils.file_reader import FileReader
from sunodata.data_monitor.utils.cache_manager import CacheManager
from sunodata.data_monitor.utils.chunk_processor import ChunkProcessor

# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


class DataMonitorPipeline:
    """Main pipeline for data monitoring and analysis."""

    def __init__(self, config_path: str = "config.yaml", s3_category_filter: Optional[List[str]] = None):
        """Initialize pipeline.

        Args:
            config_path: Path to configuration file
            s3_category_filter: Optional list of keywords to filter records by s3_filepath
                               (e.g., ['discogs', 'pond5']). Only records matching these keywords
                               will be processed, significantly improving performance.
        """
        self.config = self.load_config(config_path)
        self.cache_manager = CacheManager()
        self.s3_category_filter = s3_category_filter
        self.setup_directories()

    def load_config(self, config_path: str) -> Dict:
        """Load configuration from YAML file."""
        import yaml

        config_path = Path(__file__).parent / config_path
        with open(config_path, "r") as f:
            return yaml.safe_load(f)

    def setup_directories(self):
        """Setup output directories."""
        # Find best data directory
        for path in self.config["data_base_paths"]:
            expanded = Path(path).expanduser()
            if expanded.parent.exists():
                self.data_dir = expanded
                break
        else:
            self.data_dir = Path("/tmp/suno_data_monitor")

        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.output_dir = self.data_dir / "outputs"
        self.output_dir.mkdir(exist_ok=True)

        logger.info(f"Data directory: {self.data_dir}")

    def get_analyzers(
        self, analyzer_names: Optional[List[str]] = None, output_dir: Optional[Path] = None
    ) -> List[Any]:
        """Get list of analyzer instances.

        Args:
            analyzer_names: List of analyzer names to use. None means all.
            output_dir: Directory for outputs

        Returns:
            List of analyzer instances
        """
        if output_dir is None:
            output_dir = self.output_dir

        all_analyzers = {
            "id": IDAnalyzer,
            "tags": TagsAnalyzer,
            "text": TextAnalyzer,
            "lang": LanguageAnalyzer,
            "weight": WeightAnalyzer,
            "stems": StemsAnalyzer,
            "paths": PathsAnalyzer,
            "lists": ListsAnalyzer,
            "samples": SampleCollector,
        }

        if analyzer_names is None:
            analyzer_names = list(all_analyzers.keys())

        analyzers = []
        for name in analyzer_names:
            if name in all_analyzers:
                analyzers.append(all_analyzers[name](output_dir))
            else:
                logger.warning(f"Unknown analyzer: {name}")

        return analyzers

    def analyze_dataset(
        self,
        dataset_name: str,
        split: str = "train",
        analyzers: Optional[List[str]] = None,
        use_cache: bool = True,
    ) -> Dict[str, Any]:
        """Analyze a dataset.

        Args:
            dataset_name: Name of dataset (e.g., 'v9')
            split: Dataset split ('train' or 'val')
            analyzers: List of analyzer names to use
            use_cache: Whether to use cached results

        Returns:
            Analysis results
        """
        # Get dataset configuration
        if dataset_name not in self.config["datasets"]:
            raise ValueError(f"Unknown dataset: {dataset_name}")

        dataset_config = self.config["datasets"][dataset_name]
        file_path = Path(dataset_config[split])

        if not file_path.exists():
            raise FileNotFoundError(f"Dataset file not found: {file_path}")

        # Create output directory for this run
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir = self.output_dir / f"{dataset_name}_{split}_{timestamp}"
        run_dir.mkdir(parents=True, exist_ok=True)

        logger.info(f"Analyzing {dataset_name} {split} dataset")
        logger.info(f"File: {file_path}")
        logger.info(f"Output: {run_dir}")

        # Check cache
        if use_cache:
            cached_result = self.cache_manager.get_analysis_result(file_path, f"full_{split}")
            if cached_result:
                logger.info("Using cached results")
                return cached_result

        # Initialize analyzers
        analyzer_instances = self.get_analyzers(analyzers, run_dir)

        # Process with ChunkProcessor
        processor = ChunkProcessor(
            file_path,
            chunk_size=self.config["processing"]["chunk_size"],
            num_workers=self.config["processing"]["parallel_workers"],
            cache_manager=self.cache_manager,
            s3_category_filter=self.s3_category_filter,
        )

        start_time = time.time()
        results = processor.process_file(analyzer_instances, run_dir)
        elapsed = time.time() - start_time

        logger.info(f"Analysis completed in {elapsed:.2f}s")

        # Save individual analyzer results
        for analyzer in analyzer_instances:
            if analyzer.name in results and not analyzer.name.startswith("_"):
                # Remove non-serializable fields before saving individual files
                result_copy = results[analyzer.name].copy()
                if "id_set" in result_copy:
                    del result_copy["id_set"]  # Don't save ID set in individual files
                if "duplicate_ids_list" in result_copy:
                    del result_copy[
                        "duplicate_ids_list"
                    ]  # Don't save duplicate list in individual files
                analyzer.save_results(result_copy, suffix=f"_{split}")

        # Save ID set as pickle for efficient reuse
        if "id" in results and "id_set" in results["id"]:
            id_set_file = run_dir / f"id_set_{split}.pkl"
            with open(id_set_file, "wb") as f:
                pickle.dump(results["id"]["id_set"], f)
            logger.info(f"ID set saved to: {id_set_file}")

        # Save combined results (convert sets to lists for JSON)
        results_serializable = {}
        for key, value in results.items():
            if isinstance(value, dict):
                value_copy = value.copy()
                if "id_set" in value_copy:
                    value_copy["id_set"] = (
                        f"<set with {len(value_copy['id_set'])} IDs, saved to id_set_{split}.pkl>"
                    )
                if "duplicate_ids_list" in value_copy:
                    value_copy["duplicate_ids_list"] = (
                        f"<{len(value_copy['duplicate_ids_list'])} duplicates, saved to duplicate_ids_{split}.json>"
                    )
                results_serializable[key] = value_copy
            else:
                results_serializable[key] = value

        results_file = run_dir / "analysis_results.json"
        with open(results_file, "w") as f:
            json.dump(results_serializable, f, indent=2, ensure_ascii=False)

        logger.info(f"Results saved to: {results_file}")

        # Cache results
        if use_cache:
            self.cache_manager.save_analysis_result(file_path, f"full_{split}", results)

        return results

    def analyze_sft(self, dataset_name: str, train_results: Optional[Dict] = None) -> Dict[str, Any]:
        """Analyze SFT ID sets.

        Args:
            dataset_name: Name of dataset
            train_results: Pre-computed training set results (for ID comparison)

        Returns:
            SFT analysis results
        """
        dataset_config = self.config["datasets"][dataset_name]

        if "sft_ids" not in dataset_config:
            logger.warning(f"No SFT IDs configured for {dataset_name}")
            return {}

        sft_file = Path(dataset_config["sft_ids"])

        if not sft_file.exists():
            logger.error(f"SFT file not found: {sft_file}")
            return {}

        # Get training IDs
        if train_results and "id" in train_results:
            train_ids = train_results["id"].get("id_set", set())
        else:
            # Load training IDs
            logger.info("Loading training IDs for SFT analysis...")
            train_file = Path(dataset_config["train"])
            reader = FileReader(train_file, self.cache_manager)
            train_ids = reader.extract_ids()

        # Analyze SFT
        analyzer = SFTAnalyzer(self.output_dir)
        results = analyzer.analyze_sft_file(sft_file, train_ids)

        # Save results
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = self.output_dir / f"{dataset_name}_sft_analysis_{timestamp}.json"
        with open(output_file, "w") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)

        logger.info(f"SFT analysis saved to: {output_file}")

        return results

    def analyze_dataset_to_dir(self, dataset_name: str, split: str, output_dir: Path) -> Dict[str, Any]:
        """Analyze dataset and save to specific directory.

        Args:
            dataset_name: Name of dataset
            split: Dataset split ('train' or 'val')
            output_dir: Output directory for results

        Returns:
            Analysis results
        """
        # Get dataset configuration
        if dataset_name not in self.config["datasets"]:
            raise ValueError(f"Unknown dataset: {dataset_name}")

        dataset_config = self.config["datasets"][dataset_name]
        file_path = Path(dataset_config[split])

        if not file_path.exists():
            raise FileNotFoundError(f"Dataset file not found: {file_path}")

        logger.info(f"Analyzing {dataset_name} {split} dataset")
        logger.info(f"File: {file_path}")
        logger.info(f"Output: {output_dir}")

        # Initialize analyzers
        analyzer_instances = self.get_analyzers(None, output_dir)

        # Process with ChunkProcessor
        processor = ChunkProcessor(
            file_path,
            chunk_size=self.config["processing"]["chunk_size"],
            num_workers=self.config["processing"]["parallel_workers"],
            cache_manager=self.cache_manager,
            s3_category_filter=self.s3_category_filter,
        )

        start_time = time.time()
        results = processor.process_file(analyzer_instances, output_dir)
        elapsed = time.time() - start_time

        logger.info(f"Analysis completed in {elapsed:.2f}s")

        # Save individual analyzer results
        for analyzer in analyzer_instances:
            if analyzer.name in results and not analyzer.name.startswith("_"):
                # Remove non-serializable fields before saving individual files
                result_copy = results[analyzer.name].copy()
                if "id_set" in result_copy:
                    del result_copy["id_set"]  # Don't save ID set in individual files
                if "duplicate_ids_list" in result_copy:
                    del result_copy[
                        "duplicate_ids_list"
                    ]  # Don't save duplicate list in individual files
                analyzer.save_results(result_copy, suffix=f"_{split}")

        # Save ID set as pickle for efficient reuse
        if "id" in results and "id_set" in results["id"]:
            id_set_file = output_dir / f"id_set_{split}.pkl"
            with open(id_set_file, "wb") as f:
                pickle.dump(results["id"]["id_set"], f)
            logger.info(f"ID set saved to: {id_set_file}")

            # Save duplicate IDs list separately
            if "duplicate_ids_list" in results["id"]:
                duplicate_ids_file = output_dir / f"duplicate_ids_{split}.json"
                with open(duplicate_ids_file, "w") as f:
                    json.dump(
                        {
                            "duplicate_ids": results["id"]["duplicate_ids_list"],
                            "total_duplicates": len(results["id"]["duplicate_ids_list"]),
                        },
                        f,
                        indent=2,
                        ensure_ascii=False,
                    )
                logger.info(f"Duplicate IDs saved to: {duplicate_ids_file}")

        # Save combined results (convert sets to lists for JSON)
        results_serializable = {}
        for key, value in results.items():
            if isinstance(value, dict):
                value_copy = value.copy()
                if "id_set" in value_copy:
                    value_copy["id_set"] = (
                        f"<set with {len(value_copy['id_set'])} IDs, saved to id_set_{split}.pkl>"
                    )
                if "duplicate_ids_list" in value_copy:
                    value_copy["duplicate_ids_list"] = (
                        f"<{len(value_copy['duplicate_ids_list'])} duplicates, saved to duplicate_ids_{split}.json>"
                    )
                results_serializable[key] = value_copy
            else:
                results_serializable[key] = value

        results_file = output_dir / "analysis_results.json"
        with open(results_file, "w") as f:
            json.dump(results_serializable, f, indent=2, ensure_ascii=False)

        logger.info(f"Results saved to: {results_file}")

        return results

    def analyze_overlap_from_results(self, train_results: Dict, val_results: Dict) -> Dict[str, Any]:
        """Analyze overlap from pre-computed results.

        Args:
            train_results: Training set results
            val_results: Validation set results

        Returns:
            Overlap analysis results
        """
        # Calculate overlap
        id_analyzer = IDAnalyzer(self.output_dir)
        overlap_results = id_analyzer.analyze_overlap(train_results["id"], val_results["id"])

        return overlap_results

    def analyze_sft_from_results(self, dataset_name: str, train_results: Dict) -> Dict[str, Any]:
        """Analyze SFT from pre-computed train results.

        Args:
            dataset_name: Name of dataset
            train_results: Training set results

        Returns:
            SFT analysis results
        """
        dataset_config = self.config["datasets"][dataset_name]

        if "sft_ids" not in dataset_config:
            logger.warning(f"No SFT IDs configured for {dataset_name}")
            return {}

        sft_file = Path(dataset_config["sft_ids"])

        if not sft_file.exists():
            logger.error(f"SFT file not found: {sft_file}")
            return {}

        # Get training IDs
        train_ids = train_results["id"].get("id_set", set())

        # Analyze SFT
        analyzer = SFTAnalyzer(self.output_dir)
        results = analyzer.analyze_sft_file(sft_file, train_ids)

        return results

    def analyze_overlap(self, dataset_name: str) -> Dict[str, Any]:
        """Analyze overlap between train and validation sets.

        Args:
            dataset_name: Name of dataset

        Returns:
            Overlap analysis results
        """
        logger.info(f"Analyzing train/val overlap for {dataset_name}")

        # Analyze both sets
        train_results = self.analyze_dataset(dataset_name, "train", analyzers=["id"])
        val_results = self.analyze_dataset(dataset_name, "val", analyzers=["id"])

        # Calculate overlap
        id_analyzer = IDAnalyzer(self.output_dir)
        overlap_results = id_analyzer.analyze_overlap(train_results["id"], val_results["id"])

        # Save results
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = self.output_dir / f"{dataset_name}_overlap_{timestamp}.json"
        with open(output_file, "w") as f:
            json.dump(overlap_results, f, indent=2, ensure_ascii=False)

        logger.info(f"Overlap analysis saved to: {output_file}")

        return overlap_results

    def run_full_analysis(self, dataset_name: str) -> Dict[str, Any]:
        """Run complete analysis for a dataset.

        Args:
            dataset_name: Name of dataset

        Returns:
            Complete analysis results
        """
        logger.info(f"Running full analysis for {dataset_name}")

        # Create shared timestamp and output directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Include filter name in output directory if filtering is active
        if self.s3_category_filter:
            filter_suffix = "_".join(self.s3_category_filter)
            shared_run_dir = self.output_dir / f"{dataset_name}_full_{filter_suffix}_{timestamp}"
        else:
            shared_run_dir = self.output_dir / f"{dataset_name}_full_{timestamp}"

        shared_run_dir.mkdir(parents=True, exist_ok=True)

        results = {"dataset": dataset_name, "timestamp": timestamp, "output_dir": str(shared_run_dir)}

        # Analyze training set
        logger.info("\nAnalyzing training set...")
        train_dir = shared_run_dir / "train"
        train_dir.mkdir(exist_ok=True)
        train_results = self.analyze_dataset_to_dir(dataset_name, "train", train_dir)
        results["train"] = f"See {train_dir}"

        # Analyze validation set
        logger.info("\nAnalyzing validation set...")
        val_dir = shared_run_dir / "val"
        val_dir.mkdir(exist_ok=True)
        val_results = self.analyze_dataset_to_dir(dataset_name, "val", val_dir)
        results["val"] = f"See {val_dir}"

        # Analyze overlap
        logger.info("\nAnalyzing train/val overlap...")
        overlap_results = self.analyze_overlap_from_results(train_results, val_results)
        overlap_file = shared_run_dir / "overlap_analysis.json"
        with open(overlap_file, "w") as f:
            json.dump(overlap_results, f, indent=2, ensure_ascii=False)
        results["overlap"] = f"See {overlap_file}"

        # Analyze SFT
        logger.info("\nAnalyzing SFT IDs...")
        sft_results = self.analyze_sft_from_results(dataset_name, train_results)
        sft_file = shared_run_dir / "sft_analysis.json"
        with open(sft_file, "w") as f:
            json.dump(sft_results, f, indent=2, ensure_ascii=False)
        results["sft"] = f"See {sft_file}"

        # Save summary (skip full results to avoid crashes)
        summary_file = shared_run_dir / "summary.json"
        with open(summary_file, "w") as f:
            json.dump(results, f, indent=2, ensure_ascii=False)

        logger.info(f"\nFull analysis complete!")
        logger.info(f"Results saved to: {shared_run_dir}")
        logger.info(f"Summary: {summary_file}")

        return results


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Run data monitoring pipeline")
    parser.add_argument("--dataset", required=True, help="Dataset version (e.g., v9)")
    parser.add_argument(
        "--split", choices=["train", "val", "both"], default="both", help="Dataset split to analyze"
    )
    parser.add_argument("--analyzers", nargs="+", help="Specific analyzers to run (default: all)")
    parser.add_argument("--no-cache", action="store_true", help="Disable cache usage")
    parser.add_argument("--clear-cache", action="store_true", help="Clear cache before running")
    parser.add_argument(
        "--full", action="store_true", help="Run complete analysis including overlap and SFT"
    )
    parser.add_argument(
        "--s3-filter",
        nargs="+",
        help="Filter records by s3_filepath keywords (e.g., --s3-filter discogs pond5)",
    )

    args = parser.parse_args()

    # Initialize pipeline
    pipeline = DataMonitorPipeline(s3_category_filter=args.s3_filter)

    # Clear cache if requested
    if args.clear_cache:
        logger.info("Clearing cache...")
        pipeline.cache_manager.clear_cache()

    # Run analysis
    use_cache = not args.no_cache

    if args.full:
        results = pipeline.run_full_analysis(args.dataset)
    elif args.split == "both":
        train_results = pipeline.analyze_dataset(args.dataset, "train", args.analyzers, use_cache)
        val_results = pipeline.analyze_dataset(args.dataset, "val", args.analyzers, use_cache)
        results = {"train": train_results, "val": val_results}
    else:
        results = pipeline.analyze_dataset(args.dataset, args.split, args.analyzers, use_cache)

    logger.info("\nAnalysis complete!")


if __name__ == "__main__":
    main()