import os
import re
from typing import Optional, Dict, Any, Union, Tuple
import httpx
from pydantic import BaseModel

# --- Private Data Protection ---

_PRIVATE_PATTERN = re.compile(r"\[PRIVATE\](.*?)\[/PRIVATE\]", re.DOTALL | re.IGNORECASE)

_AUTO_PROTECT_PATTERNS = [
    # Email addresses
    r"[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}",
    # Phone numbers (US-style: xxx-xxx-xxxx, (xxx) xxx-xxxx, +1xxxxxxxxxx)
    r"(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}",
    # SSN (xxx-xx-xxxx)
    r"\b\d{3}-\d{2}-\d{4}\b",
    # Credit card numbers (xxxx-xxxx-xxxx-xxxx or xxxxxxxxxxxx)
    r"\b(?:\d{4}[-\s]?){3}\d{4}\b",
    # Bearer tokens
    r"Bearer\s+[a-zA-Z0-9\-_.]+",
    # API key patterns (sk-, pk-, ak_, tk_ prefixed)
    r"\b(?:sk|pk|ak|tk)[-_][a-zA-Z0-9_\-]{16,}\b",
    # AWS keys (AKIA/ASIA + 16-20 chars)
    r"\b(?:AKIA|ASIA)[A-Z0-9]{16,20}\b",
]

_AUTO_PROTECT_COMPILED = [re.compile(p) for p in _AUTO_PROTECT_PATTERNS]


def _extract_private(text: str) -> Tuple[str, Dict[str, str]]:
    """
    Extract [PRIVATE]...[/PRIVATE] blocks from text and replace with indexed placeholders.
    Returns (sanitized_text, private_map) where private_map maps placeholders to original content.
    """
    private_map: Dict[str, str] = {}
    counter = 0

    def replacer(match: re.Match) -> str:
        nonlocal counter
        placeholder = f"__PRIVATE_{counter}__"
        private_map[placeholder] = match.group(1)
        counter += 1
        return placeholder

    sanitized = _PRIVATE_PATTERN.sub(replacer, text)
    return sanitized, private_map


def _restore_private(text: str, private_map: Dict[str, str]) -> str:
    """
    Restore __PRIVATE_N__ placeholders back to [PRIVATE]...[/PRIVATE] tagged content.
    """
    for placeholder, original in private_map.items():
        text = text.replace(placeholder, f"[PRIVATE]{original}[/PRIVATE]")
    return text


def auto_protect(text: str) -> str:
    """
    Automatically detect common sensitive patterns (emails, phone numbers, SSNs,
    credit cards, API keys, Bearer tokens) and wrap them in [PRIVATE] tags.

    This is an opt-in convenience utility. Call it before optimize() if desired:
        protected_text = auto_protect(raw_text)
        result = client.optimize(protected_text)
    """
    # Find all matches with their positions
    matches = []
    for pattern in _AUTO_PROTECT_COMPILED:
        for match in pattern.finditer(text):
            matches.append((match.start(), match.end(), match.group(0)))

    if not matches:
        return text

    # Sort by start pos descending so we can replace from end to start without shifting indices
    matches.sort(key=lambda m: m[0], reverse=True)

    # Remove overlapping matches (keep the first/longest)
    filtered = []
    min_end = float("inf")
    for start, end, value in matches:
        if end <= min_end:
            # Check the match is not already inside [PRIVATE] tags
            pre_context = text[max(0, start - 10):start]
            if "[PRIVATE]" not in pre_context.upper():
                filtered.append((start, end, value))
                min_end = start

    result = text
    for start, end, value in filtered:
        result = result[:start] + f"[PRIVATE]{value}[/PRIVATE]" + result[end:]

    return result


class OptimizationStats(BaseModel):
    original_tokens: int
    optimized_tokens: int
    reduction_percent: float
    latency_ms: float

class OptimizationResult(BaseModel):
    text: str
    stats: OptimizationStats

class Trunkate:
    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        self.api_key = api_key or os.getenv("TRUNKATE_API_KEY")
        if not self.api_key:
            raise ValueError(
                "The TRUNKATE_API_KEY environment variable is missing or empty; "
                "either provide it, or instantiate the Trunkate client with an api_key."
            )
        # Default to official API if not specified, can be overridden by env for testing
        self.base_url = base_url or os.getenv("TRUNKATE_API_URL", "https://api.trunkate.ai")
        # Ensure trailing slashes are stripped
        self.base_url = self.base_url.rstrip("/")
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        self.client = httpx.Client(base_url=self.base_url, headers=headers, timeout=10.0)

    def optimize(self, text: str, task: Optional[str] = None, budget: Optional[Union[int, str]] = None, model: str = "gpt-4o") -> OptimizationResult:
        """
        Optimize a prompt using the Trunkate AI API.
        
        [PRIVATE]...[/PRIVATE] tagged content is automatically extracted before 
        the API call (never sent to the server) and restored in the response.
        """
        # 1. Extract private content (zero-trust: private data never leaves the client)
        sanitized_text, private_map = _extract_private(text)
        sanitized_task = task
        task_private_map: Dict[str, str] = {}
        if task:
            sanitized_task, task_private_map = _extract_private(task)
        
        # 2. Build payload with sanitized text only
        payload: Dict[str, Any] = {
            "text": sanitized_text,
            "model": model
        }
        if sanitized_task:
            payload["task"] = sanitized_task
        if budget:
            payload["budget"] = budget

        response = self.client.post("/optimize", json=payload)
        response.raise_for_status()
        
        data = response.json()
        
        # 3. Restore private content in the response
        optimized_text = data["optimized_text"]
        optimized_text = _restore_private(optimized_text, private_map)
        optimized_text = _restore_private(optimized_text, task_private_map)
        
        return OptimizationResult(
            text=optimized_text,
            stats=OptimizationStats(
                original_tokens=data["original_tokens"],
                optimized_tokens=data["optimized_tokens"],
                reduction_percent=data["reduction_percent"],
                latency_ms=data["latency_ms"]
            )
        )

