src/localize.py

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

from src.config import LOCALIZED_DIR, PROMPTS_DIR
from src.generator import Generator
from src.region_registry import get_region_description


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Localize DAS encodings and context for a target language/community."
    )
    parser.add_argument(
        "--input_path",
        type=str,
        required=True,
        help="Path to encoded JSON file.",
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default=None,
        help="Path to save localized output JSON.",
    )
    parser.add_argument(
        "--language",
        type=str,
        required=True,
        help="Target language label, e.g. 'Swahili'.",
    )
    parser.add_argument(
        "--region",
        type=str,
        default="",
        help="Optional target region/community, e.g. 'Kenya - Nairobi'.",
    )
    parser.add_argument(
        "--context_prompt",
        type=str,
        default=str(PROMPTS_DIR / "das_localize_context.md"),
        help="Path to DAS context localization prompt.",
    )
    parser.add_argument(
        "--localize_prompt",
        type=str,
        default=str(PROMPTS_DIR / "das_localize.md"),
        help="Path to DAS localization prompt.",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Model alias from model_registry.py",
    )
    parser.add_argument(
        "--max_instances",
        type=int,
        default=None,
        help="Optional cap on number of dialogues to process.",
    )
    parser.add_argument(
        "--start_idx",
        type=int,
        default=0,
        help="Optional start index for slicing input data.",
    )
    parser.add_argument(
        "--end_idx",
        type=int,
        default=None,
        help="Optional end index for slicing input data.",
    )
    parser.add_argument(
        "--dont_use_cached",
        action="store_true",
        help="Disable cached prompt responses.",
    )
    return parser.parse_args()


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


def save_json(path: str, data: Any) -> None:
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(data, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )


def normalize_speaker_id(speaker_id: Any) -> str:
    speaker = str(speaker_id)
    if speaker == "1":
        return "A"
    if speaker == "2":
        return "B"
    return speaker


def stringify_functions(functions: Any) -> str:
    if isinstance(functions, list):
        return "; ".join(str(f) for f in functions)
    return str(functions)


def preprocess_conversation(das_encoding: List[Dict[str, Any]]) -> str:
    formatted_turns: List[str] = []

    for idx, turn in enumerate(das_encoding, start=1):
        speaker = normalize_speaker_id(turn.get("speaker_id", "A"))
        functions = stringify_functions(turn.get("functions", []))
        formatted_turns.append(f"{idx}: {speaker}.{functions}")

    return "\n".join(formatted_turns)


def preprocess_localize_input(
    data: List[Dict[str, Any]],
    language: str,
    region: str,
) -> List[Dict[str, Any]]:
    context_desc = get_region_description(region, "localize_context", language) or ""
    das_desc = get_region_description(region, "localize_das", language) or ""

    processed: List[Dict[str, Any]] = []

    for item in data:
        if "das_encoding" not in item:
            raise ValueError("Each input item must contain 'das_encoding'.")
        if "context" not in item:
            raise ValueError("Each input item must contain 'context'.")

        new_item = dict(item)
        new_item["language"] = language
        new_item["region"] = region
        new_item["region_description"] = context_desc
        new_item["turns"] = preprocess_conversation(item["das_encoding"])
        processed.append(new_item)

    return processed, das_desc


def merge_localized_context(
    base_data: List[Dict[str, Any]],
    responses: List[str],
) -> List[Dict[str, Any]]:
    merged: List[Dict[str, Any]] = []

    for item, response_text in zip(base_data, responses):
        if response_text is None:
            print(f"[Localize] Skipping item with failed context generation")
            continue

        response_json = Generator.parse_json_response(response_text)

        if "localized_context" not in response_json:
            raise ValueError(
                f"Missing 'localized_context' in model response:\n{response_text}"
            )

        new_item = dict(item)
        new_item["localized_context"] = response_json["localized_context"]
        merged.append(new_item)

    return merged


def merge_localized_das(
    base_data: List[Dict[str, Any]],
    responses: List[str],
) -> List[Dict[str, Any]]:
    merged: List[Dict[str, Any]] = []

    for item, response_text in zip(base_data, responses):
        if response_text is None:
            print(f"[Localize] Skipping item with failed DAS generation")
            continue

        response_json = Generator.parse_json_response(response_text)

        if "localized_das" not in response_json:
            raise ValueError(
                f"Missing 'localized_das' in model response:\n{response_text}"
            )

        new_item = dict(item)
        new_item["localized_das"] = response_json["localized_das"]
        merged.append(new_item)

    return merged


def default_output_path(input_path: str, language: str, region: str) -> str:
    stem = Path(input_path).stem
    suffix_parts = [language.strip().lower().replace(" ", "_")]

    if region.strip():
        suffix_parts.append(region.strip().lower().replace(" ", "_").replace("/", "_"))

    suffix = "_".join(suffix_parts)
    return str(LOCALIZED_DIR / f"{stem}_{suffix}_localized.json")


def main() -> None:
    args = parse_args()

    raw_data = load_json(args.input_path)
    if not isinstance(raw_data, list):
        raise ValueError("Input JSON must be a list of dialogue objects.")

    sliced_data = raw_data[args.start_idx:args.end_idx]
    if args.max_instances is not None:
        sliced_data = sliced_data[: args.max_instances]

    output_path = args.output_path or default_output_path(
        args.input_path,
        args.language,
        args.region,
    )

    generator = Generator(
        model_alias=args.model,
        use_cache=not args.dont_use_cached,
    )

    processed_data, das_description = preprocess_localize_input(
        data=sliced_data,
        language=args.language,
        region=args.region,
    )

    print(f"[Localize] Building localized contexts for {len(processed_data)} items...")
    context_prompts, context_response_format = generator.build_prompts(
        args.context_prompt,
        processed_data,
    )
    context_responses = generator.prompt(
        prompts=context_prompts,
        response_format=context_response_format,
        dont_use_cached=args.dont_use_cached,
        skip_failures=True,
    )
    data_with_localized_context = merge_localized_context(
        processed_data,
        context_responses,
    )

    # Swap region_description to the DAS-stage version for the localize prompt
    for item in data_with_localized_context:
        item["region_description"] = das_description

    print(
        f"[Localize] Building localized DAS for "
        f"{len(data_with_localized_context)} items..."
    )
    localize_prompts, localize_response_format = generator.build_prompts(
        args.localize_prompt,
        data_with_localized_context,
    )
    localize_responses = generator.prompt(
        prompts=localize_prompts,
        response_format=localize_response_format,
        dont_use_cached=args.dont_use_cached,
        skip_failures=True,
    )
    final_data = merge_localized_das(
        data_with_localized_context,
        localize_responses,
    )

    save_json(output_path, final_data)
    print(f"Saved localized data to: {output_path}")
    generator.print_usage_summary(stage="Localize")


if __name__ == "__main__":
    main()