import re
import requests
import json
import html
from typing import List, Tuple

class Translator:
    def __init__(self, ollama_url: str = "http://localhost:11434", model: str = "gemma3:12b"):
        self.ollama_url = ollama_url
        self.model = model
        self.api_endpoint = f"{ollama_url}/api/generate"

    def _clean_input(self, text: str) -> str:
        """Clean input text by removing HTML tags, unescaping entities, and markdown code blocks"""
        # Remove all markdown code block markers
        text = re.sub(r'```', '', text)
        # Unescape HTML entities
        text = html.unescape(text)
        # Remove HTML tags
        text = re.sub(r'<[^>]+>', '', text)
        return text.strip()

    def _extract_segments(self, text: str) -> list:
        """Split text into HTML tags and translatable text segments.

        Returns a list of tuples: ('tag', content) or ('text', content)
        """
        # Match HTML tags (including self-closing) and HTML entities
        pattern = r'(<[^>]+>|&[a-zA-Z]+;|&#\d+;|&#x[0-9a-fA-F]+;)'
        parts = re.split(pattern, text)

        segments = []
        for part in parts:
            if not part:
                continue
            if re.match(r'^<[^>]+>$', part) or re.match(r'^&[a-zA-Z]+;$|^&#\d+;$|^&#x[0-9a-fA-F]+;$', part):
                segments.append(('tag', part))
            else:
                segments.append(('text', part))
        return segments

    def _translate_plain_text(self, text: str, target_language: str) -> tuple[str, dict]:
        """Translate plain text (no HTML) via the LLM."""
        prompt = f"""Translate this text to {target_language}. Provide only the translated text, no explanations or additional content:

{text}"""

        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.2,
                "top_p": 0.9,
                "max_tokens": 2000,
                "stop": ["IMPORTANT", "Instructions:", "Text to translate:"]
            }
        }

        response = requests.post(
            self.api_endpoint,
            json=payload,
            headers={"Content-Type": "application/json"},
            timeout=60
        )

        if response.status_code == 200:
            result = response.json()
            translated_text = result.get('response', '').strip()
            cleaned = self._clean_response(translated_text, preserve_html=False)

            # Remove original text if prepended
            if cleaned.strip().startswith(text.strip()):
                cleaned = cleaned.strip()[len(text.strip()):].strip()

            # If translation failed (returned original or empty), retry
            if cleaned.strip() == text.strip() or not cleaned.strip():
                prompt_retry = f"""Translate this text to {target_language}. Translate everything, including proper names and titles. Provide only the translated text:

{text}"""
                payload_retry = payload.copy()
                payload_retry["prompt"] = prompt_retry
                response_retry = requests.post(
                    self.api_endpoint,
                    json=payload_retry,
                    headers={"Content-Type": "application/json"},
                    timeout=60
                )
                if response_retry.status_code == 200:
                    result_retry = response_retry.json()
                    translated_text = result_retry.get('response', '').strip()
                    cleaned = self._clean_response(translated_text, preserve_html=False)
                    if cleaned.strip().startswith(text.strip()):
                        cleaned = cleaned.strip()[len(text.strip()):].strip()
                    token_usage = {
                        "prompt_eval_count": result.get('prompt_eval_count', 0) + result_retry.get('prompt_eval_count', 0),
                        "eval_count": result.get('eval_count', 0) + result_retry.get('eval_count', 0)
                    }
                    if cleaned.strip() == text.strip() or not cleaned.strip():
                        return text, token_usage
                    return cleaned, token_usage

            token_usage = {
                "prompt_eval_count": result.get('prompt_eval_count', 0),
                "eval_count": result.get('eval_count', 0)
            }
            return cleaned, token_usage
        else:
            return text, {"prompt_eval_count": 0, "eval_count": 0}

    def translate_text(self, text: str, target_language: str = "Hindi", preserve_html: bool = False) -> tuple[str, dict]:
        try:
            if preserve_html:
                # Check if text actually contains HTML
                if re.search(r'<[^>]+>', text):
                    # Split into HTML tags and text segments
                    segments = self._extract_segments(text)

                    # Collect all translatable text segments
                    text_parts = []
                    text_indices = []
                    for idx, (seg_type, content) in enumerate(segments):
                        if seg_type == 'text' and content.strip():
                            text_parts.append(content)
                            text_indices.append(idx)

                    if not text_parts:
                        # No translatable text found, return as-is
                        return text, {"prompt_eval_count": 0, "eval_count": 0}

                    # Translate all text parts together for better context
                    combined_text = "\n|||SEPARATOR|||\n".join(text_parts)
                    translated_combined, token_usage = self._translate_plain_text(combined_text, target_language)

                    # Split translated text back
                    translated_parts = translated_combined.split("|||SEPARATOR|||")
                    translated_parts = [p.strip() for p in translated_parts]

                    # If split count doesn't match, fall back to translating one by one
                    if len(translated_parts) != len(text_parts):
                        total_tokens = {"prompt_eval_count": token_usage.get("prompt_eval_count", 0), "eval_count": token_usage.get("eval_count", 0)}
                        for i, idx in enumerate(text_indices):
                            part_translated, part_tokens = self._translate_plain_text(text_parts[i], target_language)
                            # Preserve leading/trailing whitespace from original
                            original = segments[idx][1]
                            leading = original[:len(original) - len(original.lstrip())]
                            trailing = original[len(original.rstrip()):]
                            segments[idx] = ('text', leading + part_translated.strip() + trailing)
                            total_tokens["prompt_eval_count"] += part_tokens.get("prompt_eval_count", 0)
                            total_tokens["eval_count"] += part_tokens.get("eval_count", 0)

                        result = ''.join(content for _, content in segments)
                        return result, total_tokens

                    # Map translated parts back to segments
                    for i, idx in enumerate(text_indices):
                        if i < len(translated_parts):
                            original = segments[idx][1]
                            leading = original[:len(original) - len(original.lstrip())]
                            trailing = original[len(original.rstrip()):]
                            segments[idx] = ('text', leading + translated_parts[i].strip() + trailing)

                    result = ''.join(content for _, content in segments)
                    return result, token_usage
                else:
                    # No HTML tags found, translate normally
                    cleaned_text = text
                    translated, token_usage = self._translate_plain_text(cleaned_text, target_language)
                    return translated, token_usage
            else:
                cleaned_text = self._clean_input(text)
                translated, token_usage = self._translate_plain_text(cleaned_text, target_language)
                return translated, token_usage

        except Exception:
            return text, {"prompt_eval_count": 0, "eval_count": 0}

    def _clean_response(self, text: str, preserve_html: bool = False) -> str:
        """Clean up LLM response to remove instruction bleeding and artifacts"""
        if not preserve_html:
            # Remove HTML tags and unescape entities
            text = html.unescape(text)
            text = re.sub(r'<[^>]+>', '', text)
        # Remove common instruction artifacts
        cleanup_patterns = [
             r'IMPORTANT INSTRUCTIONS?:.*?(?=\n|$)',
            r'Instructions?:.*?(?=\n|$)',
            r'Text to translate:.*?(?=\n|$)',
            r'Hindi translation:.*?(?=\n|$)',
            r'[A-Za-z]+ translation:.*?(?=\n|$)',
            r"Here's the.*?(?=\n|$)",
            r'keeping the HTML tags.*?(?=\n|$)',
            r'\*\*.*?\*\*',  # Remove bold markdown
            r'Explanation:.*?(?=\n|$)',
            r'\* .*?(?=\n|$)',  # Remove bullet points
            r'- Keep all HTML.*?(?=\n|$)',
            r'- Only translate.*?(?=\n|$)',
            r'- Preserve.*?(?=\n|$)',
            r'- If there are.*?(?=\n|$)',
            r'<[^>]*?\s+[^>]*?>',  # Fix broken HTML tags with spaces
            r'Original:.*?(?=\n|$)',
            r'So, if the original text was.*?(?=\n|$)',
            r'The translated version would be:.*?(?=\n|$)',
            r'ye word mera name variable mia add ho raha hai.*?(?=\n|$)',  # Remove unwanted Hindi text
        ]

        for pattern in cleanup_patterns:
            text = re.sub(pattern, '', text, flags=re.IGNORECASE | re.MULTILINE)

        # Fix common HTML tag issues
        text = re.sub(r'<\s+([^>]+)\s*>', r'<\1>', text)  # Remove spaces in tags
        text = re.sub(r'<([^>]+)\s+>', r'<\1>', text)     # Remove trailing spaces in tags

        # Clean up multiple newlines and extra whitespace
        text = re.sub(r'\n\s*\n\s*\n', '\n\n', text)
        text = text.strip()

        return text

    def count_tokens(self, text: str) -> int:
        """
        Estimate token count for given text
        Basic estimation: ~4 characters per token for English
        """
        if not text:
            return 0
        return len(text) // 4

    def translate_array(self, sentences: List[str], target_language: str = "Hindi", preserve_html_list: List[bool] = None) -> tuple[List[str], dict]:
        translated_sentences = []
        total_prompt_tokens = 0
        total_completion_tokens = 0

        for i, sentence in enumerate(sentences):
            p_html = preserve_html_list[i] if preserve_html_list else False
            translated, token_usage = self.translate_text(sentence, target_language, p_html)
            translated_sentences.append(translated)
            total_prompt_tokens += token_usage.get('prompt_eval_count', 0)
            total_completion_tokens += token_usage.get('eval_count', 0)

        token_summary = {
            "input_tokens": total_prompt_tokens,
            "output_tokens": total_completion_tokens,
            "total_tokens": total_prompt_tokens + total_completion_tokens
        }

        return translated_sentences, token_summary