#!/usr/bin/env python3
from __future__ import annotations

import argparse
import difflib
import json
import re
import sys
from copy import deepcopy
from pathlib import Path
from typing import Any

MAX_CHANGED_WORDS = 18


def ab_path(value: str) -> Path:
    text = str(value)
    if text.startswith("::HOME::"):
        text = str(Path.home()) + text[len("::HOME::") :]
    elif text.startswith("file://"):
        text = text[7:]
    elif text.startswith("file:"):
        text = text[5:]
    return Path(text).expanduser()


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")


def word_text(words: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    for word in words:
        text = str(word.get("text", "")).strip()
        if not text:
            continue
        if not parts:
            parts.append(text)
        elif re.match(r"^[,.;:!?%)\]}]+$", text):
            parts[-1] += text
        elif re.match(r"^[(\[{]+$", text):
            parts.append(text)
        else:
            parts.append(" " + text)
    return "".join(parts).strip()


def normalize_text(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()


def diff_word_count(old_text: str, new_text: str) -> int:
    old_tokens = re.findall(r"\S+", normalize_text(old_text))
    new_tokens = re.findall(r"\S+", normalize_text(new_text))
    prefix = 0
    while (
        prefix < len(old_tokens)
        and prefix < len(new_tokens)
        and old_tokens[prefix].casefold() == new_tokens[prefix].casefold()
    ):
        prefix += 1
    old_suffix = len(old_tokens)
    new_suffix = len(new_tokens)
    while (
        old_suffix > prefix
        and new_suffix > prefix
        and old_tokens[old_suffix - 1].casefold() == new_tokens[new_suffix - 1].casefold()
    ):
        old_suffix -= 1
        new_suffix -= 1
    return max(old_suffix - prefix, new_suffix - prefix)


def normalized_find_count(haystack: str, needle: str) -> int:
    normalized_haystack = normalize_text(haystack).casefold()
    normalized_needle = normalize_text(needle).casefold()
    if not normalized_needle:
        return 0
    count = 0
    start = 0
    while True:
        index = normalized_haystack.find(normalized_needle, start)
        if index < 0:
            return count
        count += 1
        start = index + max(1, len(normalized_needle))


def normalize_span_match_text(value: str) -> str:
    return re.sub(r"^[\s\"'“”‘’()[\]{}.,;:!?]+|[\s\"'“”‘’()[\]{}.,;:!?]+$", "", normalize_text(value)).casefold()


def extract_units(transcript: dict[str, Any]) -> list[dict[str, Any]]:
    segments = transcript.get("segments")
    if not isinstance(segments, list):
        raise ValueError("Transcript JSON does not contain a segments array.")
    units: list[dict[str, Any]] = []
    for index, segment in enumerate(segments):
        if not isinstance(segment, dict):
            continue
        words_raw = segment.get("words")
        words = [entry for entry in words_raw if isinstance(entry, dict)] if isinstance(words_raw, list) else []
        text = word_text(words)
        if not text:
            continue
        units.append(
            {
                "id": f"seg_{index:04d}",
                "segmentIndex": index,
                "start": segment.get("start"),
                "duration": segment.get("duration"),
                "speaker": segment.get("speaker"),
                "text": text,
            }
        )
    if not units:
        raise ValueError("Transcript JSON does not contain readable word text.")
    return units


def chunk_units(units: list[dict[str, Any]], batch_size: int, context_size: int) -> list[dict[str, Any]]:
    batches: list[dict[str, Any]] = []
    for start in range(0, len(units), batch_size):
        end = min(len(units), start + batch_size)
        before_start = max(0, start - context_size)
        after_end = min(len(units), end + context_size)
        batch_id = f"batch_{len(batches):04d}"
        batches.append(
            {
                "id": batch_id,
                "targetIds": [unit["id"] for unit in units[start:end]],
                "contextBefore": units[before_start:start],
                "targetSegments": units[start:end],
                "contextAfter": units[end:after_end],
            }
        )
    return batches


def command_plan(args: argparse.Namespace) -> int:
    input_path = ab_path(args.input)
    work_dir = ab_path(args.work_dir)
    request_dir = work_dir / "requests"
    response_dir = work_dir / "responses"
    state_path = work_dir / "state.json"
    transcript = load_json(input_path)
    if not isinstance(transcript, dict):
        raise ValueError("Transcript JSON root must be an object.")
    batch_size = max(1, int(args.batch_size))
    context_size = max(0, int(args.context_size))
    units = extract_units(transcript)
    batches = chunk_units(units, batch_size, context_size)
    request_dir.mkdir(parents=True, exist_ok=True)
    response_dir.mkdir(parents=True, exist_ok=True)

    rows: list[str] = []
    for batch in batches:
        request_path = request_dir / f"{batch['id']}.json"
        response_path = response_dir / f"{batch['id']}.response.json"
        request = {
            "projectContext": args.project_context,
            "knownNamesAndTerminology": args.known_names,
            "proofreadingInstructions": args.focus,
            "reportDetail": args.detail,
            "instructions": [
                "Use contextBefore and contextAfter only to understand meaning, names, terminology, and repeated phrasing.",
                "Suggest corrections only for targetSegments.",
                "Return correction patches, not a rewritten transcript.",
                "Each correction must include segmentId, oldText, newText, reason, and severity.",
                "oldText must be copied exactly from the target segment whenever possible.",
                "Prefer high-confidence transcript recognition fixes such as misspelled names, brands, technical terms, duplicate words, obvious casing, and punctuation that changes readability or meaning.",
                "Do not rewrite style, polish wording, summarize, or make general grammar edits.",
                "For Concise reports, return only high-confidence corrections. For Detailed reports, medium-confidence transcript-recognition corrections are allowed, but uncertain guesses should still be omitted.",
                "Prefer small local corrections. Do not rewrite whole sentences unless the original is unusable.",
                "Do not include timing data, word indices, markdown, or comments.",
                "If there are no likely issues in the target segments, return an empty corrections array.",
            ],
            "contextBefore": [{"id": u["id"], "text": u["text"]} for u in batch["contextBefore"]],
            "targetSegments": [{"id": u["id"], "text": u["text"]} for u in batch["targetSegments"]],
            "contextAfter": [{"id": u["id"], "text": u["text"]} for u in batch["contextAfter"]],
        }
        write_json(request_path, request)
        rows.append("\t".join([batch["id"], str(request_path), str(response_path), ",".join(batch["targetIds"])]))

    state = {
        "inputPath": str(input_path),
        "workDir": str(work_dir),
        "responseDir": str(response_dir),
        "projectContext": args.project_context,
        "knownNamesAndTerminology": args.known_names,
        "proofreadingInstructions": args.focus,
        "reportDetail": args.detail,
        "unitCount": len(units),
        "batchCount": len(batches),
        "units": units,
        "batches": [{"id": batch["id"], "targetIds": batch["targetIds"]} for batch in batches],
    }
    write_json(state_path, state)
    print("\n".join(rows))
    print(f"PLAN|units={len(units)}|batches={len(batches)}|state={state_path}", file=sys.stderr)
    return 0


def normalize_corrections_response(value: Any) -> list[dict[str, str]]:
    if not isinstance(value, dict):
        raise ValueError("Proofread response root must be an object.")
    corrections = value.get("corrections")
    if not isinstance(corrections, list):
        raise ValueError('Proofread response must contain a "corrections" array.')
    result: list[dict[str, str]] = []
    for index, entry in enumerate(corrections):
        if not isinstance(entry, dict):
            raise ValueError(f"corrections[{index}] must be an object.")
        normalized: dict[str, str] = {}
        for field in ("segmentId", "oldText", "newText", "reason", "severity"):
            value = entry.get(field)
            if not isinstance(value, str) or not value.strip():
                raise ValueError(f"corrections[{index}].{field} must be a non-empty string.")
            normalized[field] = value.strip()
        result.append(normalized)
    return result


def validate_correction(correction: dict[str, str], target_ids: set[str], units_by_id: dict[str, dict[str, Any]]) -> tuple[bool, str]:
    segment_id = correction["segmentId"]
    if segment_id not in target_ids:
        return False, f"segmentId {segment_id!r} was not part of this batch target set."
    unit = units_by_id.get(segment_id)
    if not unit:
        return False, f"segmentId {segment_id!r} does not exist in the transcript state."
    if normalize_text(correction["oldText"]) == normalize_text(correction["newText"]):
        return False, "oldText and newText are identical."
    occurrences = normalized_find_count(str(unit.get("text", "")), correction["oldText"])
    if occurrences == 0:
        return False, "oldText was not found in the source segment."
    if occurrences > 1:
        return False, "oldText appears more than once in the source segment; manual review needed."
    changed_words = diff_word_count(correction["oldText"], correction["newText"])
    if changed_words > MAX_CHANGED_WORDS:
        return False, f"changed span is {changed_words} words; keep this as manual review instead of a local correction."
    return True, f"validated; changed span is {changed_words} words"


def markdown_escape(value: Any) -> str:
    return str(value).replace("|", "\\|").replace("\n", " ").strip()


def collect_validation_records(state: dict[str, Any]) -> tuple[list[dict[str, str]], list[dict[str, str]]]:
    units_by_id = {unit["id"]: unit for unit in state["units"]}
    response_dir = Path(state["responseDir"])
    valid: list[dict[str, str]] = []
    rejected: list[dict[str, str]] = []
    for batch in state["batches"]:
        batch_id = batch["id"]
        expected = set(batch["targetIds"])
        response_path = response_dir / f"{batch_id}.response.json"
        if not response_path.exists():
            raise ValueError(f"Missing proofread response file for {batch_id}: {response_path}")
        corrections = normalize_corrections_response(load_json(response_path))
        for correction_index, correction in enumerate(corrections):
            ok, note = validate_correction(correction, expected, units_by_id)
            record = {
                **correction,
                "correctionId": f"{batch_id}:{correction_index:04d}",
                "batchId": batch_id,
                "validationNote": note,
            }
            if ok:
                valid.append(record)
            else:
                rejected.append(record)
    return valid, rejected


def build_report(state: dict[str, Any], item_name: str) -> tuple[str, int, int]:
    valid, rejected = collect_validation_records(state)

    lines: list[str] = []
    lines.append(f"## {item_name}")
    lines.append("")
    lines.append(f"- Transcript segments checked: {state['unitCount']}")
    lines.append(f"- Workflow AI batches: {state['batchCount']}")
    lines.append(f"- Validated suggestions: {len(valid)}")
    lines.append(f"- Suggestions needing manual review: {len(rejected)}")
    lines.append("")
    if valid:
        lines.append("### Suggested corrections")
        lines.append("")
        lines.append("| Segment | Severity | Old text | Suggested text | Reason | Validation |")
        lines.append("| --- | --- | --- | --- | --- | --- |")
        for entry in valid:
            lines.append(
                "| "
                + " | ".join(
                    [
                        markdown_escape(entry["segmentId"]),
                        markdown_escape(entry["severity"]),
                        markdown_escape(entry["oldText"]),
                        markdown_escape(entry["newText"]),
                        markdown_escape(entry["reason"]),
                        markdown_escape(entry["validationNote"]),
                    ]
                )
                + " |"
            )
        lines.append("")
    else:
        lines.append("No validated correction suggestions were returned.")
        lines.append("")
    if rejected:
        lines.append("### Needs manual review")
        lines.append("")
        lines.append("| Segment | Old text | Suggested text | Why it was not auto-validated |")
        lines.append("| --- | --- | --- | --- |")
        for entry in rejected:
            lines.append(
                "| "
                + " | ".join(
                    [
                        markdown_escape(entry["segmentId"]),
                        markdown_escape(entry["oldText"]),
                        markdown_escape(entry["newText"]),
                        markdown_escape(entry["validationNote"]),
                    ]
                )
                + " |"
            )
        lines.append("")
    return "\n".join(lines).strip() + "\n\n", len(valid), len(rejected)


def tsv_cell(value: Any) -> str:
    return str(value).replace("\t", " ").replace("\r", " ").replace("\n", " ").strip()


def format_seconds(value: Any) -> str:
    if not isinstance(value, (int, float)):
        return ""
    total = max(0, int(float(value)))
    hours = total // 3600
    minutes = (total % 3600) // 60
    seconds = total % 60
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}" if hours else f"{minutes:02d}:{seconds:02d}"


def command_list_validated(args: argparse.Namespace) -> int:
    state = load_json(ab_path(args.state))
    valid, _rejected = collect_validation_records(state)
    units_by_id = {unit["id"]: unit for unit in state["units"]}
    rows: list[str] = []
    for index, entry in enumerate(valid):
        unit = units_by_id.get(entry["segmentId"], {})
        rows.append(
            "\t".join(
                [
                    tsv_cell(entry["correctionId"]),
                    tsv_cell(entry["segmentId"]),
                    tsv_cell(entry["oldText"]),
                    tsv_cell(entry["newText"]),
                    tsv_cell(entry["reason"]),
                    tsv_cell(entry["severity"]),
                    tsv_cell(format_seconds(unit.get("start"))),
                    tsv_cell(unit.get("text", "")),
                    tsv_cell(index + 1),
                    tsv_cell(len(valid)),
                ]
            )
        )
    print("\n".join(rows))
    print(f"VALIDATED|count={len(valid)}", file=sys.stderr)
    return 0


def load_accepted_ids(path: Path) -> set[str]:
    if not path.exists():
        return set()
    return {line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()}


def words_for_segment(segment: dict[str, Any]) -> list[dict[str, Any]]:
    raw = segment.get("words")
    return [entry for entry in raw if isinstance(entry, dict)] if isinstance(raw, list) else []


def find_word_span(words: list[dict[str, Any]], old_text: str) -> tuple[int, int] | None:
    target = normalize_text(old_text).casefold()
    edge_punctuation_target = normalize_span_match_text(old_text)
    if not target:
        return None
    for start in range(0, len(words)):
        for end in range(start + 1, len(words) + 1):
            candidate = normalize_text(word_text(words[start:end])).casefold()
            if candidate == target or (
                edge_punctuation_target and normalize_span_match_text(candidate) == edge_punctuation_target
            ):
                return start, end
            if len(candidate) > len(target) + 20:
                break
    return None


def replacement_tokens(text: str) -> list[str]:
    tokens = re.findall(r"\S+", normalize_text(text))
    return tokens or [text.strip()]


def clone_word_with_text(word: dict[str, Any], text: str) -> dict[str, Any]:
    entry = deepcopy(word)
    entry["text"] = text
    entry["type"] = entry.get("type", "word")
    entry["confidence"] = entry.get("confidence", 1)
    if "tags" not in entry:
        entry["tags"] = []
    return entry


def word_bounds(words: list[dict[str, Any]], start: int, end: int) -> tuple[float, float]:
    span = words[start:end]
    first_start = span[0].get("start") if span else None
    last = span[-1] if span else {}
    last_start = last.get("start")
    last_duration = last.get("duration")
    if (
        isinstance(first_start, (int, float))
        and isinstance(last_start, (int, float))
        and isinstance(last_duration, (int, float))
    ):
        span_start = float(first_start)
        span_end = float(last_start) + max(0.001, float(last_duration))
        return span_start, max(0.001, span_end - span_start)
    return 0.0, max(0.001, float(max(1, end - start)))


def interpolate_words_in_span(original_span: list[dict[str, Any]], tokens: list[str]) -> list[dict[str, Any]]:
    span_start, span_duration = word_bounds(original_span, 0, len(original_span))
    per_word = span_duration / max(1, len(tokens))
    replaced = []
    for index, token in enumerate(tokens):
        template = deepcopy(original_span[min(index, len(original_span) - 1)]) if original_span else {}
        template["text"] = token
        template["type"] = template.get("type", "word")
        template["start"] = span_start + per_word * index
        template["duration"] = per_word
        template["confidence"] = template.get("confidence", 1)
        if "tags" not in template:
            template["tags"] = []
        replaced.append(template)
    return replaced


def word_time_start(word: dict[str, Any]) -> float | None:
    value = word.get("start")
    return float(value) if isinstance(value, (int, float)) else None


def word_time_end(word: dict[str, Any]) -> float | None:
    start = word.get("start")
    duration = word.get("duration")
    if isinstance(start, (int, float)) and isinstance(duration, (int, float)):
        return float(start) + max(0.001, float(duration))
    return None


def time_gap(left: dict[str, Any] | None, right: dict[str, Any] | None) -> tuple[float, float] | None:
    if left is None or right is None:
        return None
    left_end = word_time_end(left)
    right_start = word_time_start(right)
    if left_end is None or right_start is None or right_start <= left_end:
        return None
    return left_end, right_start - left_end


def append_insert_to_neighbor(result: list[dict[str, Any]], pending: list[str], next_word: dict[str, Any] | None) -> None:
    if not pending:
        return
    insert_text = " ".join(pending)
    if result:
        result[-1]["text"] = f"{result[-1].get('text', '')} {insert_text}".strip()
    elif next_word is not None:
        result.append(clone_word_with_text(next_word, f"{insert_text} {next_word.get('text', '')}".strip()))
    else:
        result.append({"text": insert_text, "type": "word", "start": 0.0, "duration": 0.001, "confidence": 1, "tags": []})


def insert_words_between(
    result: list[dict[str, Any]],
    tokens: list[str],
    left_word: dict[str, Any] | None,
    right_word: dict[str, Any] | None,
) -> None:
    if not tokens:
        return
    gap = time_gap(left_word, right_word)
    if gap is None:
        append_insert_to_neighbor(result, tokens, right_word)
        return
    gap_start, gap_duration = gap
    per_word = gap_duration / len(tokens)
    template = left_word or right_word or {}
    for index, token in enumerate(tokens):
        entry = clone_word_with_text(template, token)
        entry["start"] = gap_start + per_word * index
        entry["duration"] = per_word
        result.append(entry)


def aligned_replacement_words(original_span: list[dict[str, Any]], new_text: str) -> list[dict[str, Any]]:
    old_tokens = [str(word.get("text", "")).strip() for word in original_span]
    new_tokens = replacement_tokens(new_text)
    if not old_tokens:
        return interpolate_words_in_span(original_span, new_tokens)

    matcher = difflib.SequenceMatcher(
        a=[token.casefold() for token in old_tokens],
        b=[token.casefold() for token in new_tokens],
        autojunk=False,
    )
    result: list[dict[str, Any]] = []
    opcodes = matcher.get_opcodes()
    prefix_for_next: list[str] = []

    def consume_prefix(token: str) -> str:
        nonlocal prefix_for_next
        if not prefix_for_next:
            return token
        text = f"{' '.join(prefix_for_next)} {token}".strip()
        prefix_for_next = []
        return text

    for tag, i1, i2, j1, j2 in opcodes:
        if tag == "equal":
            for old_index, new_index in zip(range(i1, i2), range(j1, j2)):
                result.append(clone_word_with_text(original_span[old_index], consume_prefix(new_tokens[new_index])))
            continue
        if tag == "delete":
            continue
        if tag == "replace":
            old_slice = original_span[i1:i2]
            new_slice = new_tokens[j1:j2]
            if len(old_slice) == len(new_slice):
                for old_word, token in zip(old_slice, new_slice):
                    result.append(clone_word_with_text(old_word, consume_prefix(token)))
            else:
                if prefix_for_next and new_slice:
                    new_slice = [consume_prefix(new_slice[0]), *new_slice[1:]]
                result.extend(interpolate_words_in_span(old_slice, new_slice))
            continue
        if tag == "insert":
            left_word = result[-1] if result else (original_span[i1 - 1] if i1 > 0 else None)
            right_word = original_span[i1] if i1 < len(original_span) else None
            gap = time_gap(left_word, right_word)
            inserted_tokens = new_tokens[j1:j2]
            if gap is None and not result and right_word is not None:
                prefix_for_next.extend(inserted_tokens)
            elif gap is None and result:
                append_insert_to_neighbor(result, inserted_tokens, right_word)
            else:
                insert_words_between(result, inserted_tokens, left_word, right_word)
            continue
    return result


def make_replacement_words(words: list[dict[str, Any]], start: int, end: int, new_text: str) -> list[dict[str, Any]]:
    return aligned_replacement_words(words[start:end], new_text)


def normalize_segment_eos(segment: dict[str, Any]) -> None:
    words = words_for_segment(segment)
    for index, word in enumerate(words):
        word["eos"] = index == len(words) - 1


def apply_correction_to_transcript(transcript: dict[str, Any], state: dict[str, Any], correction: dict[str, str]) -> bool:
    units_by_id = {unit["id"]: unit for unit in state["units"]}
    unit = units_by_id.get(correction["segmentId"])
    if not unit:
        return False
    segments = transcript.get("segments")
    if not isinstance(segments, list):
        raise ValueError("Transcript JSON does not contain a segments array.")
    segment_index = int(unit["segmentIndex"])
    if segment_index < 0 or segment_index >= len(segments) or not isinstance(segments[segment_index], dict):
        return False
    segment = segments[segment_index]
    words = words_for_segment(segment)
    span = find_word_span(words, correction["oldText"])
    if span is None:
        return False
    start, end = span
    replacement = make_replacement_words(words, start, end, correction["newText"])
    segment["words"] = words[:start] + replacement + words[end:]
    normalize_segment_eos(segment)
    return True


def command_apply(args: argparse.Namespace) -> int:
    state = load_json(ab_path(args.state))
    transcript = load_json(Path(state["inputPath"]))
    if not isinstance(transcript, dict):
        raise ValueError("Transcript JSON root must be an object.")
    accepted_ids = load_accepted_ids(ab_path(args.accepted))
    valid, _rejected = collect_validation_records(state)
    applied = 0
    skipped = 0
    for correction in valid:
        if correction["correctionId"] not in accepted_ids:
            continue
        if apply_correction_to_transcript(transcript, state, correction):
            applied += 1
        else:
            skipped += 1
            print(
                f"APPLY_SKIP|correctionId={correction['correctionId']}|segmentId={correction['segmentId']}|reason=\"word span not found\"",
                file=sys.stderr,
            )
    output_path = ab_path(args.output)
    write_json(output_path, transcript)
    print(f"APPLY|accepted={len(accepted_ids)}|applied={applied}|skipped={skipped}|output={output_path}")
    return 0


def command_report(args: argparse.Namespace) -> int:
    state_path = ab_path(args.state)
    state = load_json(state_path)
    markdown, valid_count, rejected_count = build_report(state, args.item_name)
    print(markdown, end="")
    print(
        f"REPORT|valid={valid_count}|manualReview={rejected_count}|batches={state['batchCount']}",
        file=sys.stderr,
    )
    return 0


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description="Automation Agent transcript proofreading pipeline")
    subparsers = parser.add_subparsers(dest="command", required=True)

    plan = subparsers.add_parser("plan")
    plan.add_argument("--input", required=True)
    plan.add_argument("--work-dir", required=True)
    plan.add_argument("--focus", default="")
    plan.add_argument("--project-context", default="")
    plan.add_argument("--known-names", default="")
    plan.add_argument("--detail", required=True)
    plan.add_argument("--batch-size", required=True, type=int)
    plan.add_argument("--context-size", required=True, type=int)
    plan.set_defaults(func=command_plan)

    report = subparsers.add_parser("report")
    report.add_argument("--state", required=True)
    report.add_argument("--item-name", required=True)
    report.set_defaults(func=command_report)

    list_validated = subparsers.add_parser("list-validated")
    list_validated.add_argument("--state", required=True)
    list_validated.set_defaults(func=command_list_validated)

    apply = subparsers.add_parser("apply")
    apply.add_argument("--state", required=True)
    apply.add_argument("--accepted", required=True)
    apply.add_argument("--output", required=True)
    apply.set_defaults(func=command_apply)

    args = parser.parse_args(argv)
    return int(args.func(args))


if __name__ == "__main__":
    try:
        raise SystemExit(main(sys.argv[1:]))
    except Exception as exc:
        print(f"ERROR|{exc}", file=sys.stderr)
        raise SystemExit(1)
