# #region LLMTranslationService [C:4] [TYPE Module] [SEMANTICS translate, llm, call, orchestrate, retry]
# @BRIEF LLM interaction for batch translation: call provider with retry, handle truncation
#         by recursive splitting, enforce dictionary post-processing. Orchestrates HTTP calls
#         (_llm_http) and response parsing (_llm_parse).
#         Language detection is now handled locally (lingua) — LLM prompt no longer
#         requests detected_source_language; local _detected_lang takes priority.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [TranslationRecord], [TranslationLanguage]
# @RELATION DEPENDS_ON -> [ContextAwarePromptBuilder]
# @RELATION DEPENDS_ON -> [EXT:method:_llm_http], [_llm_parse]
# @PRE DB session is available. LLM provider is configured on the job.
# @POST TranslationRecord rows created for LLM-processed rows (success/fail/skip).
# @SIDE_EFFECT HTTP calls to LLM provider API; DB writes.
# @RATIONALE Core orchestration logic — HTTP and parse logic extracted to _llm_http.py and
#            _llm_parse.py to meet INV_7 module limit (< 400 lines).
# @REJECTED Single monolithic call_llm_for_batch at 327 lines — split into focused sub-methods.

from datetime import UTC, datetime
import json
import time
from typing import Any
import uuid

from sqlalchemy.orm import Session

from ...core.logger import belief_scope, logger
from ...models.translate import TranslationBatch, TranslationJob, TranslationLanguage, TranslationRecord
from ...services.llm_prompt_templates import render_prompt
from ...services.llm_provider import LLMProviderService
from ._llm_http import call_openai_compatible
from ._llm_parse import parse_llm_response
from ._utils import _enforce_dictionary
from .preview import DEFAULT_EXECUTION_PROMPT_TEMPLATE
from .prompt_builder import ContextAwarePromptBuilder

MAX_RETRIES_PER_BATCH = 3


# #region LLMTranslationService [C:4] [TYPE Class]
# @BRIEF Call LLM, handle retry/truncation, parse response, persist records.
class LLMTranslationService:
    """LLM interaction for batch translation with retry, truncation handling, and parsing."""

    def __init__(self, db: Session) -> None:
        self.db = db

    # #region call_llm_for_batch [C:3] [TYPE Function] [SEMANTICS translate, llm, batch, orchestrate]
    # @BRIEF Call LLM for a batch of rows requiring translation. Parse and persist results.
    # @PRE job has valid provider_id. batch_rows is non-empty.
    # @POST Returns dict with successful/failed/skipped counts.
    # @SIDE_EFFECT HTTP call to LLM provider; DB writes.
    def call_llm_for_batch(
        self, job: TranslationJob, run_id: str,
        batch_rows: list[dict[str, Any]], dict_matches: list[dict[str, Any]],
        batch_id: str, max_tokens: int = 8192, _recursion_depth: int = 0,
    ) -> dict[str, int]:
        """Call LLM for a batch of rows; parse response; create records."""
        with belief_scope("LLMTranslationService.call_llm_for_batch"):
            dictionary_section = self._build_dictionary_section(dict_matches, batch_rows)
            target_languages = self._resolve_target_languages(job)
            prompt = self._build_prompt(job, batch_rows, dictionary_section, target_languages)

            llm_response, finish_reason, retries, last_error = self._call_llm_with_retry(
                job, prompt, batch_id, max_tokens,
            )
            if llm_response is None:
                return self._handle_llm_failure(batch_rows, run_id, batch_id, retries, last_error)

            if finish_reason == "length" and len(batch_rows) >= 2 and run_id:
                if _recursion_depth < MAX_RETRIES_PER_BATCH:
                    return self._split_and_retry(job, run_id, batch_rows, dict_matches,
                                                  batch_id, max_tokens, _recursion_depth, retries)
                logger.explore("Truncation recursion depth exceeded", {"batch_id": batch_id, "depth": _recursion_depth})

            try:
                translations = parse_llm_response(llm_response, len(batch_rows), target_languages=target_languages)
            except ValueError as e:
                return self._handle_parse_failure(batch_rows, run_id, batch_id, retries, e)

            return self._create_records_from_translations(
                batch_rows, run_id, batch_id, target_languages, translations, dict_matches, retries,
            )
    # #endregion call_llm_for_batch

    def _build_dictionary_section(self, dict_matches, batch_rows) -> str:
        if not dict_matches:
            return ""
        row_context = batch_rows[0].get("source_data") if batch_rows else None
        annotated = ContextAwarePromptBuilder.build_context_entries(dict_matches, row_context)
        return "Terminology dictionary (use these translations when applicable):\n" + \
            "\n".join(f"- {a}" for a in annotated) + "\n\n"

    @staticmethod
    def _resolve_target_languages(job):
        langs = job.target_languages or [job.target_dialect or "en"]
        return [str(langs)] if not isinstance(langs, list) else langs

    @staticmethod
    def _build_prompt(job, batch_rows, dictionary_section, target_languages):
        target_languages_str = ", ".join(target_languages)
        rows_json = json.dumps([
            {"row_id": str(row.get("row_index", idx)), "text": row.get("source_text", "")}
            for idx, row in enumerate(batch_rows)
        ], indent=2)
        return render_prompt(DEFAULT_EXECUTION_PROMPT_TEMPLATE, {
            "source_language": job.source_dialect or "SQL",
            "target_language": target_languages_str,
            "target_languages": target_languages_str,
            "source_dialect": job.source_dialect or "",
            "target_dialect": job.target_dialect or "",
            "translation_column": job.translation_column or "",
            "dictionary_section": dictionary_section,
            "rows_json": rows_json,
            "row_count": str(len(batch_rows)),
        })

    def _call_llm_with_retry(self, job, prompt, batch_id, max_tokens):
        llm_response = None
        last_error = None
        retries = 0
        finish_reason = None
        for attempt in range(1, MAX_RETRIES_PER_BATCH + 1):
            try:
                llm_response, finish_reason = self.call_llm(job, prompt, max_tokens=max_tokens)
                break
            except Exception as e:
                last_error = str(e)
                retries += 1
                logger.explore(f"LLM call failed (attempt {attempt})", {"batch_id": batch_id, "error": last_error})
                if attempt < MAX_RETRIES_PER_BATCH:
                    time.sleep(2 ** attempt)
        return llm_response, finish_reason, retries, last_error

    def _handle_llm_failure(self, batch_rows, run_id, batch_id, retries, last_error):
        for row in batch_rows:
            self.db.add(TranslationRecord(
                id=str(uuid.uuid4()), batch_id=batch_id, run_id=run_id,
                source_sql=row.get("source_text", ""), target_sql=None,
                source_object_type="table_row", source_object_id=row.get("row_index"),
                source_object_name=row.get("source_object_name", ""),
                source_data=row.get("source_data"), status="FAILED",
                error_message=f"LLM call failed after {retries} retries: {last_error}",
            ))
        return {"successful": 0, "failed": len(batch_rows), "skipped": 0, "retries": retries}

    def _split_and_retry(self, job, run_id, batch_rows, dict_matches, batch_id, max_tokens, depth, retries):
        mid = len(batch_rows) // 2
        logger.explore("LLM output truncated — splitting batch",
                       {"batch_id": batch_id, "batch_size": len(batch_rows), "split_at": mid, "depth": depth})

        # Create proper TranslationBatch rows for child halves (fixes FK violation
        # where _L/_R string suffixes referenced non-existent translation_batches rows).
        # @RATIONALE Previous code appended "_L"/"_R" to batch_id string, violating
        #            FK translation_records.batch_id → translation_batches.id.
        # @REJECTED Continuing with string-suffixed batch_ids — causes FK violation
        #           when any child record is flushed.
        left_batch = TranslationBatch(
            id=str(uuid.uuid4()), run_id=run_id, batch_index=-1,
            status="RUNNING", total_records=len(batch_rows[:mid]),
            started_at=datetime.now(UTC),
        )
        right_batch = TranslationBatch(
            id=str(uuid.uuid4()), run_id=run_id, batch_index=-1,
            status="RUNNING", total_records=len(batch_rows[mid:]),
            started_at=datetime.now(UTC),
        )
        self.db.add_all([left_batch, right_batch])
        self.db.flush()

        left = self.call_llm_for_batch(job, run_id, batch_rows[:mid], dict_matches,
                                        left_batch.id, max_tokens, depth + 1)
        right = self.call_llm_for_batch(job, run_id, batch_rows[mid:], dict_matches,
                                         right_batch.id, max_tokens, depth + 1)

        # Finalise child batch stats
        left_batch.completed_at = datetime.now(UTC)
        right_batch.completed_at = datetime.now(UTC)
        left_batch.successful_records = left["successful"]
        right_batch.successful_records = right["successful"]
        left_batch.failed_records = left["failed"]
        right_batch.failed_records = right["failed"]
        left_batch.status = "COMPLETED" if left["failed"] == 0 else "COMPLETED_WITH_ERRORS"
        right_batch.status = "COMPLETED" if right["failed"] == 0 else "COMPLETED_WITH_ERRORS"
        self.db.flush()

        return {"successful": left["successful"] + right["successful"],
                "failed": left["failed"] + right["failed"],
                "skipped": left["skipped"] + right["skipped"],
                "retries": retries + left.get("retries", 0) + right.get("retries", 0)}

    def _handle_parse_failure(self, batch_rows, run_id, batch_id, retries, error):
        for row in batch_rows:
            self.db.add(TranslationRecord(
                id=str(uuid.uuid4()), batch_id=batch_id, run_id=run_id,
                source_sql=row.get("source_text", ""), target_sql=None,
                source_object_type="table_row", source_object_id=row.get("row_index"),
                source_object_name=row.get("source_object_name", ""),
                source_data=row.get("source_data"), status="SKIPPED",
                error_message=f"LLM parse failure: {error}",
            ))
        return {"successful": 0, "failed": 0, "skipped": len(batch_rows), "retries": retries}

    def _create_records_from_translations(self, batch_rows, run_id, batch_id, target_languages, translations, dict_matches, retries):
        successful = failed = skipped = 0
        for row in batch_rows:
            row_id = str(row.get("row_index", ""))
            td = translations.get(row_id)
            source_text = row.get("source_text", "")
            # ★ Local detection (from _batch_proc._detect_languages) takes priority.
            #    Fallback to LLM response field for backward compatibility.
            detected_lang = row.get("_detected_lang", "und") or "und"
            if detected_lang == "und" and td:
                detected_lang = td.get("detected_source_language", "und") or "und"
            if td is None:
                skipped += 1
                self._add_skipped(row, run_id, batch_id, source_text, "NULL translation")
                continue
            plv = self._extract_per_lang_values(td, target_languages)
            if dict_matches and source_text:
                _enforce_dictionary(source_text, plv, dict_matches, batch_id, row_id)
            if not plv:
                skipped += 1
                self._add_skipped(row, run_id, batch_id, source_text, "Empty translation")
                continue
            successful += 1
            primary = next(iter(plv.values()), "")
            record = TranslationRecord(
                id=str(uuid.uuid4()), batch_id=batch_id, run_id=run_id,
                source_sql=source_text, target_sql=primary,
                source_object_type="table_row", source_object_id=row.get("row_index"),
                source_object_name=row.get("source_object_name", ""),
                source_data=row.get("source_data"), source_hash=row.get("_source_hash"),
                status="SUCCESS",
            )
            self.db.add(record)
            for lang_code in target_languages:
                if detected_lang != "und" and str(lang_code).lower() == str(detected_lang).lower():
                    continue
                val = plv.get(lang_code, "")
                needs_review = (detected_lang == "und")
                if needs_review:
                    logger.explore("undetected language", {"record_id": row_id, "language_code": lang_code, "text": source_text[:100]})
                self.db.add(TranslationLanguage(
                    id=str(uuid.uuid4()), record_id=record.id, language_code=lang_code,
                    source_language_detected=detected_lang, translated_value=val or "",
                    final_value=val or "", status="translated", needs_review=needs_review,
                ))
        return {"successful": successful, "failed": failed, "skipped": skipped, "retries": retries}

    @staticmethod
    def _extract_per_lang_values(td, target_languages):
        plv = {}
        has_any = False
        for lc in target_languages:
            lv = td.get(lc)
            if lv is not None and str(lv).strip():
                plv[lc] = str(lv)
                has_any = True
        if not has_any:
            t = td.get("translation", "")
            if t.strip():
                plv[target_languages[0]] = t
                has_any = True
        return plv if has_any else {}

    def _add_skipped(self, row, run_id, batch_id, source_text, reason):
        self.db.add(TranslationRecord(
            id=str(uuid.uuid4()), batch_id=batch_id, run_id=run_id,
            source_sql=source_text, target_sql="",
            source_object_type="table_row", source_object_id=row.get("row_index"),
            source_object_name=row.get("source_object_name", ""),
            source_data=row.get("source_data"), status="SKIPPED",
            error_message=reason,
        ))

    # #region call_llm [C:3] [TYPE Function] [SEMANTICS translate, llm, call]
    # @BRIEF Route to provider-specific LLM call implementation.
    def call_llm(self, job: TranslationJob, prompt: str, max_tokens: int = 8192) -> tuple[str, str | None]:
        """Call the configured LLM provider with the batch prompt."""
        with belief_scope("LLMTranslationService.call_llm"):
            if not job.provider_id:
                raise ValueError("Job has no LLM provider configured")
            provider_svc = LLMProviderService(self.db)
            provider = provider_svc.get_provider(job.provider_id)
            if not provider:
                raise ValueError(f"LLM provider '{job.provider_id}' not found")
            api_key = provider_svc.get_decrypted_api_key(job.provider_id)
            if not api_key:
                raise ValueError(f"Could not decrypt API key for provider '{job.provider_id}'")
            model = provider.default_model or "gpt-4o-mini"
            provider_type = provider.provider_type.lower() if provider.provider_type else "openai"
            disable_reasoning = getattr(job, 'disable_reasoning', False)

            if provider_type not in ("openai", "openai_compatible", "openrouter", "kilo", "litellm"):
                raise ValueError(f"Unsupported provider type '{provider_type}'")
            return call_openai_compatible(
                base_url=provider.base_url, api_key=api_key, model=model, prompt=prompt,
                provider_type=provider_type, max_tokens=max_tokens, disable_reasoning=disable_reasoning,
            )
    # #endregion call_llm

    # -- Static methods delegated to sub-modules for backward compat --
    @staticmethod
    def call_openai_compatible(*a, **kw):
        return call_openai_compatible(*a, **kw)

    @staticmethod
    def _parse_llm_response(*a, **kw):
        return parse_llm_response(*a, **kw)
# #endregion LLMTranslationService
# #endregion LLMTranslationService
