consolidate_data.py

"""Data consolidation and validation for HabitHub calendar extraction.

This module handles merging data from all months, validating quality,
and exporting to CSV format.
"""

import csv
import json
import os
import calendar
from collections import Counter
from datetime import datetime
from typing import List, Dict, Any

import config


class DataConsolidator:
    """Handles consolidation and validation of extracted habit data."""

    def __init__(self, extracted_data_path: str):
        """Initialize the consolidator with extracted data.

        Args:
            extracted_data_path: Path to extracted_data.json file
        """
        self.extracted_data_path = extracted_data_path
        self.all_days: List[Dict[str, str]] = []
        self.issues: List[str] = []

    def load_data(self) -> bool:
        """Load extracted data from JSON file.

        Returns:
            True if successful, False otherwise
        """
        print(f"Loading extracted data from: {self.extracted_data_path}")

        if not os.path.exists(self.extracted_data_path):
            print(f"✗ Error: File not found: {self.extracted_data_path}")
            return False

        try:
            with open(self.extracted_data_path, 'r') as f:
                data = json.load(f)

            if "months" not in data:
                print("✗ Error: Invalid data format - 'months' key not found")
                return False

            print(f"✓ Loaded data for {len(data['months'])} month(s)")
            self.raw_data = data
            return True

        except json.JSONDecodeError as e:
            print(f"✗ Error: Failed to parse JSON: {str(e)}")
            return False
        except Exception as e:
            print(f"✗ Error: {str(e)}")
            return False

    def merge_all_months(self) -> List[Dict[str, str]]:
        """Merge data from all months into a single list.

        Returns:
            List of {date, status} dictionaries
        """
        print("\nMerging data from all months...")

        all_days = []

        for month_data in self.raw_data["months"]:
            month_name = month_data.get("month", "Unknown")
            year = month_data.get("year", "Unknown")
            days = month_data.get("data", [])

            print(f"  Processing {month_name} {year}: {len(days)} day(s)")

            for day in days:
                all_days.append({
                    "date": day["date"],
                    "status": day["status"]
                })

        # Sort by date
        all_days.sort(key=lambda x: x["date"])

        print(f"✓ Merged {len(all_days)} total days")
        self.all_days = all_days
        return all_days

    def validate_data(self) -> bool:
        """Validate data quality and check for issues.

        Returns:
            True if validation passes, False if critical issues found
        """
        print("\nValidating data quality...")

        self.issues = []
        critical_issues = 0

        # Check 1: Date format validation
        print("  Checking date formats...")
        invalid_dates = []
        for day in self.all_days:
            date_str = day["date"]
            try:
                datetime.strptime(date_str, "%Y-%m-%d")
            except ValueError:
                invalid_dates.append(date_str)

        if invalid_dates:
            issue = f"Invalid date format: {len(invalid_dates)} dates"
            self.issues.append(issue)
            print(f"    ✗ {issue}")
            print(f"      Examples: {invalid_dates[:3]}")
            critical_issues += 1
        else:
            print("    ✓ All dates have valid format")

        # Check 2: Status validation
        print("  Checking status values...")
        valid_statuses = {"green", "red", "blue", "gray"}
        invalid_statuses = [
            day for day in self.all_days
            if day["status"] not in valid_statuses
        ]

        if invalid_statuses:
            issue = f"Invalid status values: {len(invalid_statuses)} entries"
            self.issues.append(issue)
            print(f"    ✗ {issue}")
            print(f"      Examples: {invalid_statuses[:3]}")
            critical_issues += 1
        else:
            print("    ✓ All status values are valid")

        # Check 3: Month entry count validation
        print("  Checking entry counts per month...")
        month_entry_issues = []

        for month_data in self.raw_data["months"]:
            month_name = month_data.get("month", "Unknown")
            year = month_data.get("year", 0)
            days = month_data.get("data", [])

            # Get expected number of days in the month
            try:
                # Convert month name to month number
                month_num = datetime.strptime(month_name, "%B").month
                expected_days = calendar.monthrange(year, month_num)[1]
                actual_days = len(days)

                if actual_days != expected_days:
                    month_entry_issues.append({
                        "month": f"{month_name} {year}",
                        "expected": expected_days,
                        "actual": actual_days,
                        "missing": expected_days - actual_days
                    })
            except (ValueError, AttributeError):
                # Skip if month name is invalid
                pass

        if month_entry_issues:
            issue = f"Incorrect entry counts: {len(month_entry_issues)} month(s)"
            self.issues.append(issue)
            print(f"    ✗ {issue}")
            for m in month_entry_issues:
                print(f"      {m['month']}: expected {m['expected']}, got {m['actual']} "
                      f"(missing {m['missing']})")
            # Store detailed month entry issues for summary report
            self.month_entry_issues = month_entry_issues
            critical_issues += 1
        else:
            print("    ✓ All months have correct entry counts")
            self.month_entry_issues = []

        # Check 4: Duplicate dates
        print("  Checking for duplicate dates...")
        dates = [day["date"] for day in self.all_days]
        date_counts = Counter(dates)
        duplicates = {date: count for date, count in date_counts.items()
                      if count > 1}

        if duplicates:
            issue = f"Duplicate dates found: {len(duplicates)} dates"
            self.issues.append(issue)
            print(f"    ⚠ {issue}")
            print(f"      Examples: {list(duplicates.items())[:3]}")
            print("    Note: Only the first occurrence will be kept")
        else:
            print("    ✓ No duplicate dates")

        # Check 5: Date continuity (check for large gaps)
        print("  Checking date continuity...")
        if len(self.all_days) >= 2:
            gaps = []
            for i in range(1, len(self.all_days)):
                date1 = datetime.strptime(self.all_days[i-1]["date"],
                                          "%Y-%m-%d")
                date2 = datetime.strptime(self.all_days[i]["date"],
                                          "%Y-%m-%d")
                gap_days = (date2 - date1).days

                # Flag gaps larger than 60 days
                if gap_days > 60:
                    gaps.append((self.all_days[i-1]["date"],
                                self.all_days[i]["date"], gap_days))

            if gaps:
                issue = f"Large gaps found: {len(gaps)} gap(s) > 60 days"
                self.issues.append(issue)
                print(f"    ⚠ {issue}")
                print(f"      Examples: {gaps[:3]}")
            else:
                print("    ✓ No large gaps in date sequence")

        # Summary
        print("\n" + "-" * 60)
        if critical_issues > 0:
            print(f"✗ Validation FAILED: {critical_issues} critical issue(s)")
            return False
        elif self.issues:
            print(f"⚠ Validation passed with {len(self.issues)} warning(s)")
            return True
        else:
            print("✓ Validation passed - no issues found")
            return True

    def remove_duplicates(self) -> List[Dict[str, str]]:
        """Remove duplicate dates, keeping only the first occurrence.

        Returns:
            Deduplicated list of days
        """
        if not self.all_days:
            return []

        print("\nRemoving duplicates...")
        seen_dates = set()
        deduplicated = []

        for day in self.all_days:
            if day["date"] not in seen_dates:
                deduplicated.append(day)
                seen_dates.add(day["date"])

        removed = len(self.all_days) - len(deduplicated)
        if removed > 0:
            print(f"  Removed {removed} duplicate(s)")

        self.all_days = deduplicated
        return deduplicated

    def export_to_csv(self, output_path: str) -> bool:
        """Export consolidated data to CSV file.

        Args:
            output_path: Path to output CSV file

        Returns:
            True if successful, False otherwise
        """
        print(f"\nExporting to CSV: {output_path}")

        try:
            with open(output_path, 'w', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=["date", "status"])
                writer.writeheader()
                writer.writerows(self.all_days)

            print(f"✓ Exported {len(self.all_days)} rows to CSV")
            return True

        except Exception as e:
            print(f"✗ Error writing CSV: {str(e)}")
            return False

    def generate_summary_report(self) -> Dict[str, Any]:
        """Generate a summary report of the extracted data.

        Returns:
            Dictionary containing summary statistics
        """
        print("\nGenerating summary report...")

        if not self.all_days:
            return {}

        # Count status values
        status_counts = Counter(day["status"] for day in self.all_days)

        # Date range
        dates = [datetime.strptime(day["date"], "%Y-%m-%d")
                 for day in self.all_days]
        min_date = min(dates).strftime("%Y-%m-%d")
        max_date = max(dates).strftime("%Y-%m-%d")
        date_range_days = (max(dates) - min(dates)).days + 1

        # Monthly breakdown
        monthly_breakdown = {}
        months_with_gray = []
        months_with_blue = []

        for day in self.all_days:
            date = datetime.strptime(day["date"], "%Y-%m-%d")
            month_key = date.strftime("%Y-%m")
            month_name = date.strftime("%B %Y")

            if month_key not in monthly_breakdown:
                monthly_breakdown[month_key] = {
                    "green": 0, "red": 0, "blue": 0, "gray": 0, "total": 0
                }

            monthly_breakdown[month_key][day["status"]] += 1
            monthly_breakdown[month_key]["total"] += 1

            # Track months with gray or blue entries
            if day["status"] == "gray" and month_name not in months_with_gray:
                months_with_gray.append(month_name)
            if day["status"] == "blue" and month_name not in months_with_blue:
                months_with_blue.append(month_name)

        # Calculate success rate
        total_days = len(self.all_days)
        green_days = status_counts.get("green", 0)
        success_rate = (green_days / total_days * 100) if total_days > 0 else 0

        summary = {
            "total_days": total_days,
            "green_days": status_counts.get("green", 0),
            "red_days": status_counts.get("red", 0),
            "blue_days": status_counts.get("blue", 0),
            "gray_days": status_counts.get("gray", 0),
            "success_rate": round(success_rate, 2),
            "date_range": {
                "start": min_date,
                "end": max_date,
                "days": date_range_days
            },
            "coverage": round(total_days / date_range_days * 100, 2),
            "monthly_breakdown": monthly_breakdown,
            "months_with_gray": months_with_gray,
            "months_with_blue": months_with_blue,
            "month_entry_validation_failures": getattr(self, 'month_entry_issues', []),
            "validation_issues": self.issues,
            "source_file": self.extracted_data_path
        }

        # Print summary
        print(f"  Total days tracked: {summary['total_days']}")
        print(f"  Green (success): {summary['green_days']} "
              f"({summary['success_rate']}%)")
        print(f"  Red (missed): {summary['red_days']}")
        print(f"  Blue (skipped): {summary['blue_days']}")
        print(f"  Gray (unmarked): {summary['gray_days']}")
        print(f"  Date range: {min_date} to {max_date} "
              f"({date_range_days} days)")
        print(f"  Coverage: {summary['coverage']}% of date range")

        # Print months with gray/blue entries
        if months_with_gray:
            print(f"\n  Months with gray entries ({len(months_with_gray)}):")
            for month in months_with_gray:
                print(f"    - {month}")

        if months_with_blue:
            print(f"\n  Months with blue entries ({len(months_with_blue)}):")
            for month in months_with_blue:
                print(f"    - {month}")

        # Print validation failure details
        if hasattr(self, 'month_entry_issues') and self.month_entry_issues:
            print(f"\n  Month validation failures ({len(self.month_entry_issues)}):")
            for issue in self.month_entry_issues:
                print(f"    - {issue['month']}: expected {issue['expected']} days, "
                      f"got {issue['actual']} days (missing {issue['missing']})")

        return summary

    def save_summary_report(self, summary: Dict[str, Any],
                            output_path: str) -> bool:
        """Save summary report to JSON file.

        Args:
            summary: Summary dictionary
            output_path: Path to output JSON file

        Returns:
            True if successful, False otherwise
        """
        print(f"\nSaving summary report: {output_path}")

        try:
            with open(output_path, 'w') as f:
                json.dump(summary, f, indent=2)

            print("✓ Summary report saved")
            return True

        except Exception as e:
            print(f"✗ Error writing summary: {str(e)}")
            return False

    def run(self, csv_output: str, summary_output: str) -> bool:
        """Run the complete consolidation pipeline.

        Args:
            csv_output: Path for CSV output
            summary_output: Path for summary JSON output

        Returns:
            True if successful, False otherwise
        """
        print("=" * 60)
        print("PHASE 3: Data Consolidation")
        print("=" * 60)

        # Load data
        if not self.load_data():
            return False

        # Merge all months
        self.merge_all_months()

        if not self.all_days:
            print("\n✗ No data to consolidate!")
            return False

        # Remove duplicates
        self.remove_duplicates()

        # Validate
        validation_passed = self.validate_data()

        # Export to CSV
        if not self.export_to_csv(csv_output):
            return False

        # Generate and save summary
        summary = self.generate_summary_report()
        if not self.save_summary_report(summary, summary_output):
            return False

        # Final summary
        print("\n" + "=" * 60)
        print("Phase 3 Complete!")
        print("=" * 60)
        print(f"✓ CSV exported: {csv_output}")
        print(f"✓ Summary saved: {summary_output}")

        if not validation_passed:
            print("\n⚠ Note: Some validation issues were found")
            print("  Review the summary report for details")

        return True


def main():
    """Main entry point for the consolidation script."""
    import sys

    # Check for extracted data
    extracted_data_path = os.path.join(config.OUTPUT_DIR, "extracted_data.json")
    if not os.path.exists(extracted_data_path):
        print(f"✗ Error: Extracted data not found: {extracted_data_path}")
        print("\nPlease run Phase 2 (extract_with_llm.py) first")
        sys.exit(1)

    # Create consolidator
    consolidator = DataConsolidator(extracted_data_path)

    # Run consolidation
    csv_output = config.OUTPUT_CSV
    summary_output = config.OUTPUT_SUMMARY

    success = consolidator.run(csv_output, summary_output)

    if not success:
        print("\n✗ Consolidation failed!")
        sys.exit(1)
    else:
        print("\n✓ Consolidation complete!")
        sys.exit(0)


if __name__ == "__main__":
    main()