#!/usr/bin/env python3
"""
Sort JSON shlokas into correct Charaka Samhita order:
SU → NI → VI → SHA → IND → CHI → KAL → SID
Handles decimal chapters like 1.1, 2.2
"""

import json
import re

def sort_shlokas(data):
    """Sort by section order, then chapter (including decimals), then verse"""

    section_order = {
        'SU': 0, 'NI': 1, 'VI': 2, 'SHA': 3,
        'IND': 4, 'IN': 4, 'CI': 5, 'CHI': 5,
        'KA': 6, 'KAL': 6, 'SI': 7, 'SID': 7
    }

    def parse_chapter(chap_str):
        """Parse chapter number including decimals: 1, 1.1, 2.2"""
        parts = chap_str.split('.')
        if len(parts) == 1:
            return (int(parts[0]), 0)
        return (int(parts[0]), int(parts[1]))

    def sort_key(item):
        sid = item['shloka_id']
        # Handle decimal chapters like CH_CHI_1.1/7
        match = re.match(r'CH_([A-Z]+)_([\d.]+)/(.*)', sid)
        if not match:
            return (99, (999, 0), 999)

        sec, chap, verse = match.groups()
        sec_idx = section_order.get(sec, 99)

        # Parse chapter as tuple (main, sub)
        chap_tuple = parse_chapter(chap)

        # For verse ranges like 4-5, use first number
        verse_match = re.match(r'(\d+)', verse)
        verse_num = int(verse_match.group(1)) if verse_match else 999

        return (sec_idx, chap_tuple, verse_num)

    return sorted(data, key=sort_key)

def main():
    input_file = "shlokas_export_2026-05-14_FIXED.json"
    output_file = "shlokas_export_2026-05-14_SORTED.json"

    print(f"Loading {input_file}...")
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)

    print(f"Loaded {len(data)} shlokas")

    # Check current order
    current_sections = []
    for s in data:
        sec = s['shloka_id'].split('_')[1]
        if sec not in current_sections:
            current_sections.append(sec)
    print(f"Current order: {current_sections}")

    # Sort
    sorted_data = sort_shlokas(data)

    # Verify new order
    new_sections = []
    for s in sorted_data:
        sec = s['shloka_id'].split('_')[1]
        if sec not in new_sections:
            new_sections.append(sec)
    print(f"New order: {new_sections}")

    # Check first and last
    print(f"\nFirst: {sorted_data[0]['shloka_id']}")
    print(f"Last: {sorted_data[-1]['shloka_id']}")

    # Show transitions between sections
    print("\n--- Section transitions ---")
    prev_sec = None
    for s in sorted_data:
        sec = s['shloka_id'].split('_')[1]
        if sec != prev_sec:
            print(f"{s['shloka_id']} (section: {sec})")
            prev_sec = sec

    # Save
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(sorted_data, f, ensure_ascii=False, indent=2)

    print(f"\nSaved sorted file: {output_file}")

if __name__ == "__main__":
    main()