# #region LLMResponseParser [C:3] [TYPE Module] [SEMANTICS translate, llm, parse, json, recovery]
# @BRIEF Parse LLM JSON response into per-row translations with support for markdown code
#         blocks, truncated JSON recovery, and multi-language format.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [TranslationJob]
# @RATIONALE Extracted from _llm_call.py (793 lines) to comply with INV_7 (< 400 lines/module).
# @REJECTED Inline parsing inside call_llm_for_batch — made the function 327 lines.

import json
import re

from ...core.logger import logger


# #region parse_llm_response [C:3] [TYPE Function] [SEMANTICS translate, llm, parse, json]
# @BRIEF Parse LLM JSON response into dict of row_id -> per-language translations.
# @PRE response_text is valid JSON with {"rows": [...]} structure.
# @POST Returns dict mapping row_id to dict with 'detected_source_language' and per-language codes.
# @RAISES ValueError if JSON cannot be parsed and no rows recovered.
def parse_llm_response(
    response_text: str,
    expected_count: int,
    target_languages: list[str] | None = None,
    finish_reason: str | None = None,
) -> dict[str, dict[str, str]]:
    """Parse LLM JSON response into dict of row_id -> translations."""
    try:
        data = json.loads(response_text)
    except json.JSONDecodeError:
        data = _recover_from_markdown(response_text)

    if data is None:
        data = _recover_truncated_rows(response_text, expected_count, finish_reason)

    if data is None:
        logger.explore("No complete rows found in truncated response",
                       extra={"src": "executor", "response_preview": response_text[:1000]})
        raise ValueError("LLM response was not valid JSON")

    rows = data.get("rows", [])
    if not isinstance(rows, list):
        raise ValueError("LLM response missing 'rows' array")

    translations: dict[str, dict[str, str]] = {}
    for item in rows:
        row_id = str(item.get("row_id", ""))
        if not row_id:
            continue

        detected_lang = str(item.get("detected_source_language", "und")) if item.get("detected_source_language") else "und"
        result: dict[str, str] = {"detected_source_language": detected_lang}

        has_language_data = False
        if target_languages:
            for lang_code in target_languages:
                lang_val = item.get(lang_code)
                if lang_val is not None and str(lang_val).strip():
                    result[lang_code] = str(lang_val)
                    has_language_data = True

        if not has_language_data:
            translation = item.get("translation")
            if translation is not None:
                result["translation"] = str(translation)
                has_language_data = True

        if has_language_data:
            translations[row_id] = result

    if len(translations) < expected_count:
        logger.explore(f"LLM returned fewer translations expected={expected_count} got={len(translations)}")

    return translations
# #endregion parse_llm_response


# #region _recover_from_markdown [C:1] [TYPE Function]
def _recover_from_markdown(response_text: str) -> dict | None:
    """Try to extract JSON from markdown code block."""
    match = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', response_text, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group(1))
            rows = data.get("rows", [])
            if isinstance(rows, list) and rows:
                logger.reason("Parsed JSON from markdown code block", {"rows": len(rows)})
            return data
        except json.JSONDecodeError:
            pass
    return None
# #endregion _recover_from_markdown


# #region _recover_truncated_rows [C:1] [TYPE Function]
def _recover_truncated_rows(
    response_text: str, expected_count: int, finish_reason: str | None,
) -> dict | None:
    """Try to recover complete rows from truncated JSON response."""
    logger.explore("LLM truncated, trying partial row recovery",
                   extra={"src": "executor", "finish_reason": finish_reason, "response_length": len(response_text)})
    rows_match = re.findall(r'\{\s*"row_id"\s*:\s*(?:\d+|"\d+").*?\}\s*', response_text, re.DOTALL)
    if rows_match:
        partial_rows = []
        for row_text in rows_match:
            try:
                row_data = json.loads(row_text)
                partial_rows.append(row_data)
            except json.JSONDecodeError:
                continue
        if partial_rows:
            logger.explore(f"Recovered {len(partial_rows)}/{expected_count} complete rows from truncated response",
                           extra={"src": "executor"})
            return {"rows": partial_rows}
    return None
# #endregion _recover_truncated_rows
# #endregion LLMResponseParser
