"""JSON Translator tool for SDKRouter.

Translates JSON objects and text using LLM with smart caching.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING

if TYPE_CHECKING:
    from ..._config import SDKConfig

from .cache import TranslationCacheManager
from .detectors import LanguageDetector, ScriptDetector
from .utils import PromptBuilder, TextUtils

logger = logging.getLogger(__name__)


class TranslationError(Exception):
    """Translation error."""
    pass


class TranslatorResource:
    """
    JSON Translator tool (sync).

    Translates JSON objects and text using LLM with smart text-level caching.

    Example:
        ```python
        from sdkrouter import SDKRouter

        client = SDKRouter(api_key="your-key")

        # Translate JSON
        data = {"title": "Hello", "description": "World"}
        translated = client.translator.translate_json(data, target_language="es")
        # {"title": "Hola", "description": "Mundo"}

        # Translate text
        text = client.translator.translate("Hello, world!", target_language="ru")
        # "Привет, мир!"
        ```
    """

    def __init__(
        self,
        config: "SDKConfig",
        cache_dir: Optional[str] = None,
        default_model: str = "openai/gpt-4o-mini",
    ):
        """
        Initialize translator.

        Args:
            config: SDK configuration
            cache_dir: Optional cache directory path
            default_model: Default model for translation
        """
        self._config = config
        self._default_model = default_model

        # Initialize components
        self.cache = TranslationCacheManager(cache_dir=cache_dir)
        self.script_detector = ScriptDetector()
        self.language_detector = LanguageDetector()
        self.text_utils = TextUtils()
        self.prompt_builder = PromptBuilder()

        # Lazy-loaded OpenAI client
        self._openai_client = None

    @property
    def openai_client(self):
        """Get OpenAI client (lazy-loaded)."""
        if self._openai_client is None:
            from openai import OpenAI
            self._openai_client = OpenAI(
                api_key=self._config.api_key,
                base_url=self._config.llm_url,
                timeout=self._config.timeout,
            )
        return self._openai_client

    def detect_language(self, text: str) -> str:
        """
        Detect language of text.

        Args:
            text: Text to analyze

        Returns:
            Language code (ko, ja, zh, ru, en, etc.)
        """
        return self.script_detector.detect_language(text)

    def needs_translation(
        self,
        text: str,
        source_language: str,
        target_language: str
    ) -> bool:
        """
        Check if text needs translation.

        Args:
            text: Text to check
            source_language: Source language code
            target_language: Target language code

        Returns:
            True if translation is needed
        """
        return self.text_utils.needs_translation(
            text, source_language, target_language, self.script_detector
        )

    def translate(
        self,
        text: str,
        target_language: str = "en",
        source_language: str = "auto",
        fail_silently: bool = False,
        model: Optional[str] = None,
        temperature: float = 0.1,
    ) -> str:
        """
        Translate single text.

        Args:
            text: Text to translate
            target_language: Target language code
            source_language: Source language code ('auto' for detection)
            fail_silently: Don't raise exceptions on failure
            model: Optional model override
            temperature: LLM temperature (0-1)

        Returns:
            Translated text

        Raises:
            TranslationError: If translation fails and fail_silently is False
        """
        try:
            if not text or not text.strip():
                return text

            # Detect source language if auto
            actual_source = source_language
            if source_language == 'auto':
                actual_source = self.script_detector.detect_language(text)
                if actual_source == 'unknown':
                    actual_source = 'en'

            # Check if translation needed
            if not self.text_utils.needs_translation(text, actual_source, target_language, self.script_detector):
                return text

            # Check cache
            cached = self.cache.get(text, actual_source, target_language)
            if cached:
                logger.debug("Translation cache hit")
                return cached

            # Build prompt and call LLM
            prompt = self.prompt_builder.build_text_translation_prompt(
                text, actual_source, target_language
            )

            response = self.openai_client.chat.completions.create(
                model=model or self._default_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=4000,
            )

            content = response.choices[0].message.content
            translated = content.strip() if content else text

            # Cache result
            self.cache.set(text, actual_source, target_language, translated)

            logger.info("Translated text: %d -> %d chars", len(text), len(translated))
            return translated

        except Exception as e:
            error_msg = f"Translation failed: {e}"
            logger.error(error_msg)
            if not fail_silently:
                raise TranslationError(error_msg) from e
            return text

    def translate_json(
        self,
        data: Dict[str, Any],
        target_language: str = "en",
        source_language: str = "auto",
        fail_silently: bool = False,
        model: Optional[str] = None,
        temperature: float = 0.1,
    ) -> Dict[str, Any]:
        """
        Translate JSON object with smart caching.

        Translates ONLY values, never keys. Uses text-level caching
        to minimize LLM calls.

        Args:
            data: JSON object to translate
            target_language: Target language code
            source_language: Source language code ('auto' for detection)
            fail_silently: Don't raise exceptions on failure
            model: Optional model override
            temperature: LLM temperature (0-1)

        Returns:
            Translated JSON object

        Raises:
            TranslationError: If translation fails and fail_silently is False

        Example:
            ```python
            data = {
                "title": "Привет",
                "items": ["Мир", "Земля"]
            }
            result = translator.translate_json(data, target_language="en")
            # {"title": "Hello", "items": ["World", "Earth"]}
            ```
        """
        try:
            # Extract translatable texts
            translatable_texts = self._extract_translatable_texts(
                data, source_language, target_language
            )

            if not translatable_texts:
                logger.debug("No texts need translation in JSON")
                return data

            logger.info("Found %d texts to translate", len(translatable_texts))

            # Detect source language from first text if auto
            actual_source = source_language
            if source_language == 'auto' and translatable_texts:
                first_text = list(translatable_texts)[0]
                detected = self.language_detector.detect_language(first_text)
                if detected and detected != 'unknown':
                    actual_source = detected
                else:
                    actual_source = 'en'

            # Check cache for each text
            cached_translations: Dict[str, str] = {}
            uncached_texts: List[str] = []

            for text in translatable_texts:
                cached = self.cache.get(text, actual_source, target_language)
                if cached:
                    cached_translations[text] = cached
                else:
                    uncached_texts.append(text)

            logger.info("Cache: %d hits, %d misses", len(cached_translations), len(uncached_texts))

            # If everything cached, just apply
            if not uncached_texts:
                return self._apply_translations(data, cached_translations)

            # Create partial JSON with only uncached texts
            uncached_json = self._create_partial_json(data, uncached_texts)
            json_str = json.dumps(uncached_json, ensure_ascii=False, indent=2)

            # Build prompt
            prompt = self.prompt_builder.build_json_translation_prompt(
                json_str, actual_source, target_language
            )

            # Call LLM
            response = self.openai_client.chat.completions.create(
                model=model or self._default_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=4000,
            )

            content = response.choices[0].message.content
            translated_json_str = content.strip() if content else "{}"

            # Parse response
            try:
                # Remove markdown formatting
                if translated_json_str.startswith("```json"):
                    translated_json_str = translated_json_str.replace("```json", "").replace("```", "").strip()
                elif translated_json_str.startswith("```"):
                    translated_json_str = translated_json_str.replace("```", "").strip()

                translated_partial = json.loads(translated_json_str)

                # Extract new translations
                new_translations = self._extract_translations_by_comparison(
                    uncached_json, translated_partial, uncached_texts
                )

                # Cache new translations
                for orig, trans in new_translations.items():
                    self.cache.set(orig, actual_source, target_language, trans)

                # Combine all translations
                all_translations = {**cached_translations, **new_translations}

                # Apply to original data
                result = self._apply_translations(data, all_translations)

                logger.info("Translation complete: %d cached, %d new",
                           len(cached_translations), len(new_translations))
                return result

            except json.JSONDecodeError as e:
                logger.error("LLM returned invalid JSON: %s", e)
                if fail_silently:
                    return self._apply_translations(data, cached_translations)
                raise TranslationError(f"LLM returned invalid JSON: {e}") from e

        except TranslationError:
            raise
        except Exception as e:
            error_msg = f"JSON translation failed: {e}"
            logger.error(error_msg)
            if not fail_silently:
                raise TranslationError(error_msg) from e
            return data

    def _extract_translatable_texts(
        self,
        obj: Any,
        source_language: str,
        target_language: str
    ) -> Set[str]:
        """Extract texts that need translation from JSON object."""
        translatable_texts: Set[str] = set()

        def _extract_recursive(item):
            if isinstance(item, str):
                if self.text_utils.needs_translation(item, source_language, target_language):
                    translatable_texts.add(item)
            elif isinstance(item, list):
                for sub_item in item:
                    _extract_recursive(sub_item)
            elif isinstance(item, dict):
                for key, value in item.items():
                    # Note: we don't translate keys, only values
                    _extract_recursive(value)

        _extract_recursive(obj)
        return translatable_texts

    def _apply_translations(self, obj: Any, translations: Dict[str, str]) -> Any:
        """Apply translations to JSON object."""
        if isinstance(obj, str):
            return translations.get(obj, obj)
        elif isinstance(obj, list):
            return [self._apply_translations(item, translations) for item in obj]
        elif isinstance(obj, dict):
            return {
                key: self._apply_translations(value, translations)
                for key, value in obj.items()
            }
        else:
            return obj

    def _create_partial_json(self, data: Any, texts_to_include: List[str]) -> Any:
        """Create JSON containing only texts that need translation."""
        texts_set = set(texts_to_include)

        def filter_recursive(obj):
            if isinstance(obj, str):
                return obj if obj in texts_set else "SKIP_TRANSLATION"
            elif isinstance(obj, dict):
                result = {}
                for key, value in obj.items():
                    filtered = filter_recursive(value)
                    if self._contains_translatable(filtered, texts_set):
                        result[key] = filtered
                return result
            elif isinstance(obj, list):
                result = []
                for item in obj:
                    filtered = filter_recursive(item)
                    if self._contains_translatable(filtered, texts_set):
                        result.append(filtered)
                return result
            else:
                return obj

        return filter_recursive(data)

    def _contains_translatable(self, obj: Any, texts_set: set) -> bool:
        """Check if object contains any translatable text."""
        if isinstance(obj, str):
            return obj in texts_set
        elif isinstance(obj, dict):
            return any(self._contains_translatable(v, texts_set) for v in obj.values())
        elif isinstance(obj, list):
            return any(self._contains_translatable(item, texts_set) for item in obj)
        return False

    def _extract_translations_by_comparison(
        self,
        original: Any,
        translated: Any,
        uncached_texts: List[str]
    ) -> Dict[str, str]:
        """Extract translations by comparing original and translated JSON."""
        translations: Dict[str, str] = {}
        uncached_set = set(uncached_texts)

        def compare_recursive(orig, trans):
            if isinstance(orig, str) and isinstance(trans, str):
                if orig in uncached_set and orig != trans:
                    translations[orig] = trans
            elif isinstance(orig, list) and isinstance(trans, list):
                for o, t in zip(orig, trans):
                    compare_recursive(o, t)
            elif isinstance(orig, dict) and isinstance(trans, dict):
                for key in orig:
                    if key in trans:
                        compare_recursive(orig[key], trans[key])

        compare_recursive(original, translated)
        logger.debug("Extracted %d translations from LLM response", len(translations))
        return translations

    def get_stats(self) -> Dict[str, Any]:
        """Get translation statistics."""
        return self.cache.get_stats()

    def clear_cache(
        self,
        source_lang: Optional[str] = None,
        target_lang: Optional[str] = None
    ):
        """
        Clear translation cache.

        Args:
            source_lang: Optional source language (clears specific pair if both provided)
            target_lang: Optional target language
        """
        self.cache.clear(source_lang, target_lang)


class AsyncTranslatorResource:
    """
    JSON Translator tool (async).

    Async version of TranslatorResource for use in async contexts.

    Example:
        ```python
        from sdkrouter import AsyncSDKRouter

        client = AsyncSDKRouter(api_key="your-key")

        # Translate JSON
        data = {"title": "Hello", "description": "World"}
        translated = await client.translator.translate_json(data, target_language="es")
        ```
    """

    def __init__(
        self,
        config: "SDKConfig",
        cache_dir: Optional[str] = None,
        default_model: str = "openai/gpt-4o-mini",
    ):
        """
        Initialize async translator.

        Args:
            config: SDK configuration
            cache_dir: Optional cache directory path
            default_model: Default model for translation
        """
        self._config = config
        self._default_model = default_model

        # Initialize components (shared with sync)
        self.cache = TranslationCacheManager(cache_dir=cache_dir)
        self.script_detector = ScriptDetector()
        self.language_detector = LanguageDetector()
        self.text_utils = TextUtils()
        self.prompt_builder = PromptBuilder()

        # Lazy-loaded AsyncOpenAI client
        self._openai_client = None

    @property
    def openai_client(self):
        """Get AsyncOpenAI client (lazy-loaded)."""
        if self._openai_client is None:
            from openai import AsyncOpenAI
            self._openai_client = AsyncOpenAI(
                api_key=self._config.api_key,
                base_url=self._config.llm_url,
                timeout=self._config.timeout,
            )
        return self._openai_client

    def detect_language(self, text: str) -> str:
        """Detect language of text."""
        return self.script_detector.detect_language(text)

    def needs_translation(
        self,
        text: str,
        source_language: str,
        target_language: str
    ) -> bool:
        """Check if text needs translation."""
        return self.text_utils.needs_translation(
            text, source_language, target_language, self.script_detector
        )

    async def translate(
        self,
        text: str,
        target_language: str = "en",
        source_language: str = "auto",
        fail_silently: bool = False,
        model: Optional[str] = None,
        temperature: float = 0.1,
    ) -> str:
        """
        Translate single text (async).

        Args:
            text: Text to translate
            target_language: Target language code
            source_language: Source language code ('auto' for detection)
            fail_silently: Don't raise exceptions on failure
            model: Optional model override
            temperature: LLM temperature (0-1)

        Returns:
            Translated text
        """
        try:
            if not text or not text.strip():
                return text

            actual_source = source_language
            if source_language == 'auto':
                actual_source = self.script_detector.detect_language(text)
                if actual_source == 'unknown':
                    actual_source = 'en'

            if not self.text_utils.needs_translation(text, actual_source, target_language, self.script_detector):
                return text

            cached = self.cache.get(text, actual_source, target_language)
            if cached:
                return cached

            prompt = self.prompt_builder.build_text_translation_prompt(
                text, actual_source, target_language
            )

            response = await self.openai_client.chat.completions.create(
                model=model or self._default_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=4000,
            )

            content = response.choices[0].message.content
            translated = content.strip() if content else text
            self.cache.set(text, actual_source, target_language, translated)

            return translated

        except Exception as e:
            logger.error("Translation failed: %s", e)
            if not fail_silently:
                raise TranslationError(f"Translation failed: {e}") from e
            return text

    async def translate_json(
        self,
        data: Dict[str, Any],
        target_language: str = "en",
        source_language: str = "auto",
        fail_silently: bool = False,
        model: Optional[str] = None,
        temperature: float = 0.1,
    ) -> Dict[str, Any]:
        """
        Translate JSON object (async).

        Args:
            data: JSON object to translate
            target_language: Target language code
            source_language: Source language code ('auto' for detection)
            fail_silently: Don't raise exceptions on failure
            model: Optional model override
            temperature: LLM temperature (0-1)

        Returns:
            Translated JSON object
        """
        try:
            translatable_texts = self._extract_translatable_texts(
                data, source_language, target_language
            )

            if not translatable_texts:
                return data

            actual_source = source_language
            if source_language == 'auto' and translatable_texts:
                first_text = list(translatable_texts)[0]
                detected = self.language_detector.detect_language(first_text)
                if detected and detected != 'unknown':
                    actual_source = detected
                else:
                    actual_source = 'en'

            cached_translations: Dict[str, str] = {}
            uncached_texts: List[str] = []

            for text in translatable_texts:
                cached = self.cache.get(text, actual_source, target_language)
                if cached:
                    cached_translations[text] = cached
                else:
                    uncached_texts.append(text)

            if not uncached_texts:
                return self._apply_translations(data, cached_translations)

            uncached_json = self._create_partial_json(data, uncached_texts)
            json_str = json.dumps(uncached_json, ensure_ascii=False, indent=2)

            prompt = self.prompt_builder.build_json_translation_prompt(
                json_str, actual_source, target_language
            )

            response = await self.openai_client.chat.completions.create(
                model=model or self._default_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=4000,
            )

            content = response.choices[0].message.content
            translated_json_str = content.strip() if content else "{}"

            try:
                if translated_json_str.startswith("```json"):
                    translated_json_str = translated_json_str.replace("```json", "").replace("```", "").strip()
                elif translated_json_str.startswith("```"):
                    translated_json_str = translated_json_str.replace("```", "").strip()

                translated_partial = json.loads(translated_json_str)

                new_translations = self._extract_translations_by_comparison(
                    uncached_json, translated_partial, uncached_texts
                )

                for orig, trans in new_translations.items():
                    self.cache.set(orig, actual_source, target_language, trans)

                all_translations = {**cached_translations, **new_translations}
                return self._apply_translations(data, all_translations)

            except json.JSONDecodeError as e:
                if fail_silently:
                    return self._apply_translations(data, cached_translations)
                raise TranslationError(f"LLM returned invalid JSON: {e}") from e

        except TranslationError:
            raise
        except Exception as e:
            if not fail_silently:
                raise TranslationError(f"JSON translation failed: {e}") from e
            return data

    def _extract_translatable_texts(
        self,
        obj: Any,
        source_language: str,
        target_language: str
    ) -> Set[str]:
        """Extract texts that need translation."""
        translatable: Set[str] = set()

        def extract(item):
            if isinstance(item, str):
                if self.text_utils.needs_translation(item, source_language, target_language):
                    translatable.add(item)
            elif isinstance(item, list):
                for i in item:
                    extract(i)
            elif isinstance(item, dict):
                for v in item.values():
                    extract(v)

        extract(obj)
        return translatable

    def _apply_translations(self, obj: Any, translations: Dict[str, str]) -> Any:
        """Apply translations to JSON object."""
        if isinstance(obj, str):
            return translations.get(obj, obj)
        elif isinstance(obj, list):
            return [self._apply_translations(i, translations) for i in obj]
        elif isinstance(obj, dict):
            return {k: self._apply_translations(v, translations) for k, v in obj.items()}
        return obj

    def _create_partial_json(self, data: Any, texts_to_include: List[str]) -> Any:
        """Create JSON with only translatable texts."""
        texts_set = set(texts_to_include)

        def filter_obj(obj):
            if isinstance(obj, str):
                return obj if obj in texts_set else "SKIP_TRANSLATION"
            elif isinstance(obj, dict):
                result = {}
                for k, v in obj.items():
                    filtered = filter_obj(v)
                    if self._has_translatable(filtered, texts_set):
                        result[k] = filtered
                return result
            elif isinstance(obj, list):
                return [filter_obj(i) for i in obj if self._has_translatable(filter_obj(i), texts_set)]
            return obj

        return filter_obj(data)

    def _has_translatable(self, obj: Any, texts_set: set) -> bool:
        """Check if object has translatable content."""
        if isinstance(obj, str):
            return obj in texts_set
        elif isinstance(obj, dict):
            return any(self._has_translatable(v, texts_set) for v in obj.values())
        elif isinstance(obj, list):
            return any(self._has_translatable(i, texts_set) for i in obj)
        return False

    def _extract_translations_by_comparison(
        self,
        original: Any,
        translated: Any,
        uncached_texts: List[str]
    ) -> Dict[str, str]:
        """Extract translations by comparison."""
        translations: Dict[str, str] = {}
        uncached_set = set(uncached_texts)

        def compare(orig, trans):
            if isinstance(orig, str) and isinstance(trans, str):
                if orig in uncached_set and orig != trans:
                    translations[orig] = trans
            elif isinstance(orig, list) and isinstance(trans, list):
                for o, t in zip(orig, trans):
                    compare(o, t)
            elif isinstance(orig, dict) and isinstance(trans, dict):
                for k in orig:
                    if k in trans:
                        compare(orig[k], trans[k])

        compare(original, translated)
        return translations

    def get_stats(self) -> Dict[str, Any]:
        """Get translation statistics."""
        return self.cache.get_stats()

    def clear_cache(
        self,
        source_lang: Optional[str] = None,
        target_lang: Optional[str] = None
    ):
        """Clear translation cache."""
        self.cache.clear(source_lang, target_lang)
