"""
Checkpoint and resume functionality for wikiq output.

This module handles:
- Finding resume points in existing output (JSONL or Parquet)
- Merging resumed data with existing output (for Parquet, streaming, memory-efficient)
- Checkpoint file management for fast resume point lookup
"""

import json
import os
import sys
from collections import deque

import pyarrow.parquet as pq


def get_checkpoint_path(output_file, partition_namespaces=False):
    """Get the path to the checkpoint file for a given output file.

    For partitioned output, the checkpoint is placed outside the partition directory
    to avoid pyarrow trying to read it as a parquet file. The filename includes
    the output filename to keep it unique per input file (for parallel jobs).
    """
    if partition_namespaces:
        partition_dir = os.path.dirname(output_file)
        output_filename = os.path.basename(output_file)
        parent_dir = os.path.dirname(partition_dir)
        return os.path.join(parent_dir, output_filename + ".checkpoint")
    return str(output_file) + ".checkpoint"


def read_checkpoint(checkpoint_path, partition_namespaces=False):
    """
    Read resume point from checkpoint file if it exists.

    Checkpoint format:
        Single file: {"pageid": 54, "revid": 325} or {"pageid": 54, "revid": 325, "part": 2}
        Partitioned: {"0": {"pageid": 54, "revid": 325, "part": 1}, ...}

    Returns:
        For single files: A tuple (pageid, revid) or (pageid, revid, part), or None if not found.
        For partitioned: A dict mapping namespace -> (pageid, revid, part), or None.
    """
    if not os.path.exists(checkpoint_path):
        return None

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

        if not data:
            return None

        # Single-file format: {"pageid": ..., "revid": ..., "part": ...}
        if "pageid" in data and "revid" in data:
            part = data.get("part", 0)
            if part > 0:
                return (data["pageid"], data["revid"], part)
            return (data["pageid"], data["revid"])

        # Partitioned format: {"0": {"pageid": ..., "revid": ..., "part": ...}, ...}
        result = {}
        for key, value in data.items():
            part = value.get("part", 0)
            result[int(key)] = (value["pageid"], value["revid"], part)

        return result if result else None

    except (json.JSONDecodeError, IOError, KeyError, TypeError) as e:
        print(f"Warning: Could not read checkpoint file {checkpoint_path}: {e}", file=sys.stderr)
        return None


def cleanup_interrupted_resume(output_file, partition_namespaces):
    """
    Merge any leftover .resume_temp files from a previous interrupted run.

    This should be called BEFORE get_resume_point() so the resume point
    is calculated from the merged data.

    Returns:
        None - no temp files found or normal merge completed
        "start_fresh" - both original and temp were corrupted and deleted
    """
    import shutil

    if partition_namespaces:
        partition_dir = os.path.dirname(output_file)
        output_filename = os.path.basename(output_file)
        temp_suffix = ".resume_temp"

        if not os.path.isdir(partition_dir):
            return

        has_old_temp_files = False
        for ns_dir in os.listdir(partition_dir):
            if ns_dir.startswith('namespace='):
                temp_path = os.path.join(partition_dir, ns_dir, output_filename + temp_suffix)
                if os.path.exists(temp_path):
                    has_old_temp_files = True
                    break

        if has_old_temp_files:
            print(f"Found leftover temp files in {partition_dir} from previous interrupted partitioned run, merging first...", file=sys.stderr)
            had_corruption = merge_partitioned_namespaces(partition_dir, temp_suffix, output_filename)

            has_valid_data = False
            for ns_dir in os.listdir(partition_dir):
                if ns_dir.startswith('namespace='):
                    ns_path = os.path.join(partition_dir, ns_dir)
                    parquet_files = [f for f in os.listdir(ns_path) if f.endswith('.parquet') and not f.endswith('.resume_temp')]
                    if parquet_files:
                        has_valid_data = True
                        break

            if had_corruption and not has_valid_data:
                checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
                if os.path.exists(checkpoint_path):
                    os.remove(checkpoint_path)
                print("All partitioned files were corrupted, will start fresh.", file=sys.stderr)
                return "start_fresh"

            print("Previous temp files merged successfully.", file=sys.stderr)
    else:
        temp_output_file = output_file + ".resume_temp"
        if os.path.exists(temp_output_file) and not os.path.isdir(temp_output_file):
            print(f"Found leftover temp file {temp_output_file} from previous interrupted run, merging first...", file=sys.stderr)
            merged_path = output_file + ".merged"
            merged = merge_parquet_files(output_file, temp_output_file, merged_path)
            if merged == "original_only":
                os.remove(temp_output_file)
            elif merged == "temp_only":
                if os.path.exists(output_file):
                    os.remove(output_file)
                os.rename(temp_output_file, output_file)
                print("Recovered from temp file (original was corrupted or missing).", file=sys.stderr)
            elif merged == "both_invalid":
                if os.path.exists(output_file):
                    os.remove(output_file)
                if os.path.exists(temp_output_file):
                    os.remove(temp_output_file)
                checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
                if os.path.exists(checkpoint_path):
                    os.remove(checkpoint_path)
                print("Both files were corrupted, will start fresh.", file=sys.stderr)
                return "start_fresh"
            elif merged == "merged":
                os.remove(output_file)
                os.rename(merged_path, output_file)
                os.remove(temp_output_file)
                print("Previous temp file merged successfully.", file=sys.stderr)
            else:
                os.remove(temp_output_file)


def get_jsonl_resume_point(output_file, input_file=None):
    """Get resume point from last complete line of JSONL file.

    For .jsonl.d directories, derives the file path from input_file using get_output_filename.
    """
    # Handle .jsonl.d directory output
    if output_file.endswith('.jsonl.d'):
        if input_file is None:
            return None
        if os.path.isdir(output_file):
            # Import here to avoid circular import
            from wikiq import get_output_filename
            jsonl_filename = os.path.basename(get_output_filename(input_file, 'jsonl'))
            output_file = os.path.join(output_file, jsonl_filename)
        else:
            return None

    if not os.path.exists(output_file):
        return None

    try:
        with open(output_file) as f:
            # Stream through file, keeping only last 2 lines in memory
            for line in reversed(deque(f, maxlen=2)):
                try:
                    record = json.loads(line)
                    return (record['articleid'], record['revid'])
                except (json.JSONDecodeError, KeyError):
                    continue
        return None
    except IOError as e:
        print(f"Warning: Could not read {output_file}: {e}", file=sys.stderr)
        return None


def get_resume_point(output_file, partition_namespaces=False, input_file=None):
    """
    Find the resume point(s) from existing output.

    For JSONL: reads last line of file (no checkpoint needed).
    For Parquet: checks checkpoint file, falls back to scanning parquet.

    Args:
        output_file: Path to the output file.
        partition_namespaces: Whether the output uses namespace partitioning.
        input_file: Path to input file (needed for .jsonl.d directory output).

    Returns:
        For single files: A tuple (pageid, revid) or (pageid, revid, part), or None.
        For partitioned: A dict mapping namespace -> (pageid, revid, part), or None.
    """
    # For JSONL, read resume point directly from last line (no checkpoint needed)
    if output_file.endswith('.jsonl') or output_file.endswith('.jsonl.d'):
        result = get_jsonl_resume_point(output_file, input_file)
        if result:
            print(f"Resume point found from JSONL: pageid={result[0]}, revid={result[1]}", file=sys.stderr)
        return result

    # For Parquet, use checkpoint file (fast)
    checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
    checkpoint_result = read_checkpoint(checkpoint_path, partition_namespaces)
    if checkpoint_result is not None:
        print(f"Resume point found in checkpoint file {checkpoint_path}", file=sys.stderr)
        return checkpoint_result

    # Fall back to scanning parquet (slow, for backwards compatibility)
    print(f"No checkpoint file found at {checkpoint_path}, scanning parquet output...", file=sys.stderr)
    try:
        if partition_namespaces:
            return _get_resume_point_partitioned(output_file)
        else:
            return _get_resume_point_single_file(output_file)
    except Exception as e:
        print(f"Error reading resume point from {output_file}: {e}", file=sys.stderr)
        return None


def _get_last_row_resume_point(pq_path):
    """Get resume point by reading only the last row group of a parquet file."""
    pf = pq.ParquetFile(pq_path)
    if pf.metadata.num_row_groups == 0:
        return None

    last_rg_idx = pf.metadata.num_row_groups - 1
    table = pf.read_row_group(last_rg_idx, columns=['articleid', 'revid'])
    if table.num_rows == 0:
        return None

    max_pageid = table['articleid'][-1].as_py()
    max_revid = table['revid'][-1].as_py()
    return (max_pageid, max_revid, 0)


def _get_resume_point_partitioned(output_file):
    """Find per-namespace resume points from partitioned output."""
    partition_dir = os.path.dirname(output_file)
    output_filename = os.path.basename(output_file)

    if not os.path.exists(partition_dir) or not os.path.isdir(partition_dir):
        return None

    namespace_dirs = [d for d in os.listdir(partition_dir) if d.startswith('namespace=')]
    if not namespace_dirs:
        return None

    resume_points = {}
    for ns_dir in namespace_dirs:
        ns = int(ns_dir.split('=')[1])
        pq_path = os.path.join(partition_dir, ns_dir, output_filename)

        if not os.path.exists(pq_path):
            continue

        try:
            result = _get_last_row_resume_point(pq_path)
            if result is not None:
                resume_points[ns] = result
        except Exception as e:
            print(f"Warning: Could not read {pq_path}: {e}", file=sys.stderr)
            continue

    return resume_points if resume_points else None


def _get_resume_point_single_file(output_file):
    """Find resume point from a single parquet file."""
    if not os.path.exists(output_file):
        return None

    if os.path.isdir(output_file):
        return None

    return _get_last_row_resume_point(output_file)


def merge_parquet_files(original_path, temp_path, merged_path):
    """
    Merge two parquet files by streaming row groups.

    Returns:
        "merged" - merged file was created from both sources
        "original_only" - temp was invalid, keep original unchanged
        "temp_only" - original was corrupted but temp is valid
        "both_invalid" - both files invalid
        False - both files were valid but empty
    """
    original_valid = False
    temp_valid = False
    original_pq = None
    temp_pq = None

    try:
        original_pq = pq.ParquetFile(original_path)
        original_valid = True
    except Exception as e:
        print(f"Warning: Original file {original_path} is corrupted or invalid: {e}", file=sys.stderr)

    try:
        if not os.path.exists(temp_path):
            print(f"Note: Temp file {temp_path} does not exist", file=sys.stderr)
        else:
            temp_pq = pq.ParquetFile(temp_path)
            temp_valid = True
    except Exception:
        print(f"Note: No new data in temp file {temp_path}", file=sys.stderr)

    if not original_valid and not temp_valid:
        print(f"Both original and temp files are invalid, will start fresh", file=sys.stderr)
        return "both_invalid"

    if not original_valid and temp_valid:
        print(f"Original file corrupted but temp file is valid, recovering from temp", file=sys.stderr)
        return "temp_only"

    if original_valid and not temp_valid:
        return "original_only"

    merged_writer = None

    for i in range(original_pq.num_row_groups):
        row_group = original_pq.read_row_group(i)
        if merged_writer is None:
            merged_writer = pq.ParquetWriter(
                merged_path,
                row_group.schema,
                flavor="spark"
            )
        merged_writer.write_table(row_group)

    for i in range(temp_pq.num_row_groups):
        row_group = temp_pq.read_row_group(i)
        if merged_writer is None:
            merged_writer = pq.ParquetWriter(
                merged_path,
                row_group.schema,
                flavor="spark"
            )
        merged_writer.write_table(row_group)

    if merged_writer is not None:
        merged_writer.close()
        return "merged"
    return False


def merge_partitioned_namespaces(partition_dir, temp_suffix, file_filter):
    """
    Merge partitioned namespace directories after resume.

    Returns:
        True if at least one namespace has valid data after merge
        False if all namespaces ended up with corrupted/deleted data
    """
    namespace_dirs = [d for d in os.listdir(partition_dir) if d.startswith('namespace=')]
    had_corruption = False
    expected_temp = file_filter + temp_suffix

    for ns_dir in namespace_dirs:
        ns_path = os.path.join(partition_dir, ns_dir)
        temp_path = os.path.join(ns_path, expected_temp)

        if not os.path.exists(temp_path):
            continue

        original_file = file_filter
        original_path = os.path.join(ns_path, original_file)

        if os.path.exists(original_path):
            merged_path = original_path + ".merged"
            merged = merge_parquet_files(original_path, temp_path, merged_path)

            if merged == "original_only":
                if os.path.exists(temp_path):
                    os.remove(temp_path)
            elif merged == "temp_only":
                os.remove(original_path)
                os.rename(temp_path, original_path)
            elif merged == "both_invalid":
                if os.path.exists(original_path):
                    os.remove(original_path)
                if os.path.exists(temp_path):
                    os.remove(temp_path)
                had_corruption = True
            elif merged == "merged":
                os.remove(original_path)
                os.rename(merged_path, original_path)
                if os.path.exists(temp_path):
                    os.remove(temp_path)
            else:
                if os.path.exists(original_path):
                    os.remove(original_path)
                if os.path.exists(temp_path):
                    os.remove(temp_path)
        else:
            try:
                pq.ParquetFile(temp_path)
                os.rename(temp_path, original_path)
            except Exception:
                if os.path.exists(temp_path):
                    os.remove(temp_path)
                had_corruption = True

    return had_corruption


def finalize_resume_merge(
    original_output_file,
    temp_output_file,
    partition_namespaces,
    original_partition_dir
):
    """
    Finalize the resume by merging temp output with original output.
    """
    import shutil

    print("Merging resumed data with existing output...", file=sys.stderr)
    try:
        if partition_namespaces and original_partition_dir is not None:
            file_filter = os.path.basename(original_output_file)
            merge_partitioned_namespaces(original_partition_dir, ".resume_temp", file_filter)
            if os.path.exists(temp_output_file) and os.path.isdir(temp_output_file):
                shutil.rmtree(temp_output_file)
        else:
            merged_output_file = original_output_file + ".merged"
            merged = merge_parquet_files(original_output_file, temp_output_file, merged_output_file)

            if merged == "original_only":
                if os.path.exists(temp_output_file):
                    os.remove(temp_output_file)
            elif merged == "temp_only":
                os.remove(original_output_file)
                os.rename(temp_output_file, original_output_file)
            elif merged == "both_invalid":
                os.remove(original_output_file)
                if os.path.exists(temp_output_file):
                    os.remove(temp_output_file)
            elif merged == "merged":
                os.remove(original_output_file)
                os.rename(merged_output_file, original_output_file)
                if os.path.exists(temp_output_file):
                    os.remove(temp_output_file)
            else:
                os.remove(original_output_file)
                if os.path.exists(temp_output_file):
                    os.remove(temp_output_file)

        print("Merge complete.", file=sys.stderr)
    except Exception as e:
        print(f"Error merging resume data for {original_output_file}: {e}", file=sys.stderr)
        print(f"New data saved in: {temp_output_file}", file=sys.stderr)
        raise


def setup_resume_temp_output(output_file, partition_namespaces):
    """
    Set up temp output for resume mode (Parquet only).

    Returns:
        Tuple of (original_output_file, temp_output_file, original_partition_dir)
        or (None, None, None) if no existing output to resume from.
    """
    import shutil

    original_output_file = None
    temp_output_file = None
    original_partition_dir = None

    if partition_namespaces:
        partition_dir = os.path.dirname(output_file)
        output_filename = os.path.basename(output_file)
        output_exists = False
        if os.path.isdir(partition_dir):
            for d in os.listdir(partition_dir):
                if d.startswith('namespace='):
                    if os.path.exists(os.path.join(partition_dir, d, output_filename)):
                        output_exists = True
                        break
        if output_exists:
            original_partition_dir = partition_dir
    else:
        output_exists = isinstance(output_file, str) and os.path.exists(output_file)

    if output_exists:
        original_output_file = output_file
        temp_output_file = output_file + ".resume_temp"

        if os.path.exists(temp_output_file):
            if os.path.isdir(temp_output_file):
                shutil.rmtree(temp_output_file)
            else:
                os.remove(temp_output_file)

        if partition_namespaces:
            os.makedirs(temp_output_file, exist_ok=True)

    return original_output_file, temp_output_file, original_partition_dir