"""
json_to_viewer_md.py
--------------------
Converts the structured olmOCR output (.md with HTML tables) into clean
readable prose for the OCR viewer.

Reads the _olmocr.md file produced by ocr_wardiaries.py.
- Diary pages (HTML tables) are flattened to readable prose.
- Non-diary pages (plain text) are passed through unchanged.

Output: a clean .md file with ## Page N markers, no HTML tags,
suitable for display in the OCR viewer.

No API calls. Pure Python. Runs locally in seconds.

Usage:
    python json_to_viewer_md.py --input "G:/path/to/outputs_step1_olmocr/Calgary-Highlanders_War-Diary_Sep44_olmocr.md" --output_dir "G:/path/to/outputs_step2_json_to_viewer_md"
"""

import argparse
import re
import sys
from html.parser import HTMLParser
from pathlib import Path


# ---------------------------------------------------------------------------
# HTML table parser
# ---------------------------------------------------------------------------

class TableParser(HTMLParser):
    """
    Parses a single HTML table into a list of row dicts.
    Each row: { place, date, hour, summary, remarks }
    """

    def __init__(self):
        super().__init__()
        self.rows = []
        self._in_table = False
        self._in_header = False
        self._current_row = None
        self._current_cell = None
        self._cell_index = 0
        self._cell_parts = []

    def handle_starttag(self, tag, attrs):
        tag = tag.lower()
        if tag == 'table':
            self._in_table = True
        elif tag == 'tr' and self._in_table:
            self._current_row = []
            self._cell_index = 0
        elif tag in ('td', 'th') and self._current_row is not None:
            self._in_header = (tag == 'th')
            self._current_cell = []
            self._cell_parts = []
        elif tag == 'br' and self._current_cell is not None:
            self._cell_parts.append('\n')

    def handle_endtag(self, tag):
        tag = tag.lower()
        if tag in ('td', 'th') and self._current_cell is not None:
            cell_text = ''.join(self._cell_parts).strip()
            self._current_row.append(cell_text)
            self._current_cell = None
            self._cell_parts = []
            self._cell_index += 1
        elif tag == 'tr' and self._current_row is not None:
            if not self._in_header and len(self._current_row) >= 4:
                self.rows.append({
                    'place':   self._current_row[0] if len(self._current_row) > 0 else '',
                    'date':    self._current_row[1] if len(self._current_row) > 1 else '',
                    'hour':    self._current_row[2] if len(self._current_row) > 2 else '',
                    'summary': self._current_row[3] if len(self._current_row) > 3 else '',
                    'remarks': self._current_row[4] if len(self._current_row) > 4 else '',
                })
            self._current_row = None
            self._in_header = False
        elif tag == 'table':
            self._in_table = False

    def handle_data(self, data):
        if self._current_cell is not None:
            self._cell_parts.append(data)


def parse_table(html: str) -> list[dict]:
    """Parse an HTML table string into a list of row dicts."""
    parser = TableParser()
    parser.feed(html)
    return parser.rows


# ---------------------------------------------------------------------------
# Text cleaning
# ---------------------------------------------------------------------------

def clean_cell(text: str) -> str:
    """Strip leading/trailing whitespace and collapse multiple newlines to one."""
    text = re.sub(r'\n{2,}', '\n', text)
    return text.strip()


def format_place(place: str) -> str:
    """
    Format the Place column for readability.
    Lines separated by newlines → joined with ' / ' for display.
    Removes blank lines.
    """
    lines = [l.strip() for l in place.split('\n') if l.strip()]
    return ' / '.join(lines)


def format_date(date: str) -> str:
    """Collapse multiline date (e.g. '1 Sep\n44') to single line."""
    parts = [p.strip() for p in date.split('\n') if p.strip()]
    return ' '.join(parts)


def row_to_prose(row: dict) -> str:
    """
    Convert a single diary table row to readable prose for the viewer.

    Format:
        DATE — PLACE
        [HOUR]
        SUMMARY TEXT

        Remarks: REMARKS
    """
    parts = []

    date  = format_date(clean_cell(row.get('date', '')))
    place = format_place(clean_cell(row.get('place', '')))
    hour  = clean_cell(row.get('hour', ''))
    summary = clean_cell(row.get('summary', ''))
    remarks = clean_cell(row.get('remarks', ''))

    # Header line: date and place
    header_parts = []
    if date:
        header_parts.append(date)
    if place:
        header_parts.append(place)
    if header_parts:
        parts.append(' — '.join(header_parts))

    # Hour on its own line if present
    if hour:
        parts.append(f'Hour: {hour}')

    # Blank line before summary
    if summary:
        parts.append('')
        parts.append(summary)

    # Remarks if present
    if remarks:
        parts.append('')
        parts.append(f'Remarks: {remarks}')

    return '\n'.join(parts)


# ---------------------------------------------------------------------------
# MD page parser
# ---------------------------------------------------------------------------

def parse_md_pages(md_text: str) -> dict[int, str]:
    """
    Parse a .md file into a dict of {page_number: page_text}.
    """
    pages = {}
    page_re = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
    matches = list(page_re.finditer(md_text))

    for i, m in enumerate(matches):
        page_num = int(m.group(1))
        start = m.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(md_text)
        page_text = md_text[start:end].strip()
        pages[page_num] = page_text

    return pages


# ---------------------------------------------------------------------------
# Page flattener
# ---------------------------------------------------------------------------

def flatten_page(page_num: int, page_text: str) -> str:
    """
    Flatten one page of OCR output to clean prose for the viewer.

    - If the page contains HTML tables: parse and flatten each table row.
    - If the page is plain text: pass through as-is.
    """

    # Extract any weather line that sits outside the table
    weather_match = re.search(
        r'(?:^|\n)(Weather[:\s].+?)(?:\n|$)',
        page_text,
        re.IGNORECASE
    )
    weather_line = weather_match.group(1).strip() if weather_match else None

    # Find all tables on this page
    table_re = re.compile(r'<table[\s\S]*?</table>', re.IGNORECASE)
    tables = table_re.findall(page_text)

    if not tables:
        # No table — plain text page, pass through unchanged
        return page_text

    # Diary page — parse each table and flatten rows
    output_parts = []

    for table_html in tables:
        rows = parse_table(table_html)
        if not rows:
            continue
        for row in rows:
            prose = row_to_prose(row)
            if prose.strip():
                output_parts.append(prose)
                output_parts.append('')  # blank line between entries

    # Append weather if present
    if weather_line:
        output_parts.append(weather_line)

    return '\n'.join(output_parts).strip()


# ---------------------------------------------------------------------------
# Main converter
# ---------------------------------------------------------------------------

def convert(input_path: Path, output_dir: Path):
    """
    Read the olmOCR .md file and write a clean viewer .md file.
    """
    print(f"\n{'='*60}")
    print(f"Converting: {input_path.name}")
    print(f"{'='*60}")

    md_text = input_path.read_text(encoding='utf-8')
    pages = parse_md_pages(md_text)

    if not pages:
        print("  ERROR: No pages found. Check file format.")
        sys.exit(1)

    print(f"  Found {len(pages)} pages.")

    # Build output
    stem = input_path.stem
    output_path = output_dir / f"{stem}_viewer.md"

    lines = [
        f"# {stem}",
        f"Generated by json_to_viewer_md.py from {input_path.name}",
        "",
    ]

    diary_count = 0
    plain_count = 0

    for page_num in sorted(pages.keys()):
        page_text = pages[page_num]

        lines.append(f"## Page {page_num}")
        lines.append("")

        if not page_text.strip() or page_text.strip() == '[BLANK]':
            lines.append('[BLANK]')
            plain_count += 1
        else:
            flattened = flatten_page(page_num, page_text)
            lines.append(flattened)
            if re.search(r'<table', page_text, re.IGNORECASE):
                diary_count += 1
            else:
                plain_count += 1

        lines.append("")

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path.write_text('\n'.join(lines), encoding='utf-8')

    print(f"  Diary pages flattened : {diary_count}")
    print(f"  Plain text pages      : {plain_count}")
    print(f"  Saved: {output_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Flatten olmOCR HTML table output to clean prose for the OCR viewer."
    )
    parser.add_argument(
        "--input",
        required=True,
        help="Path to the _olmocr.md file from ocr_wardiaries.py.",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Folder to write the _viewer.md file.",
    )
    args = parser.parse_args()

    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: input file not found: {input_path}")
        sys.exit(1)

    output_dir = Path(args.output_dir)
    convert(input_path, output_dir)


if __name__ == "__main__":
    main()