import logging
import re
from functools import lru_cache
from typing import Any, Dict, List, Optional, Tuple

NL2SQL_MAX_CONTEXT_TOKENS = 4000
CHAT_HISTORY_MAX_TOKENS = 2000
SCHEMA_CONTEXT_MAX_TOKENS = 2000
FEW_SHOT_MAX_TOKENS = 1500

logger = logging.getLogger(__name__)


class TokenOptimizer:

    def __init__(self, max_tokens: int = 4000):
        self._max_tokens = max_tokens
        self._encoder = None

    @property
    def encoder(self):
        if self._encoder is None:
            try:
                import tiktoken
                self._encoder = tiktoken.get_encoding("cl100k_base")
            except ImportError:
                logger.warning("tiktoken not available, using approximate token counting")
                self._encoder = _ApproximateEncoder()
        return self._encoder

    def count_tokens(self, text: str) -> int:
        if not text:
            return 0
        return self._count_tokens_cached(text)

    @lru_cache(maxsize=1000)
    def _count_tokens_cached(self, text: str) -> int:
        return len(self.encoder.encode(text))

    def truncate_chat_history(
        self,
        messages: List[Dict[str, str]],
        max_tokens: Optional[int] = None,
    ) -> Tuple[List[Dict[str, str]], int]:
        if not messages:
            return [], 0

        budget = max_tokens or CHAT_HISTORY_MAX_TOKENS
        token_counts = [(msg, self.count_tokens(msg.get("content", ""))) for msg in messages]

        kept = []
        total_tokens = 0

        for msg, tokens in reversed(token_counts):
            if total_tokens + tokens > budget and kept:
                break
            kept.insert(0, msg)
            total_tokens += tokens

        kept = self._preserve_message_pairs(kept, messages)
        final_tokens = sum(self.count_tokens(m.get("content", "")) for m in kept)
        return kept, final_tokens

    def _preserve_message_pairs(
        self,
        kept: List[Dict[str, str]],
        original: List[Dict[str, str]],
    ) -> List[Dict[str, str]]:
        if not kept or len(kept) < 2:
            return kept

        if kept[0].get("role") == "assistant" and len(original) > 1:
            idx = original.index(kept[0]) if kept[0] in original else -1
            if idx > 0 and original[idx - 1].get("role") == "user":
                kept.insert(0, original[idx - 1])

        return kept

    def truncate_context(
        self,
        context: str,
        max_tokens: Optional[int] = None,
        preserve_start: bool = True,
    ) -> Tuple[str, int]:
        if not context:
            return "", 0

        budget = max_tokens or self._max_tokens
        current_tokens = self.count_tokens(context)

        if current_tokens <= budget:
            return context, current_tokens

        sentences = re.split(r'(?<=[.!?])\s+', context)
        if len(sentences) <= 1:
            return self._hard_truncate(context, budget, preserve_start)

        if preserve_start:
            return self._truncate_from_end(sentences, budget)
        return self._truncate_from_start(sentences, budget)

    def _truncate_from_end(self, sentences: List[str], budget: int) -> Tuple[str, int]:
        result = []
        total = 0
        for sentence in sentences:
            tokens = self.count_tokens(sentence)
            if total + tokens > budget and result:
                break
            result.append(sentence)
            total += tokens
        text = " ".join(result) + "..."
        return text, self.count_tokens(text)

    def _truncate_from_start(self, sentences: List[str], budget: int) -> Tuple[str, int]:
        result = []
        total = 0
        for sentence in reversed(sentences):
            tokens = self.count_tokens(sentence)
            if total + tokens > budget and result:
                break
            result.insert(0, sentence)
            total += tokens
        text = "..." + " ".join(result)
        return text, self.count_tokens(text)

    def _hard_truncate(self, text: str, budget: int, preserve_start: bool) -> Tuple[str, int]:
        tokens = self.encoder.encode(text)
        truncated_tokens = tokens[:budget - 1] if preserve_start else tokens[-(budget - 1):]
        truncated = self.encoder.decode(truncated_tokens)
        suffix = "..." if preserve_start else ""
        prefix = "..." if not preserve_start else ""
        result = prefix + truncated + suffix
        return result, self.count_tokens(result)

    def optimize_schema_context(
        self,
        schema: str,
        relevant_tables: Optional[List[str]] = None,
        max_tokens: Optional[int] = None,
    ) -> str:
        if not schema:
            return ""

        budget = max_tokens or SCHEMA_CONTEXT_MAX_TOKENS
        current = self.count_tokens(schema)
        if current <= budget:
            return schema

        table_blocks = self._parse_table_blocks(schema)
        if not table_blocks:
            truncated, _ = self.truncate_context(schema, budget)
            return truncated

        prioritized = []
        remaining = []
        relevant_lower = [t.lower() for t in (relevant_tables or [])]

        for name, block in table_blocks:
            if relevant_lower and name.lower() in relevant_lower:
                prioritized.append(block)
            else:
                remaining.append(block)

        result_parts = []
        total = 0

        for block in prioritized + remaining:
            tokens = self.count_tokens(block)
            if total + tokens > budget and result_parts:
                break
            result_parts.append(block)
            total += tokens

        return "\n\n".join(result_parts)

    def _parse_table_blocks(self, schema: str) -> List[Tuple[str, str]]:
        blocks = []
        current_name = ""
        current_lines: List[str] = []

        for line in schema.split("\n"):
            table_match = re.match(r'^Table:\s*(\w+)', line, re.IGNORECASE)
            if table_match:
                if current_name and current_lines:
                    blocks.append((current_name, "\n".join(current_lines)))
                current_name = table_match.group(1)
                current_lines = [line]
            elif current_name:
                current_lines.append(line)

        if current_name and current_lines:
            blocks.append((current_name, "\n".join(current_lines)))

        return blocks

    def select_few_shot_examples(
        self,
        examples: List[Dict[str, Any]],
        query: str,
        max_tokens: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        if not examples:
            return []

        budget = max_tokens or FEW_SHOT_MAX_TOKENS
        scored = [(ex, self._keyword_similarity(query, self._example_text(ex))) for ex in examples]
        scored.sort(key=lambda x: x[1], reverse=True)

        selected = []
        total = 0

        for ex, score in scored:
            ex_text = self._example_text(ex)
            tokens = self.count_tokens(ex_text)
            if total + tokens > budget and selected:
                break
            selected.append(ex)
            total += tokens

        return selected

    def _example_text(self, example: Dict[str, Any]) -> str:
        parts = []
        for key in ("query", "question", "input", "sql", "output", "answer", "context"):
            if key in example:
                parts.append(str(example[key]))
        return " ".join(parts) if parts else str(example)

    def _keyword_similarity(self, query: str, text: str) -> float:
        query_words = set(query.lower().split())
        text_words = set(text.lower().split())
        if not query_words:
            return 0.0
        intersection = query_words & text_words
        return len(intersection) / len(query_words)

    def optimize_for_llm_call(
        self,
        components: Dict[str, str],
        total_budget: Optional[int] = None,
    ) -> Dict[str, str]:
        budget = total_budget or NL2SQL_MAX_CONTEXT_TOKENS
        result = {}

        query = components.get("query", "")
        system_prompt = components.get("system_prompt", "")

        result["query"] = query
        result["system_prompt"] = system_prompt

        reserved = self.count_tokens(query) + self.count_tokens(system_prompt)
        remaining = max(0, budget - reserved)

        schema_budget = int(remaining * 0.40)
        few_shots_budget = int(remaining * 0.30)
        history_budget = remaining - schema_budget - few_shots_budget

        schema = components.get("schema", "")
        if schema:
            truncated, _ = self.truncate_context(schema, schema_budget)
            result["schema"] = truncated
        else:
            result["schema"] = ""

        few_shots = components.get("few_shots", "")
        if few_shots:
            truncated, _ = self.truncate_context(few_shots, few_shots_budget)
            result["few_shots"] = truncated
        else:
            result["few_shots"] = ""

        chat_history = components.get("chat_history", "")
        if chat_history:
            truncated, _ = self.truncate_context(chat_history, history_budget, preserve_start=False)
            result["chat_history"] = truncated
        else:
            result["chat_history"] = ""

        return result

    def get_optimization_stats(self, original: Dict[str, str], optimized: Dict[str, str]) -> Dict[str, Any]:
        original_tokens = sum(self.count_tokens(v) for v in original.values())
        optimized_tokens = sum(self.count_tokens(v) for v in optimized.values())
        reduction = original_tokens - optimized_tokens
        ratio = reduction / original_tokens if original_tokens > 0 else 0.0

        return {
            "original_tokens": original_tokens,
            "optimized_tokens": optimized_tokens,
            "tokens_saved": reduction,
            "reduction_ratio": round(ratio, 4),
        }


class _ApproximateEncoder:

    def encode(self, text: str) -> List[int]:
        return list(range(len(text) // 4))

    def decode(self, tokens: List[int]) -> str:
        return " " * (len(tokens) * 4)


_token_optimizer_instance: Optional[TokenOptimizer] = None


def get_token_optimizer(max_tokens: int = 4000) -> TokenOptimizer:
    global _token_optimizer_instance
    if _token_optimizer_instance is None:
        _token_optimizer_instance = TokenOptimizer(max_tokens=max_tokens)
    return _token_optimizer_instance


def reset_token_optimizer() -> None:
    global _token_optimizer_instance
    _token_optimizer_instance = None