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

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


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 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 word_time_bounds(words: list[dict[str, Any]], start: int, end: int, fallback_start: Any, fallback_duration: Any) -> tuple[Any, Any]:
    span = words[start:end]
    starts = [entry.get("start") for entry in span if isinstance(entry.get("start"), (int, float))]
    ends = [
        float(entry.get("start")) + float(entry.get("duration"))
        for entry in span
        if isinstance(entry.get("start"), (int, float)) and isinstance(entry.get("duration"), (int, float))
    ]
    if starts and ends:
        first = min(float(value) for value in starts)
        last = max(ends)
        return first, max(0.001, last - first)
    return fallback_start, fallback_duration


def split_sentence_word_ranges(words: list[dict[str, Any]]) -> list[tuple[int, int]]:
    ranges: list[tuple[int, int]] = []
    start = 0
    for index, word in enumerate(words):
        text = str(word.get("text", "")).strip()
        is_sentence_end = bool(word.get("eos")) or bool(re.search(r'[.!?]["\')\]}]*$', text))
        if is_sentence_end and index + 1 > start:
            ranges.append((start, index + 1))
            start = index + 1
    if start < len(words):
        ranges.append((start, len(words)))
    return ranges


def normalize_timing_mode(value: str) -> str:
    return "sentence" if str(value).strip().casefold().startswith("sentence") else "segment"


def extract_units(transcript: dict[str, Any], timing_mode: str = "segment") -> 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 = words_for_segment(segment)
        if timing_mode == "sentence":
            ranges = split_sentence_word_ranges(words)
        else:
            ranges = [(0, len(words))]
        for range_index, (word_start, word_end) in enumerate(ranges):
            span_words = words[word_start:word_end]
            text = word_text(span_words)
            if not text:
                continue
            start_value, duration_value = word_time_bounds(
                words,
                word_start,
                word_end,
                segment.get("start"),
                segment.get("duration"),
            )
            units.append(
                {
                    "id": f"seg_{index:04d}" if timing_mode == "segment" else f"seg_{index:04d}_sent_{range_index:03d}",
                    "segmentIndex": index,
                    "wordStart": word_start,
                    "wordEnd": word_end,
                    "start": start_value,
                    "duration": duration_value,
                    "speaker": segment.get("speaker"),
                    "text": text,
                }
            )
    if not units:
        raise ValueError("Transcript JSON does not contain translatable 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))
    timing_mode = normalize_timing_mode(args.timing_mode)
    units = extract_units(transcript, timing_mode)
    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 = {
            "targetLanguage": args.target_language,
            "targetLanguageCode": args.target_language_code or "",
            "instructions": [
                "Use contextBefore and contextAfter only to understand meaning, pronouns, terminology, continuity, and tone.",
                "Translate only targetSegments.",
                "Return exactly one translation for each target segment id.",
                "Do not add, remove, rename, or reorder segment ids.",
                "Preserve names, product terms, and formatting when they should not be localized.",
                "Do not summarize, shorten, merge segments, split segments, add comments, or add markdown.",
                "Do not return timing data.",
            ],
            "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),
        "targetLanguage": args.target_language,
        "targetLanguageCode": args.target_language_code or "",
        "timingMode": timing_mode,
        "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_translation_response(value: Any) -> dict[str, str]:
    if not isinstance(value, dict):
        raise ValueError("Translation response root must be an object.")
    segments = value.get("segments")
    if not isinstance(segments, list):
        raise ValueError('Translation response must contain a "segments" array.')
    result: dict[str, str] = {}
    for index, entry in enumerate(segments):
        if not isinstance(entry, dict):
            raise ValueError(f"segments[{index}] must be an object.")
        seg_id = entry.get("id")
        translation = entry.get("translation")
        if not isinstance(seg_id, str) or not seg_id:
            raise ValueError(f"segments[{index}].id must be a non-empty string.")
        if not isinstance(translation, str) or not translation.strip():
            raise ValueError(f"segments[{index}].translation must be a non-empty string.")
        if seg_id in result:
            raise ValueError(f"Duplicate translated segment id: {seg_id}")
        result[seg_id] = translation.strip()
    return result


def translation_tokens(text: str) -> list[str]:
    return re.findall(r"\S+", text.strip())


def segment_bounds(segment: dict[str, Any]) -> tuple[float, float]:
    start = segment.get("start")
    duration = segment.get("duration")
    if isinstance(start, (int, float)) and isinstance(duration, (int, float)) and duration > 0:
        return float(start), float(duration)
    words = [entry for entry in segment.get("words", []) if isinstance(entry, dict)]
    word_starts = [entry.get("start") for entry in words if isinstance(entry.get("start"), (int, float))]
    word_ends = [
        float(entry.get("start")) + float(entry.get("duration"))
        for entry in words
        if isinstance(entry.get("start"), (int, float)) and isinstance(entry.get("duration"), (int, float))
    ]
    if word_starts and word_ends:
        first = min(float(value) for value in word_starts)
        last = max(word_ends)
        return first, max(0.001, last - first)
    return 0.0, 0.001


def make_translated_words(segment: dict[str, Any], translation: str, unit: dict[str, Any] | None = None) -> list[dict[str, Any]]:
    tokens = translation_tokens(translation)
    if not tokens:
        tokens = [translation.strip()]
    template_words = [entry for entry in segment.get("words", []) if isinstance(entry, dict)]
    if unit is not None and isinstance(unit.get("wordStart"), int) and isinstance(unit.get("wordEnd"), int):
        word_start = int(unit["wordStart"])
        word_end = int(unit["wordEnd"])
        start, duration = word_time_bounds(template_words, word_start, word_end, segment.get("start"), segment.get("duration"))
        unit_templates = template_words[word_start:word_end] or template_words
    else:
        start, duration = segment_bounds(segment)
        unit_templates = template_words
    per_word = duration / max(1, len(tokens))
    translated: list[dict[str, Any]] = []
    for index, token in enumerate(tokens):
        template = deepcopy(unit_templates[min(index, len(unit_templates) - 1)]) if unit_templates else {}
        template["text"] = token
        template["type"] = template.get("type", "word")
        template["start"] = start + per_word * index
        template["duration"] = per_word
        template["confidence"] = template.get("confidence", 1)
        template["eos"] = index == len(tokens) - 1
        if "tags" not in template:
            template["tags"] = []
        translated.append(template)
    return translated


def command_merge(args: argparse.Namespace) -> int:
    state_path = ab_path(args.state)
    state = load_json(state_path)
    transcript = load_json(Path(state["inputPath"]))
    if not isinstance(transcript, dict):
        raise ValueError("Transcript JSON root must be an object.")
    segments = transcript.get("segments")
    if not isinstance(segments, list):
        raise ValueError("Transcript JSON does not contain a segments array.")
    response_dir = Path(state["responseDir"])
    translations: 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 translation response file for {batch_id}: {response_path}")
        actual = normalize_translation_response(load_json(response_path))
        actual_ids = set(actual.keys())
        missing = sorted(expected - actual_ids)
        extra = sorted(actual_ids - expected)
        if missing or extra:
            raise ValueError(
                f"{batch_id} returned the wrong ids. Missing={missing or 'none'} Extra={extra or 'none'}"
            )
        translations.update(actual)

    units_by_id = {unit["id"]: unit for unit in state["units"]}
    units_by_segment: dict[int, list[dict[str, Any]]] = {}
    for unit in state["units"]:
        if unit["id"] in translations:
            units_by_segment.setdefault(int(unit["segmentIndex"]), []).append(unit)

    for segment_index, translated_units in units_by_segment.items():
        segment = segments[segment_index]
        if not isinstance(segment, dict):
            continue
        translated_units.sort(key=lambda entry: (int(entry.get("wordStart", 0)), str(entry.get("id", ""))))
        if state.get("timingMode") == "sentence":
            original_words = words_for_segment(segment)
            merged_words: list[dict[str, Any]] = []
            cursor = 0
            for unit in translated_units:
                word_start = int(unit.get("wordStart", cursor))
                word_end = int(unit.get("wordEnd", word_start))
                if cursor < word_start:
                    merged_words.extend(deepcopy(original_words[cursor:word_start]))
                merged_words.extend(make_translated_words(segment, translations[unit["id"]], unit))
                cursor = max(cursor, word_end)
            if cursor < len(original_words):
                merged_words.extend(deepcopy(original_words[cursor:]))
            segment["words"] = merged_words
        else:
            unit = translated_units[0]
            segment["words"] = make_translated_words(segment, translations[unit["id"]], unit)
        target_code = str(state.get("targetLanguageCode") or "").strip()
        if target_code:
            segment["language"] = target_code
    target_code = str(state.get("targetLanguageCode") or "").strip()
    if target_code:
        transcript["language"] = target_code

    output_path = ab_path(args.output)
    write_json(output_path, transcript)
    print(
        f"MERGE|translated={len(translations)}|batches={len(state['batches'])}|output={output_path}"
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Prepare and merge Automation Agent transcript translation batches.")
    sub = parser.add_subparsers(dest="command", required=True)

    plan = sub.add_parser("plan")
    plan.add_argument("--input", required=True)
    plan.add_argument("--work-dir", required=True)
    plan.add_argument("--target-language", required=True)
    plan.add_argument("--target-language-code", default="")
    plan.add_argument("--batch-size", type=int, default=8)
    plan.add_argument("--context-size", type=int, default=2)
    plan.add_argument("--timing-mode", default="segment")
    plan.set_defaults(func=command_plan)

    merge = sub.add_parser("merge")
    merge.add_argument("--state", required=True)
    merge.add_argument("--output", required=True)
    merge.set_defaults(func=command_merge)

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


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