#!/usr/bin/env python3
"""
Voice Designer Captioning Progress Monitor
Monitors captioning progress by analyzing output JSONL and log files
"""

import os
import json
import re
import time
import argparse
from datetime import datetime, timedelta
from pathlib import Path
from collections import defaultdict, Counter
import statistics


class CaptioningProgressMonitor:
    def __init__(self, output_jsonl, log_file=None, total_records=None):
        self.output_jsonl = output_jsonl
        self.log_file = log_file
        self.total_records = total_records
        self.start_time = None
        self.last_update = None

    def parse_log_file(self):
        """Parse log file for timing and error information"""
        if not self.log_file or not os.path.exists(self.log_file):
            return {}

        log_data = {
            "start_time": None,
            "worker_times": defaultdict(list),
            "errors": [],
            "api_errors": 0,
            "validation_errors": 0,
            "exhausted_attempts": 0,
            "batches_completed": 0,
            "last_activity": None,
        }

        try:
            with open(self.log_file, "r") as f:
                for line in f:
                    line = line.strip()

                    # Extract timestamp
                    timestamp_match = re.match(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})", line)
                    if timestamp_match:
                        timestamp_str = timestamp_match.group(1)
                        timestamp = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
                        log_data["last_activity"] = timestamp

                        if log_data["start_time"] is None:
                            log_data["start_time"] = timestamp

                    # Track worker processing times
                    if "Worker" in line and "completed" in line:
                        worker_match = re.search(r"Worker (\d+)", line)
                        if worker_match and timestamp_match:
                            worker_id = int(worker_match.group(1))
                            log_data["worker_times"][worker_id].append(timestamp)

                    # Track errors
                    if "❌" in line or "ERROR" in line:
                        log_data["errors"].append(line)
                        if "API error" in line:
                            log_data["api_errors"] += 1
                        elif "Validation failed" in line:
                            log_data["validation_errors"] += 1
                        elif "Exhausted all" in line:
                            log_data["exhausted_attempts"] += 1

                    # Track batch completion
                    if "Saved batch" in line:
                        batch_match = re.search(r"batch (\d+)", line)
                        if batch_match:
                            log_data["batches_completed"] = max(
                                log_data["batches_completed"], int(batch_match.group(1))
                            )

        except Exception as e:
            print(f"Error parsing log file: {e}")

        return log_data

    def analyze_jsonl_output(self):
        """Analyze the output JSONL file for completion statistics"""
        if not os.path.exists(self.output_jsonl):
            return {
                "total_processed": 0,
                "records_with_captions": 0,
                "records_without_captions": 0,
                "records_with_empty_captions": 0,
                "records_with_errors": 0,
                "success_rate": 0.0,
                "keyword_counts": [],
                "avg_keywords": 0.0,
                "caption_types": Counter(),
                "file_size": 0,
            }

        stats = {
            "total_processed": 0,
            "records_with_captions": 0,
            "records_without_captions": 0,
            "records_with_empty_captions": 0,
            "records_with_errors": 0,
            "keyword_counts": [],
            "caption_types": Counter(),
            "file_size": os.path.getsize(self.output_jsonl),
        }

        try:
            with open(self.output_jsonl, "r") as f:
                for line_num, line in enumerate(f, 1):
                    if not line.strip():
                        continue

                    try:
                        record = json.loads(line.strip())

                        # Check for captions
                        stems_captions = record.get("stems_captions", {})
                        has_valid_captions = False
                        has_empty_captions = False
                        has_errors = False
                        has_voice_description_keywords = False

                        for stem_name, captions in stems_captions.items():
                            if captions:
                                for caption_obj in captions:
                                    caption_text = caption_obj.get("caption", "")
                                    error = caption_obj.get("error")
                                    prompt_type = caption_obj.get("prompt_type", "unknown")

                                    # Only consider "voice_description_keywords" captions
                                    if prompt_type == "voice_description_keywords":
                                        has_voice_description_keywords = True
                                        # Caption is valid if it's >= 10 characters and no error
                                        if (
                                            caption_text
                                            and len(caption_text.strip()) >= 10
                                            and not error
                                        ):
                                            has_valid_captions = True
                                            # Count keywords
                                            keywords = [
                                                kw.strip()
                                                for kw in caption_text.split(",")
                                                if kw.strip()
                                            ]
                                            stats["keyword_counts"].append(len(keywords))

                                            # Track caption type
                                            stats["caption_types"][prompt_type] += 1
                                        elif prompt_type == "voice_description_keywords" and (
                                            not caption_text or len(caption_text.strip()) < 10
                                        ):
                                            # Record has voice_description_keywords prompt but empty/short caption
                                            has_empty_captions = True

                                        if error:
                                            has_errors = True

                        # Only count records that have voice_description_keywords prompt type
                        if has_voice_description_keywords:
                            stats["total_processed"] += 1

                            if has_valid_captions:
                                stats["records_with_captions"] += 1
                            elif has_empty_captions:
                                stats["records_with_empty_captions"] += 1
                            else:
                                stats["records_without_captions"] += 1
                            if has_errors:
                                stats["records_with_errors"] += 1

                    except json.JSONDecodeError as e:
                        print(f"JSON decode error on line {line_num}: {e}")
                        continue

        except Exception as e:
            print(f"Error reading output file: {e}")

        # Calculate averages
        if stats["keyword_counts"]:
            stats["avg_keywords"] = statistics.mean(stats["keyword_counts"])
            stats["median_keywords"] = statistics.median(stats["keyword_counts"])
            stats["min_keywords"] = min(stats["keyword_counts"])
            stats["max_keywords"] = max(stats["keyword_counts"])
        else:
            stats["avg_keywords"] = 0.0
            stats["median_keywords"] = 0.0
            stats["min_keywords"] = 0
            stats["max_keywords"] = 0

        stats["success_rate"] = (
            (stats["records_with_captions"] / stats["total_processed"] * 100)
            if stats["total_processed"] > 0
            else 0.0
        )

        return stats

    def calculate_processing_speed(self, log_data, jsonl_stats):
        """Calculate processing speed and time estimates"""
        speed_data = {
            "records_per_minute": 0.0,
            "estimated_finish_time": None,
            "elapsed_time": None,
            "remaining_time": None,
            "worker_performance": {},
        }

        if log_data.get("start_time") and jsonl_stats["total_processed"] > 0:
            current_time = datetime.now()
            elapsed = current_time - log_data["start_time"]
            speed_data["elapsed_time"] = elapsed

            # Calculate overall speed
            elapsed_minutes = elapsed.total_seconds() / 60
            if elapsed_minutes > 0:
                speed_data["records_per_minute"] = jsonl_stats["total_processed"] / elapsed_minutes

            # Estimate completion time
            if self.total_records and speed_data["records_per_minute"] > 0:
                remaining_records = self.total_records - jsonl_stats["total_processed"]
                remaining_minutes = remaining_records / speed_data["records_per_minute"]
                speed_data["remaining_time"] = timedelta(minutes=remaining_minutes)
                speed_data["estimated_finish_time"] = current_time + speed_data["remaining_time"]

            # Analyze worker performance
            for worker_id, timestamps in log_data["worker_times"].items():
                if len(timestamps) > 1:
                    worker_elapsed = timestamps[-1] - timestamps[0]
                    worker_records = len(timestamps)
                    worker_minutes = worker_elapsed.total_seconds() / 60
                    if worker_minutes > 0:
                        speed_data["worker_performance"][worker_id] = worker_records / worker_minutes

        return speed_data

    def format_time_duration(self, td):
        """Format timedelta as human-readable string"""
        if not td:
            return "Unknown"

        total_seconds = int(td.total_seconds())
        days = total_seconds // 86400
        hours = (total_seconds % 86400) // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60

        parts = []
        if days > 0:
            parts.append(f"{days}d")
        if hours > 0:
            parts.append(f"{hours}h")
        if minutes > 0:
            parts.append(f"{minutes}m")
        if seconds > 0 or not parts:
            parts.append(f"{seconds}s")

        return " ".join(parts)

    def generate_report(self):
        """Generate comprehensive progress report"""
        print("=" * 80)
        print("🎵 VOICE DESIGNER CAPTIONING PROGRESS MONITOR")
        print("=" * 80)
        print(f"📁 Output file: {self.output_jsonl}")
        if self.log_file:
            print(f"📋 Log file: {self.log_file}")
        if self.total_records:
            print(f"📊 Total records: {self.total_records:,}")
        print(f"⏰ Report time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print()

        # Parse data
        log_data = self.parse_log_file()
        jsonl_stats = self.analyze_jsonl_output()
        speed_data = self.calculate_processing_speed(log_data, jsonl_stats)

        # Processing Status
        print("📈 PROCESSING STATUS")
        print("-" * 40)
        print(f"Records processed:     {jsonl_stats['total_processed']:,}")
        print(f"Records with captions: {jsonl_stats['records_with_captions']:,}")
        print(f"Records with empty captions: {jsonl_stats['records_with_empty_captions']:,}")
        print(f"Records without captions: {jsonl_stats['records_without_captions']:,}")
        print(f"Records with errors:   {jsonl_stats['records_with_errors']:,}")
        print(f"Success rate:          {jsonl_stats['success_rate']:.1f}%")

        if self.total_records:
            completion_pct = (jsonl_stats["total_processed"] / self.total_records) * 100
            print(f"Overall completion:    {completion_pct:.1f}%")

        print(f"Output file size:      {jsonl_stats['file_size'] / (1024*1024):.1f} MB")
        print()

        # Keyword Statistics
        if jsonl_stats["keyword_counts"]:
            print("🎯 KEYWORD STATISTICS")
            print("-" * 40)
            print(f"Average keywords:      {jsonl_stats['avg_keywords']:.1f}")
            print(f"Median keywords:       {jsonl_stats['median_keywords']:.1f}")
            print(f"Min keywords:          {jsonl_stats['min_keywords']}")
            print(f"Max keywords:          {jsonl_stats['max_keywords']}")

            # Keyword distribution
            keyword_ranges = Counter()
            for count in jsonl_stats["keyword_counts"]:
                if count < 20:
                    keyword_ranges["< 20"] += 1
                elif count < 25:
                    keyword_ranges["20-24"] += 1
                elif count < 30:
                    keyword_ranges["25-29"] += 1
                elif count < 40:
                    keyword_ranges["30-39"] += 1
                else:
                    keyword_ranges["40+"] += 1

            print("Keyword count ranges:")
            for range_name, count in sorted(keyword_ranges.items()):
                pct = (count / len(jsonl_stats["keyword_counts"])) * 100
                print(f"  {range_name:>8}: {count:>6} ({pct:4.1f}%)")
            print()

        # Timing Information
        if speed_data["elapsed_time"]:
            print("⏱️  TIMING & PERFORMANCE")
            print("-" * 40)
            print(f"Elapsed time:          {self.format_time_duration(speed_data['elapsed_time'])}")
            print(f"Processing speed:      {speed_data['records_per_minute']:.2f} records/min")

            if speed_data["remaining_time"]:
                print(
                    f"Estimated remaining:   {self.format_time_duration(speed_data['remaining_time'])}"
                )
                print(
                    f"Estimated finish:      {speed_data['estimated_finish_time'].strftime('%Y-%m-%d %H:%M:%S')}"
                )

            print()

            # Worker performance
            if speed_data["worker_performance"]:
                print("👥 WORKER PERFORMANCE")
                print("-" * 40)
                for worker_id, rate in sorted(speed_data["worker_performance"].items()):
                    print(f"Worker {worker_id:>2}: {rate:>8.2f} records/min")
                print()
        elif jsonl_stats["total_processed"] > 0:
            # Fallback timing calculation based on output file modification time
            try:
                import os

                current_time = datetime.now()
                file_creation_time = datetime.fromtimestamp(os.path.getctime(self.output_jsonl))
                elapsed = current_time - file_creation_time
                elapsed_minutes = elapsed.total_seconds() / 60

                if elapsed_minutes > 0:
                    records_per_minute = jsonl_stats["total_processed"] / elapsed_minutes

                    print("⏱️  ESTIMATED TIMING & PERFORMANCE")
                    print("-" * 40)
                    print(f"File age:              {self.format_time_duration(elapsed)}")
                    print(f"Est. processing speed: {records_per_minute:.2f} records/min")

                    if self.total_records:
                        remaining_records = self.total_records - jsonl_stats["total_processed"]
                        if remaining_records > 0 and records_per_minute > 0:
                            remaining_minutes = remaining_records / records_per_minute
                            remaining_time = timedelta(minutes=remaining_minutes)
                            estimated_finish = current_time + remaining_time
                            print(f"Est. remaining time:   {self.format_time_duration(remaining_time)}")
                            print(
                                f"Est. finish time:      {estimated_finish.strftime('%Y-%m-%d %H:%M:%S')}"
                            )

                    print()
            except Exception as e:
                # If file timing fails, skip timing section
                pass

        # Error Analysis
        if log_data.get("errors"):
            print("❌ ERROR ANALYSIS")
            print("-" * 40)
            print(f"API errors:            {log_data.get('api_errors', 0)}")
            print(f"Validation errors:     {log_data.get('validation_errors', 0)}")
            print(f"Exhausted attempts:    {log_data.get('exhausted_attempts', 0)}")
            print(f"Total error events:    {len(log_data.get('errors', []))}")

            if log_data.get("api_errors", 0) > 0:
                error_rate = (
                    (log_data["api_errors"] / jsonl_stats["total_processed"]) * 100
                    if jsonl_stats["total_processed"] > 0
                    else 0
                )
                print(f"API error rate:        {error_rate:.2f}%")
            print()

        # Recent Activity
        if log_data.get("last_activity"):
            time_since_activity = datetime.now() - log_data["last_activity"]
            print("🔄 RECENT ACTIVITY")
            print("-" * 40)
            print(f"Last log entry:        {log_data['last_activity'].strftime('%Y-%m-%d %H:%M:%S')}")
            print(f"Time since activity:   {self.format_time_duration(time_since_activity)}")

            if log_data.get("batches_completed", 0) > 0:
                print(f"Batches completed:     {log_data['batches_completed']}")
            print()

        # Caption Types
        if jsonl_stats["caption_types"]:
            print("📝 CAPTION TYPES")
            print("-" * 40)
            for caption_type, count in jsonl_stats["caption_types"].most_common():
                print(f"{caption_type:>20}: {count:>6}")
            print()

        print("=" * 80)


def main():
    parser = argparse.ArgumentParser(description="Monitor voice designer captioning progress")
    parser.add_argument("--output-jsonl", required=True, help="Path to output JSONL file")
    parser.add_argument("--log-file", help="Path to log file")
    parser.add_argument("--total-records", type=int, help="Total number of records to process")
    parser.add_argument(
        "--watch", "-w", action="store_true", help="Watch mode - refresh every 30 seconds"
    )
    parser.add_argument(
        "--interval", type=int, default=30, help="Refresh interval in seconds (default: 30)"
    )

    args = parser.parse_args()

    monitor = CaptioningProgressMonitor(
        output_jsonl=args.output_jsonl, log_file=args.log_file, total_records=args.total_records
    )

    if args.watch:
        try:
            while True:
                # Clear screen (works on most terminals)
                os.system("clear" if os.name == "posix" else "cls")
                monitor.generate_report()
                print(f"🔄 Refreshing in {args.interval} seconds... (Ctrl+C to exit)")
                time.sleep(args.interval)
        except KeyboardInterrupt:
            print("\n👋 Monitoring stopped.")
    else:
        monitor.generate_report()


if __name__ == "__main__":
    main()