#!/usr/bin/env python
"""Ingest and clean the MTSamples CSV into structured JSONL."""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
from collections import Counter
from datetime import datetime, timezone
from typing import Dict, Iterable, List, Tuple

SECTION_RE = re.compile(r"([A-Z][A-Z0-9 /&\-]{1,}):")


def normalize_whitespace(text: str) -> str:
    text = text.replace("\u0000", "").replace("\ufffd", "")
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def clean_section_content(text: str) -> str:
    text = normalize_whitespace(text)
    text = re.sub(r"([.!?])\s*,", r"\1", text)
    text = re.sub(r"[ ,;:]+$", "", text)
    return text


def split_sections(text: str) -> List[Tuple[str, str]]:
    matches = list(SECTION_RE.finditer(text))
    if not matches:
        cleaned = clean_section_content(text)
        return [("FULL_TEXT", cleaned)] if cleaned else []

    sections: List[Tuple[str, str]] = []
    for idx, match in enumerate(matches):
        start = match.end()
        end = matches[idx + 1].start() if idx + 1 < len(matches) else len(text)
        heading = match.group(1).strip()
        content = text[start:end]
        content = content.lstrip(" ,;:-\n\t")
        content = clean_section_content(content)
        if content:
            sections.append((heading, content))
    return sections


def build_doc_id(row: Dict[str, str], fallback_index: int) -> str:
    raw_id = (row.get("") or "").strip()
    if not raw_id:
        raw_id = str(fallback_index)
    return f"mtsamples_{raw_id}"


def iter_mtsamples(path: str) -> Iterable[Dict[str, str]]:
    with open(path, newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            yield row


def ingest(input_path: str, output_dir: str) -> Dict[str, object]:
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "cleaned_mtsamples.jsonl")

    section_counts = Counter()
    doc_count = 0
    section_count = 0
    empty_section_count = 0

    with open(output_path, "w", encoding="utf-8") as out_handle:
        for idx, row in enumerate(iter_mtsamples(input_path)):
            doc_id = build_doc_id(row, idx)
            transcription = (row.get("transcription") or "").strip()
            sections = split_sections(transcription)

            if not sections and transcription:
                empty_section_count += 1

            for section_idx, (section, content) in enumerate(sections):
                record = {
                    "doc_id": doc_id,
                    "page": None,
                    "section": section,
                    "section_index": section_idx,
                    "content": content,
                    "source": {
                        "filename": os.path.basename(input_path),
                        "row_id": row.get("") or str(idx),
                        "sample_name": (row.get("sample_name") or "").strip(),
                        "medical_specialty": (row.get("medical_specialty") or "").strip(),
                        "description": (row.get("description") or "").strip(),
                        "keywords": (row.get("keywords") or "").strip(),
                    },
                }
                out_handle.write(json.dumps(record, ensure_ascii=True) + "\n")
                section_counts[section] += 1
                section_count += 1

            doc_count += 1

    stats = {
        "input_file": os.path.abspath(input_path),
        "output_file": os.path.abspath(output_path),
        "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "document_count": doc_count,
        "section_count": section_count,
        "empty_section_count": empty_section_count,
        "top_sections": section_counts.most_common(10),
        "notes": [
            "Page numbers are not present in the MTSamples CSV; page is emitted as null.",
        ],
    }

    stats_path = os.path.join(output_dir, "ingestion_stats.json")
    with open(stats_path, "w", encoding="utf-8") as stats_handle:
        json.dump(stats, stats_handle, indent=2, ensure_ascii=True)

    return stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Ingest and clean MTSamples CSV.")
    parser.add_argument(
        "--input",
        default=os.path.join("00_Data", "mtsamples.csv"),
        help="Path to mtsamples.csv",
    )
    parser.add_argument(
        "--output-dir",
        default=os.path.join("01_data_ingestion", "output"),
        help="Output directory for JSONL and stats.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ingest(args.input, args.output_dir)


if __name__ == "__main__":
    main()