import os
import sys
import argparse
import anthropic
from pptx import Presentation

from translate_lecture.translation import (
    TranslationConfig, load_config, build_translation_system_prompt,
    extract_records, build_translation_cache, apply_translations, audit_files,
)


def process_file(src_path: str, dst_path: str, client: anthropic.Anthropic,
                 model: str, batch_size: int, tcfg: TranslationConfig):
    print(f"\n{'='*60}")
    print(f"Processing: {os.path.basename(src_path)}")
    print(f"{'='*60}")

    prs = Presentation(src_path)
    print(f"  Slides: {len(prs.slides)}")

    file_label = os.path.basename(src_path)
    records = extract_records(prs, file_label)
    print(f"  Paragraphs extracted: {len(records)}")

    cache = build_translation_cache(records, client, model, batch_size, tcfg)

    print(f"  Writing translations...")
    apply_translations(prs, records, cache, tcfg)

    prs.save(dst_path)
    print(f"  Saved: {os.path.basename(dst_path)}")

    # Verify slide count
    prs2 = Presentation(dst_path)
    assert len(prs2.slides) == len(prs.slides), "Slide count mismatch!"
    print(f"  Slide count verified: {len(prs2.slides)}")


def main():
    parser = argparse.ArgumentParser(
        description="Translate lecture PPTX files using a JSON config file."
    )
    parser.add_argument("config", help="Path to the JSON config file")
    parser.add_argument("--audit", action="store_true",
                        help="Print extracted texts without calling the API")
    parser.add_argument(
        "--model",
        default="claude-haiku-4-5",
        help="Claude model to use (default: claude-haiku-4-5)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=40,
        help="Number of texts to translate in each API call (default: 40)",
    )
    args = parser.parse_args()

    cfg = load_config(os.path.abspath(args.config))
    config_dir = os.path.dirname(os.path.abspath(args.config))
    system_prompt = build_translation_system_prompt(cfg)
    model = args.model
    batch_size = args.batch_size

    tcfg = TranslationConfig(
        source_language=cfg["lecture"]["source_language"],
        target_language=cfg["lecture"]["target_language"],
        skip_exact=set(cfg["translation"].get("skip_exact", [])),
        spacing_fix_prefixes=cfg["translation"].get("spacing_fix_prefixes", []),
        system_prompt=system_prompt,
    )

    if args.audit:
        audit_files(cfg["files"], config_dir)
        return

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("ERROR: ANTHROPIC_API_KEY environment variable not set.")
        sys.exit(1)

    client = anthropic.Anthropic(api_key=api_key)
    for entry in cfg["files"]:
        src_path = os.path.join(config_dir, entry["source"])
        dst_path = os.path.join(config_dir, entry["target"])

        if not os.path.exists(src_path):
            print(f"WARNING: source file not found: {src_path}")
            continue

        process_file(src_path, dst_path, client, model, batch_size, tcfg)

    print("\nDone.")
