"""
LLM client abstraction supporting OpenAI and Anthropic.

Provides a unified interface for classification requests with:
- Provider abstraction (OpenAI/Anthropic)
- Structured output (JSON mode)
- Retry handling
- Cost tracking
- Adaptive batch sizing based on context window
"""

from __future__ import annotations

import json
import logging
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, TypedDict

if TYPE_CHECKING:
    from reviewiq_pipeline.config import Config
    from reviewiq_pipeline.contracts import LLMClassificationResponse

logger = logging.getLogger(__name__)


# =============================================================================
# Exceptions
# =============================================================================


class PartialBatchResult(Exception):
    """
    Exception raised when batch JSON parsing partially fails but some results were recovered.

    Carries the partial results and indices of missing reviews so the caller can
    only reprocess the missing ones instead of the entire batch.
    """

    def __init__(
        self,
        message: str,
        partial_results: list[dict[str, Any]],
        missing_indices: list[int],
        metadata: dict[str, Any] | None = None,
    ):
        super().__init__(message)
        self.partial_results = partial_results
        self.missing_indices = missing_indices
        self.metadata = metadata or {}


# =============================================================================
# Model Context Windows and Token Estimation
# =============================================================================

MODEL_CONTEXT_WINDOWS = {
    # OpenAI models
    "gpt-4o": 128_000,
    "gpt-4o-mini": 128_000,
    "gpt-4-turbo": 128_000,
    "gpt-4": 8_192,
    "gpt-3.5-turbo": 16_385,
    # Anthropic models
    "claude-3-opus-20240229": 200_000,
    "claude-3-sonnet-20240229": 200_000,
    "claude-3-haiku-20240307": 200_000,
    "claude-3-5-sonnet-20241022": 200_000,
    "claude-sonnet-4-20250514": 200_000,
}

# Average tokens per character (rough estimate, varies by language)
CHARS_PER_TOKEN = 4

# Output tokens per review (classification response)
OUTPUT_TOKENS_PER_REVIEW = 450  # Conservative estimate


@dataclass
class BatchSizeCalculation:
    """Result of batch size calculation."""
    batch_size: int
    system_prompt_tokens: int
    avg_tokens_per_review: int
    output_tokens_reserved: int
    context_window: int
    utilization_target: float
    reasoning: str


@dataclass
class TokenStats:
    """Running statistics for token estimation."""
    total_reviews: int = 0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    min_review_tokens: int = 999999
    max_review_tokens: int = 0

    def update(self, review_tokens: int, output_tokens: int):
        """Update stats with new observation."""
        self.total_reviews += 1
        self.total_input_tokens += review_tokens
        self.total_output_tokens += output_tokens
        self.min_review_tokens = min(self.min_review_tokens, review_tokens)
        self.max_review_tokens = max(self.max_review_tokens, review_tokens)

    @property
    def avg_review_tokens(self) -> int:
        """Average tokens per review."""
        if self.total_reviews == 0:
            return 150  # Default estimate
        return self.total_input_tokens // self.total_reviews

    @property
    def avg_output_tokens(self) -> int:
        """Average output tokens per review."""
        if self.total_reviews == 0:
            return OUTPUT_TOKENS_PER_REVIEW
        return self.total_output_tokens // self.total_reviews


class BatchSizer:
    """
    Calculates optimal batch size based on context window and actual token usage.

    Adapts in real-time based on observed token counts from previous batches.
    """

    def __init__(
        self,
        model: str,
        system_prompt_tokens: int,
        target_utilization: float = 0.6,
    ):
        self.model = model
        self.system_prompt_tokens = system_prompt_tokens
        self.target_utilization = target_utilization
        self.context_window = MODEL_CONTEXT_WINDOWS.get(model, 128_000)
        self.stats = TokenStats()

    def estimate_tokens(self, text: str) -> int:
        """Estimate token count for text (fast approximation)."""
        # Simple heuristic: ~4 chars per token for English
        # More accurate would be to use tiktoken, but this is faster
        return max(1, len(text) // CHARS_PER_TOKEN)

    def calculate_batch_size(
        self,
        reviews: list[dict],
        fixed_size: int | None = None,
    ) -> BatchSizeCalculation:
        """
        Calculate optimal batch size for a set of reviews.

        Args:
            reviews: List of reviews with 'text' field
            fixed_size: If set, use this size (skip calculation)

        Returns:
            BatchSizeCalculation with recommended size and reasoning
        """
        if fixed_size and fixed_size > 0:
            return BatchSizeCalculation(
                batch_size=min(fixed_size, len(reviews)),
                system_prompt_tokens=self.system_prompt_tokens,
                avg_tokens_per_review=self.stats.avg_review_tokens,
                output_tokens_reserved=fixed_size * self.stats.avg_output_tokens,
                context_window=self.context_window,
                utilization_target=self.target_utilization,
                reasoning=f"Fixed batch size: {fixed_size}",
            )

        # Calculate actual token counts for these reviews
        review_tokens = [self.estimate_tokens(r.get("text", "")) for r in reviews]
        avg_review_tokens = sum(review_tokens) // len(review_tokens) if review_tokens else 150
        max_review_tokens = max(review_tokens) if review_tokens else 300

        # Use learned average if we have history, otherwise use current batch
        effective_avg = (
            (self.stats.avg_review_tokens + avg_review_tokens) // 2
            if self.stats.total_reviews > 0
            else avg_review_tokens
        )

        # Use learned output average
        output_per_review = self.stats.avg_output_tokens

        # Calculate available space
        available = int(self.context_window * self.target_utilization)
        available -= self.system_prompt_tokens
        available -= 1000  # Safety buffer for JSON overhead

        # Calculate batch size
        # Each review needs: input tokens + output tokens
        tokens_per_review = effective_avg + output_per_review

        # Use 80th percentile estimate to handle variance
        # (avg + (max - avg) * 0.3) gives room for longer reviews
        variance_adjusted = effective_avg + int((max_review_tokens - effective_avg) * 0.3)
        tokens_per_review_safe = variance_adjusted + output_per_review

        batch_size = max(1, available // tokens_per_review_safe)

        # Cap at reasonable limits
        batch_size = min(batch_size, 100, len(reviews))

        reasoning = (
            f"Context: {self.context_window:,} | "
            f"System: {self.system_prompt_tokens:,} | "
            f"Avg review: {effective_avg} (variance-adjusted: {variance_adjusted}) | "
            f"Output/review: {output_per_review} | "
            f"Target utilization: {self.target_utilization:.0%} | "
            f"→ Batch size: {batch_size}"
        )

        return BatchSizeCalculation(
            batch_size=batch_size,
            system_prompt_tokens=self.system_prompt_tokens,
            avg_tokens_per_review=effective_avg,
            output_tokens_reserved=batch_size * output_per_review,
            context_window=self.context_window,
            utilization_target=self.target_utilization,
            reasoning=reasoning,
        )

    def update_from_response(self, batch_size: int, input_tokens: int, output_tokens: int):
        """
        Update statistics from actual LLM response.

        Call this after each batch to improve future estimates.
        """
        if batch_size > 0:
            avg_input = input_tokens // batch_size
            avg_output = output_tokens // batch_size

            # Update stats for each review in batch
            for _ in range(batch_size):
                self.stats.update(avg_input, avg_output)

            logger.debug(
                f"BatchSizer updated: {batch_size} reviews, "
                f"avg input={avg_input}, avg output={avg_output}, "
                f"running avg input={self.stats.avg_review_tokens}, "
                f"running avg output={self.stats.avg_output_tokens}"
            )

    def get_stats_summary(self) -> dict:
        """Get current statistics summary."""
        return {
            "total_reviews_processed": self.stats.total_reviews,
            "avg_input_tokens": self.stats.avg_review_tokens,
            "avg_output_tokens": self.stats.avg_output_tokens,
            "min_review_tokens": self.stats.min_review_tokens if self.stats.total_reviews > 0 else 0,
            "max_review_tokens": self.stats.max_review_tokens if self.stats.total_reviews > 0 else 0,
            "model": self.model,
            "context_window": self.context_window,
        }

# System prompt for URT classification
SYSTEM_PROMPT = """You are a review classification system using URT (Universal Review Taxonomy) v5.1.

Your task is to extract semantic spans from customer reviews and classify each span independently.

## SPAN EXTRACTION RULES

**CRITICAL: Use TOPIC-BASED splitting, NOT sentence-based splitting.**

A span = all consecutive text about the SAME topic/domain, regardless of sentence count.

### When to KEEP TOGETHER (same span):
- Multiple sentences about the same topic: "The food was great. I loved the pasta. The sauce was perfect." → ONE span (all about Offering)
- Cause and effect: "The wait was long because they were understaffed" → ONE span
- Elaboration: "Staff was rude. They ignored us for 20 minutes." → ONE span (both about People)
- Single-topic reviews: Even if 5 sentences, if all about food → ONE span

### When to SPLIT (separate spans):
- Contrasting conjunctions that change topic: "Food was great BUT service was slow" → TWO spans
- Domain change: food (O) → staff (P) → ambiance (E) = split at each change
- Target change: "The waiter was nice but the manager was rude" → TWO spans (different people)

### Examples:
- "Amazing food. Best burger ever. Fries were crispy too." → 1 span (all Offering, V+)
- "Food was great but we waited an hour." → 2 spans (Offering V+, Journey V-)
- "I've been coming here for years. Always consistent quality." → 1 span (Relationship)
- "The staff are lovely and amazing with kids. More highchairs are definitely needed though." → 2 spans (People V+, Access V-)

**Guardrails**:
- Prefer FEWER, LARGER spans over many small ones
- Most reviews should have 1-3 spans, rarely more
- Min 1 span per review
- Spans must be non-overlapping

## URT TAXONOMY - COMPLETE (138 codes, use EXACT codes)

### O - OFFERING (Product/Service Quality) - 18 codes
O1.01 Works/Doesn't Work: Basic functionality success or failure
O1.02 Performance Level: How well it operates
O1.03 Durability: Longevity and resistance to wear
O1.04 Reliability: Consistency of function over time
O1.05 Outcome Achievement: Did customer accomplish their goal?
O2.01 Materials/Inputs: Quality of components or ingredients
O2.02 Craftsmanship: Skill of construction or execution
O2.03 Presentation: Visual and aesthetic quality
O2.04 Attention to Detail: Finishing touches and refinement
O2.05 Condition at Delivery: State when received
O3.01 All Components Present: Nothing missing from what was promised
O3.02 Feature Availability: Promised features actually work
O3.03 Scope Delivery: Full scope of work completed
O3.04 Documentation: Supporting materials provided
O4.01 Specification Match: Matches what was ordered
O4.02 Personalization: Adapted to individual preferences
O4.03 Flexibility: Can be modified or adjusted
O4.04 Appropriateness: Right solution for the need

### P - PEOPLE (Staff Interactions) - 20 codes
P1.01 Warmth: Friendly and welcoming manner
P1.02 Respect: Treated with dignity
P1.03 Patience: Calm and tolerant approach
P1.04 Enthusiasm: Energy and engagement
P1.05 Empathy: Understanding feelings
P2.01 Knowledge: Expertise and understanding
P2.02 Skill: Technical ability
P2.03 Problem Solving: Ability to find solutions
P2.04 Advice Quality: Helpful recommendations
P2.05 Training Level: Staff training evident
P3.01 Attentiveness: Being present and engaged
P3.02 Initiative: Proactive help
P3.03 Follow-through: Completing promised actions
P3.04 Availability: Being available when needed
P3.05 Dedication: Commitment to helping
P4.01 Clarity: Clear communication
P4.02 Listening: Understanding customer needs
P4.03 Transparency: Honest and open
P4.04 Honesty: Truthful communication
P4.05 Proactive Updates: Keeping customer informed

### J - JOURNEY (Process & Timing) - 20 codes
J1.01 Speed: How fast things happen
J1.02 Punctuality: On-time delivery
J1.03 Queue Management: Handling of waiting customers
J1.04 Punctuality: Meeting scheduled times
J1.05 Pacing: Appropriate speed (not rushed/dragged)
J2.01 Simplicity: Easy process
J2.02 Friction: Obstacles encountered
J2.03 Navigation: Finding what you need
J2.04 Booking Availability: Slots/capacity when needed
J2.05 Inventory: Stock availability
J3.01 Consistency: Same experience every time
J3.02 Accuracy: Getting it right
J3.03 Uptime: System availability
J3.04 Data Accuracy: Correct info in systems
J3.05 Integration: Systems work together
J4.01 Problem Recognition: Acknowledging issues
J4.02 Resolution Speed: How fast problems get fixed
J4.03 Resolution Fairness: Fair handling of issues
J4.04 Escalation: Getting to right person
J4.05 Closure: Issue fully resolved

### E - ENVIRONMENT (Physical & Digital Space) - 20 codes
E1.01 Cleanliness: How clean the space is
E1.02 Comfort: Physical comfort
E1.03 Space Design: Layout and organization
E1.04 Ambiance: Atmosphere and vibe
E1.05 Comfort: Physical comfort
E2.01 Lighting: Light quality and level
E2.02 Sound/Noise: Audio environment
E2.03 Temperature: Climate control
E2.04 Visual Design: Aesthetics of interface
E2.05 Mobile Experience: Mobile usability
E3.01 Interface Design: Digital UX/UI
E3.02 App/Website Speed: Digital performance
E3.03 Usability: Ease of digital use
E3.04 Health Safety: Health precautions
E3.05 Cyber Security: Digital security
E4.01 Safety: Physical safety
E4.02 Security: Protection of belongings/data
E4.03 Health/Hygiene: Health standards
E4.04 Social Responsibility: Ethical practices
E4.05 Community Impact: Local community effect

### A - ACCESS (Availability & Accessibility) - 20 codes
A1.01 Hours: Operating hours
A1.02 Booking Availability: Appointment slots
A1.03 Inventory: Product availability
A1.04 Wayfinding: Finding destination
A1.05 Physical Accessibility: Disability accommodations
A2.01 Physical Access: Mobility accessibility
A2.02 Language Access: Language accommodation
A2.03 Digital Accessibility: Screen reader/a11y
A2.04 Language Accessibility: Multilingual support
A2.05 Hours of Operation: Service availability times
A3.01 Diversity Welcome: All backgrounds welcome
A3.02 Accommodation: Special needs accommodation
A3.03 Response Time: Speed of getting answers
A3.04 Documentation Clarity: Clear instructions
A3.05 Support Accessibility: Getting help when needed
A4.01 Location: Physical location convenience
A4.02 Parking: Parking availability
A4.03 Multiple Channels: Ways to engage
A4.04 Payment Flexibility: Multiple payment options
A4.05 Refund Accessibility: Getting money back

### V - VALUE (Pricing & Costs) - 20 codes ⚠️ USE FOR ALL PRICE/COST/FEE MENTIONS
V1.01 Price Level: Cost amount ("cheap", "expensive", "affordable", "€", "$")
V1.02 Price Fairness: Fair for what you get
V1.03 Hidden Costs: Unexpected charges, surprise fees, hidden fees, extra charges
V1.04 Price Transparency: Clear pricing upfront
V1.05 Price Stability: Consistent pricing
V2.01 Clear Pricing: Easy to understand costs
V2.02 Honest Billing: Accurate charges
V2.03 Policy Clarity: Clear terms and conditions
V2.04 Quality-Price Ratio: Worth vs cost
V2.05 Competitive Value: Compared to alternatives
V3.01 Time Investment: Time required
V3.02 Hassle Factor: Difficulty and inconvenience
V3.03 Mental Load: Cognitive effort required
V3.04 Promotion Clarity: Clear offer terms
V3.05 Reward Redemption: Using points/rewards
V4.01 Value for Money: Worth what you paid
V4.02 ROI: Return on investment
V4.03 Overall Satisfaction: Happy with the exchange
V4.04 Billing Accuracy: Correct charges
V4.05 Billing Resolution: Fixing billing issues

### R - RELATIONSHIP (Trust & Loyalty) - 20 codes
R1.01 Honesty: Truthfulness
R1.02 Ethics: Ethical behavior, deceptive practices, scams
R1.03 Promises Kept: Following through on promises
R1.04 Ethics: Ethical behavior
R1.05 Accountability: Taking responsibility
R2.01 Consistency: Reliable over time
R2.02 Trustworthiness: Can be trusted
R2.03 Accountability: Takes responsibility
R2.04 Predictability: Consistent experience
R2.05 Standards: Meeting quality standards
R3.01 Error Acknowledgment: Admits mistakes
R3.02 Apology Quality: Sincere apologies
R3.03 Making It Right: Correcting mistakes
R3.04 Personal Connection: Human touch
R3.05 Going Extra Mile: Beyond expectations
R4.01 Customer Recognition: Remembers customers
R4.02 Loyalty Rewards: Rewards for loyalty
R4.03 Long-term Relationship: Builds relationships
R4.04 Service Recovery: Making things right
R4.05 Feedback Response: Acting on feedback

## CLASSIFICATION EXAMPLES (Critical Distinctions)

**PRICING/COSTS → V codes (Value), NOT P codes:**
- "Cheap prices", "good price", "€50" → V1.01 Price Level
- "Hidden charges", "surprise fees", "extra €35" → V1.03 Hidden Costs
- "Great value for money" → V4.01 Value for Money
- "Overcharged", "wrong amount" → V4.04 Billing Accuracy

**STAFF BEHAVIOR → P codes (People):**
- "Staff was friendly", "welcoming" → P1.01 Warmth
- "Rude", "disrespectful", "ignored us" → P1.02 Respect
- "Patient", "took their time" → P1.03 Patience
- "Knowledgeable", "expert" → P2.01 Knowledge

**DECEPTION/ETHICS → R codes (Relationship):**
- "They lied", "misleading" → R1.01 Honesty
- "Felt scammed", "dishonest practices" → R1.02 Ethics
- "Didn't honor the deal" → R1.03 Promises Kept

## DIMENSION CODES

### Valence
- V+ : Positive sentiment
- V- : Negative sentiment
- V0 : Neutral/factual
- V± : Mixed within the span

### Intensity
- I1 : Low ("okay", "fine", "decent")
- I2 : Moderate ("good", "bad", "slow")
- I3 : High ("amazing", "terrible", "unacceptable")

### Specificity
- S1 : Vague ("it was bad")
- S2 : Some detail ("the food was cold")
- S3 : Precise ("waited 45 minutes for appetizers")

### Actionability
- A1 : No clear action possible
- A2 : Possible actions, unclear which
- A3 : Clear, specific action ("train staff on X", "fix Y")

### Temporal
- TC : Current visit (default when no markers)
- TR : Recent pattern ("lately", "recently", "again")
- TH : Historical ("for years", "always", "used to")
- TF : Future ("won't return", "next time", "I expect")

### Evidence
- ES : Stated explicitly in text (default)
- EI : Inferred logically (not stated, but entailed)
- EC : Contextual (depends on surrounding text)

### Comparative
- CR-N : No comparison (default)
- CR-B : Better than alternatives
- CR-W : Worse than alternatives
- CR-S : Same as alternatives

## PRIMARY SPAN SELECTION

Mark exactly ONE span as is_primary=true using this order:
1. Highest intensity (I3 > I2 > I1)
2. Tie-break: negative over positive (V- > V± > V0 > V+)
3. Tie-break: earliest span_index

## USN (URT String Notation)

Generate a USN string for each span:
```
URT:S:{primary}[+{sec1}][+{sec2}]:{valence_sign}{intensity_num}:{S#}{A#}{temporal}.{evidence}.{CR_suffix}
```

Examples:
- `URT:S:J1.03:-2:22TC.ES.N` (J1.03, V-, I2, S2, A2, TC, ES, CR-N)
- `URT:S:P1.01+O2.03:+3:33TR.ES.B` (P1.01 primary, O2.03 secondary, V+, I3, S3, A3, TR, ES, CR-B)

Valence encoding: + for V+, - for V-, 0 for V0, ± for V±
CR suffix: N=CR-N, B=CR-B, W=CR-W, S=CR-S

## OUTPUT FORMAT

Return valid JSON matching this schema. No markdown, no explanations.

{
  "spans": [
    {
      "span_index": 0,
      "span_text": "exact text from review",
      "span_start": 0,
      "span_end": 25,
      "urt_primary": "O1.01",
      "urt_secondary": [],
      "valence": "V+",
      "intensity": "I2",
      "specificity": "S2",
      "actionability": "A1",
      "temporal": "TC",
      "evidence": "ES",
      "comparative": "CR-N",
      "is_primary": true,
      "confidence": "high",
      "entity": null,
      "entity_type": null,
      "relation_type": null,
      "related_span_index": null,
      "usn": "URT:S:O1.01:+2:21TC.ES.N"
    }
  ],
  "review_summary": {
    "dominant_valence": "V+",
    "dominant_domain": "O",
    "span_count": 1,
    "has_comparative": false,
    "has_entity": false
  }
}"""


class BatchReviewInput(TypedDict):
    """Input format for batch classification."""
    review_id: str
    text: str
    rating: int


class BatchClassificationResponse(TypedDict):
    """Response format for batch classification."""
    reviews: list[dict[str, Any]]  # Each contains review_index, spans, review_summary


class LLMClientBase(ABC):
    """Abstract base class for LLM clients."""

    def __init__(self, config: Config):
        self.config = config
        self.total_tokens_used = 0
        self.total_cost_usd = 0.0
        self._custom_prompt: str | None = None
        self._custom_prompt_batch: str | None = None
        self._cached_tokens: int = 0  # Track cached token usage

    def set_prompt(self, prompt: str, batch_prompt: str | None = None) -> None:
        """
        Set custom system prompts (e.g., built dynamically from database).

        Args:
            prompt: The system prompt for single review classification
            batch_prompt: The system prompt for batch classification (if different)
        """
        self._custom_prompt = prompt
        self._custom_prompt_batch = batch_prompt or prompt

    def get_prompt(self, batch_mode: bool = False) -> str:
        """Get the current system prompt (custom or default)."""
        if batch_mode:
            return self._custom_prompt_batch or self._custom_prompt or SYSTEM_PROMPT
        return self._custom_prompt or SYSTEM_PROMPT

    @abstractmethod
    async def classify(
        self,
        review_text: str,
        profile: str = "standard",
    ) -> tuple[LLMClassificationResponse, dict[str, Any]]:
        """
        Classify a single review and extract spans.

        Args:
            review_text: The review text to classify
            profile: Classification profile (lite/core/standard/full)

        Returns:
            Tuple of (classification response, metadata dict with tokens/cost)
        """
        pass

    @abstractmethod
    async def classify_batch(
        self,
        reviews: list[BatchReviewInput],
        profile: str = "standard",
    ) -> tuple[list[LLMClassificationResponse], dict[str, Any]]:
        """
        Classify multiple reviews in a single LLM call.

        Args:
            reviews: List of reviews with review_id, text, and rating
            profile: Classification profile (lite/core/standard/full)

        Returns:
            Tuple of (list of classification responses, aggregated metadata)
        """
        pass

    @abstractmethod
    async def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.7,
        max_tokens: int = 4000,
    ) -> str:
        """
        Generate text using the LLM (for synthesis, narratives, etc.).

        Args:
            system_prompt: System instructions
            user_prompt: User content/context
            temperature: Creativity level (0-1)
            max_tokens: Maximum response length

        Returns:
            Generated text response
        """
        pass

    @abstractmethod
    async def close(self) -> None:
        """Close the client and cleanup resources."""
        pass

    def _build_batch_user_prompt(self, reviews: list[BatchReviewInput]) -> str:
        """Build user prompt for batch classification."""
        lines = [
            f"Classify these {len(reviews)} reviews. Return JSON with 'reviews' array.",
            ""
        ]

        for i, review in enumerate(reviews):
            lines.append(f"---REVIEW {i} (rating={review['rating']}★)---")
            lines.append(review["text"])
            lines.append("")

        return "\n".join(lines)


class OpenAIClient(LLMClientBase):
    """OpenAI LLM client implementation with batch support and prompt caching."""

    # Pricing per 1M tokens (as of 2024)
    # Cached input tokens are 50% cheaper
    PRICING = {
        "gpt-4o": {"input": 2.50, "cached_input": 1.25, "output": 10.0},
        "gpt-4o-mini": {"input": 0.15, "cached_input": 0.075, "output": 0.60},
        "gpt-4-turbo": {"input": 10.0, "cached_input": 5.0, "output": 30.0},
        "gpt-3.5-turbo": {"input": 0.50, "cached_input": 0.25, "output": 1.50},
    }

    def __init__(self, config: Config):
        super().__init__(config)
        from openai import AsyncOpenAI

        self.client = AsyncOpenAI(api_key=config.get_llm_api_key())
        self.model = config.llm_model

    async def classify(
        self,
        review_text: str,
        profile: str = "standard",
    ) -> tuple[LLMClassificationResponse, dict[str, Any]]:
        """Classify a single review using OpenAI."""
        start_time = time.time()

        messages = [
            {"role": "system", "content": self.get_prompt()},
            {
                "role": "user",
                "content": f'Classify this review:\n\n"{review_text}"',
            },
        ]

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=self.config.llm_temperature,
            response_format={"type": "json_object"},
            timeout=self.config.llm_timeout_seconds,
        )

        # Parse response
        content = response.choices[0].message.content
        if not content:
            raise ValueError("Empty response from OpenAI")

        result = json.loads(content)

        # Calculate costs (with caching support)
        metadata = self._calculate_openai_costs(response, start_time)

        return result, metadata

    async def classify_batch(
        self,
        reviews: list[BatchReviewInput],
        profile: str = "standard",
    ) -> tuple[list[LLMClassificationResponse], dict[str, Any]]:
        """
        Classify multiple reviews in a single LLM call.

        Uses prompt caching - the system prompt is cached after first call,
        reducing input token costs by ~50% on subsequent calls.
        """
        if not reviews:
            return [], {"error": "No reviews provided"}

        start_time = time.time()

        # Build batch user prompt
        user_prompt = self._build_batch_user_prompt(reviews)

        messages = [
            {"role": "system", "content": self.get_prompt(batch_mode=True)},
            {"role": "user", "content": user_prompt},
        ]

        # Calculate max tokens based on batch size (estimate ~400 tokens per review output)
        max_output_tokens = min(16000, len(reviews) * 500)

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=self.config.llm_temperature,
            response_format={"type": "json_object"},
            max_tokens=max_output_tokens,
            timeout=self.config.llm_timeout_seconds * 2,  # Longer timeout for batch
        )

        # Parse response
        content = response.choices[0].message.content
        if not content:
            raise ValueError("Empty response from OpenAI")

        metadata = self._calculate_openai_costs(response, start_time, batch_size=len(reviews))

        # Try full JSON parse first
        try:
            batch_result = json.loads(content)
            review_results = self._parse_batch_response(batch_result, reviews)
            return review_results, metadata
        except json.JSONDecodeError as e:
            # Full parse failed - try partial extraction
            logger.warning(f"Full JSON parse failed: {e}, attempting partial recovery...")

            partial_reviews, missing_indices = self._extract_partial_batch_json(
                content, len(reviews)
            )

            if partial_reviews:
                raise PartialBatchResult(
                    message=f"Recovered {len(partial_reviews)}/{len(reviews)} reviews from malformed JSON",
                    partial_results=partial_reviews,
                    missing_indices=missing_indices,
                    metadata=metadata,
                )
            else:
                raise

    def _calculate_openai_costs(
        self,
        response: Any,
        start_time: float,
        batch_size: int = 1,
    ) -> dict[str, Any]:
        """Calculate costs from OpenAI response, accounting for cached tokens."""
        input_tokens = response.usage.prompt_tokens if response.usage else 0
        output_tokens = response.usage.completion_tokens if response.usage else 0
        total_tokens = input_tokens + output_tokens

        # Check for cached tokens (OpenAI returns this in newer API versions)
        cached_tokens = 0
        if hasattr(response.usage, "prompt_tokens_details") and response.usage.prompt_tokens_details:
            cached_tokens = getattr(response.usage.prompt_tokens_details, "cached_tokens", 0)

        uncached_input = input_tokens - cached_tokens

        pricing = self.PRICING.get(self.model, {"input": 0.15, "cached_input": 0.075, "output": 0.60})
        cost = (
            uncached_input * pricing["input"]
            + cached_tokens * pricing.get("cached_input", pricing["input"] * 0.5)
            + output_tokens * pricing["output"]
        ) / 1_000_000

        self.total_tokens_used += total_tokens
        self.total_cost_usd += cost
        self._cached_tokens += cached_tokens

        return {
            "model": self.model,
            "input_tokens": input_tokens,
            "cached_tokens": cached_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
            "cost_usd": cost,
            "latency_ms": int((time.time() - start_time) * 1000),
            "batch_size": batch_size,
            "tokens_per_review": total_tokens / batch_size if batch_size > 0 else 0,
        }

    def _parse_batch_response(
        self,
        batch_result: dict[str, Any],
        original_reviews: list[BatchReviewInput],
    ) -> list[LLMClassificationResponse]:
        """Parse batch response into individual review results."""
        results: list[LLMClassificationResponse] = []

        # Handle both formats: {"reviews": [...]} and direct list
        review_data = batch_result.get("reviews", [])
        if not review_data and isinstance(batch_result, list):
            review_data = batch_result

        # Create a lookup by review_index
        results_by_index = {r.get("review_index", i): r for i, r in enumerate(review_data)}

        for i, original in enumerate(original_reviews):
            if i in results_by_index:
                review_result = results_by_index[i]
                # Convert to standard format
                results.append({
                    "spans": review_result.get("spans", []),
                    "review_summary": review_result.get("review_summary", {
                        "dominant_valence": "V0",
                        "dominant_domain": "O",
                        "span_count": len(review_result.get("spans", [])),
                        "has_comparative": False,
                        "has_entity": False,
                    }),
                })
            else:
                # Missing review - create fallback
                logger.warning(f"Review index {i} missing from batch response, using fallback")
                results.append(create_fallback_response(original["text"]))

        return results

    async def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.7,
        max_tokens: int = 4000,
    ) -> str:
        """Generate text using OpenAI."""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            response_format={"type": "json_object"},
            timeout=self.config.llm_timeout_seconds,
        )

        content = response.choices[0].message.content
        if not content:
            raise ValueError("Empty response from OpenAI")

        # Track usage
        if response.usage:
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
            pricing = self.PRICING.get(self.model, {"input": 0.15, "output": 0.60})
            cost = (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000
            self.total_tokens_used += input_tokens + output_tokens
            self.total_cost_usd += cost

        return content

    async def close(self) -> None:
        """Close the OpenAI client."""
        await self.client.close()


class AnthropicClient(LLMClientBase):
    """Anthropic LLM client implementation with batch support and prompt caching."""

    # Pricing per 1M tokens (as of 2024)
    # Cached input tokens are 90% cheaper with Anthropic
    PRICING = {
        "claude-3-opus-20240229": {"input": 15.0, "cached_input": 1.50, "output": 75.0},
        "claude-3-sonnet-20240229": {"input": 3.0, "cached_input": 0.30, "output": 15.0},
        "claude-3-haiku-20240307": {"input": 0.25, "cached_input": 0.03, "output": 1.25},
        "claude-3-5-sonnet-20241022": {"input": 3.0, "cached_input": 0.30, "output": 15.0},
        "claude-sonnet-4-20250514": {"input": 3.0, "cached_input": 0.30, "output": 15.0},
    }

    def __init__(self, config: Config):
        super().__init__(config)
        from anthropic import AsyncAnthropic

        self.client = AsyncAnthropic(api_key=config.get_llm_api_key())
        self.model = config.llm_model

    async def classify(
        self,
        review_text: str,
        profile: str = "standard",
    ) -> tuple[LLMClassificationResponse, dict[str, Any]]:
        """Classify a single review using Anthropic."""
        start_time = time.time()

        # Use cache_control for prompt caching
        system_content = self._build_cached_system(self.get_prompt())

        response = await self.client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=system_content,
            messages=[
                {
                    "role": "user",
                    "content": f'Classify this review and return JSON only:\n\n"{review_text}"',
                },
            ],
            temperature=self.config.llm_temperature,
        )

        # Parse response
        content = response.content[0].text if response.content else ""
        if not content:
            raise ValueError("Empty response from Anthropic")

        result = self._extract_json(content)
        metadata = self._calculate_anthropic_costs(response, start_time)

        return result, metadata

    async def classify_batch(
        self,
        reviews: list[BatchReviewInput],
        profile: str = "standard",
    ) -> tuple[list[LLMClassificationResponse], dict[str, Any]]:
        """
        Classify multiple reviews in a single LLM call.

        Uses Anthropic's prompt caching with cache_control - the system prompt
        is cached after first call, reducing input costs by ~90%.
        """
        if not reviews:
            return [], {"error": "No reviews provided"}

        start_time = time.time()

        # Build batch user prompt
        user_prompt = self._build_batch_user_prompt(reviews)

        # Use cache_control for prompt caching (system prompt is cacheable)
        system_content = self._build_cached_system(self.get_prompt(batch_mode=True))

        # Calculate max tokens based on batch size
        max_output_tokens = min(16000, len(reviews) * 500)

        response = await self.client.messages.create(
            model=self.model,
            max_tokens=max_output_tokens,
            system=system_content,
            messages=[{"role": "user", "content": user_prompt}],
            temperature=self.config.llm_temperature,
        )

        # Parse response
        content = response.content[0].text if response.content else ""
        if not content:
            raise ValueError("Empty response from Anthropic")

        metadata = self._calculate_anthropic_costs(response, start_time, batch_size=len(reviews))

        # Try full JSON extraction first
        try:
            batch_result = self._extract_json(content)
            review_results = self._parse_batch_response(batch_result, reviews)
            return review_results, metadata
        except (json.JSONDecodeError, ValueError) as e:
            # Full parse failed - try partial extraction
            logger.warning(f"Full JSON parse failed: {e}, attempting partial recovery...")

            partial_reviews, missing_indices = self._extract_partial_batch_json(
                content, len(reviews)
            )

            if partial_reviews:
                # We recovered some results - raise PartialBatchResult
                raise PartialBatchResult(
                    message=f"Recovered {len(partial_reviews)}/{len(reviews)} reviews from malformed JSON",
                    partial_results=partial_reviews,
                    missing_indices=missing_indices,
                    metadata=metadata,
                )
            else:
                # Couldn't recover anything - re-raise original error
                raise

    def _build_cached_system(self, prompt: str) -> list[dict[str, Any]]:
        """Build system content with cache_control for prompt caching."""
        return [
            {
                "type": "text",
                "text": prompt,
                "cache_control": {"type": "ephemeral"},
            }
        ]

    def _calculate_anthropic_costs(
        self,
        response: Any,
        start_time: float,
        batch_size: int = 1,
    ) -> dict[str, Any]:
        """Calculate costs from Anthropic response, accounting for cached tokens."""
        input_tokens = response.usage.input_tokens
        output_tokens = response.usage.output_tokens
        total_tokens = input_tokens + output_tokens

        # Anthropic returns cache info in usage
        cached_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0
        cache_creation_tokens = getattr(response.usage, "cache_creation_input_tokens", 0) or 0

        uncached_input = input_tokens - cached_tokens

        pricing = self.PRICING.get(self.model, {"input": 3.0, "cached_input": 0.30, "output": 15.0})
        cost = (
            uncached_input * pricing["input"]
            + cached_tokens * pricing.get("cached_input", pricing["input"] * 0.1)
            + output_tokens * pricing["output"]
        ) / 1_000_000

        self.total_tokens_used += total_tokens
        self.total_cost_usd += cost
        self._cached_tokens += cached_tokens

        return {
            "model": self.model,
            "input_tokens": input_tokens,
            "cached_tokens": cached_tokens,
            "cache_creation_tokens": cache_creation_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
            "cost_usd": cost,
            "latency_ms": int((time.time() - start_time) * 1000),
            "batch_size": batch_size,
            "tokens_per_review": total_tokens / batch_size if batch_size > 0 else 0,
        }

    def _parse_batch_response(
        self,
        batch_result: dict[str, Any],
        original_reviews: list[BatchReviewInput],
    ) -> list[LLMClassificationResponse]:
        """Parse batch response into individual review results."""
        results: list[LLMClassificationResponse] = []

        # Handle both formats: {"reviews": [...]} and direct list
        review_data = batch_result.get("reviews", [])
        if not review_data and isinstance(batch_result, list):
            review_data = batch_result

        # Create a lookup by review_index
        results_by_index = {r.get("review_index", i): r for i, r in enumerate(review_data)}

        for i, original in enumerate(original_reviews):
            if i in results_by_index:
                review_result = results_by_index[i]
                results.append({
                    "spans": review_result.get("spans", []),
                    "review_summary": review_result.get("review_summary", {
                        "dominant_valence": "V0",
                        "dominant_domain": "O",
                        "span_count": len(review_result.get("spans", [])),
                        "has_comparative": False,
                        "has_entity": False,
                    }),
                })
            else:
                logger.warning(f"Review index {i} missing from batch response, using fallback")
                results.append(create_fallback_response(original["text"]))

        return results

    async def generate(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.7,
        max_tokens: int = 4000,
    ) -> str:
        """Generate text using Anthropic."""
        response = await self.client.messages.create(
            model=self.model,
            max_tokens=max_tokens,
            system=system_prompt,
            messages=[{"role": "user", "content": user_prompt}],
            temperature=temperature,
        )

        content = response.content[0].text if response.content else ""
        if not content:
            raise ValueError("Empty response from Anthropic")

        # Track usage
        input_tokens = response.usage.input_tokens
        output_tokens = response.usage.output_tokens
        pricing = self.PRICING.get(self.model, {"input": 3.0, "output": 15.0})
        cost = (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000
        self.total_tokens_used += input_tokens + output_tokens
        self.total_cost_usd += cost

        return self._extract_json_string(content)

    def _extract_json_string(self, content: str) -> str:
        """Extract JSON string from response, handling markdown code blocks."""
        import re
        content = content.strip()

        if content.startswith("{"):
            return content

        json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", content)
        if json_match:
            return json_match.group(1)

        json_match = re.search(r"\{[\s\S]*\}", content)
        if json_match:
            return json_match.group(0)

        return content

    def _extract_json(self, content: str) -> dict[str, Any]:
        """Extract JSON from response, handling markdown code blocks."""
        content = content.strip()

        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        import re

        json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", content)
        if json_match:
            return json.loads(json_match.group(1))

        json_match = re.search(r"\{[\s\S]*\}", content)
        if json_match:
            return json.loads(json_match.group(0))

        raise ValueError(f"Could not extract JSON from response: {content[:200]}")

    def _extract_partial_batch_json(
        self, content: str, expected_count: int
    ) -> tuple[list[dict[str, Any]], list[int]]:
        """
        Extract partial results from truncated/malformed batch JSON.

        Returns:
            Tuple of (successfully_parsed_reviews, missing_indices)
        """
        import re

        parsed_reviews: list[dict[str, Any]] = []
        found_indices: set[int] = set()

        # Pattern to match complete review objects with review_index
        # Matches: {"review_index": N, ... } with balanced braces
        review_pattern = r'\{\s*"review_index"\s*:\s*(\d+)[^{}]*(?:\{[^{}]*\}[^{}]*)*\}'

        # Try to find all complete review objects
        for match in re.finditer(review_pattern, content):
            try:
                # Extract the matched text and try to parse
                obj_text = match.group(0)

                # Try to parse as JSON - may need to fix trailing issues
                try:
                    obj = json.loads(obj_text)
                except json.JSONDecodeError:
                    # Try adding closing brace if truncated
                    continue

                if "review_index" in obj and "spans" in obj:
                    idx = obj["review_index"]
                    if idx not in found_indices:
                        parsed_reviews.append(obj)
                        found_indices.add(idx)
            except Exception:
                continue

        # Alternative: try parsing incrementally by finding review_index markers
        if len(parsed_reviews) < expected_count // 2:
            # Find all review_index positions
            index_matches = list(re.finditer(r'"review_index"\s*:\s*(\d+)', content))

            for i, match in enumerate(index_matches):
                idx = int(match.group(1))
                if idx in found_indices:
                    continue

                # Find the start of this review object
                start = content.rfind('{', 0, match.start())
                if start == -1:
                    continue

                # Find the end - either next review_index or end of content
                if i + 1 < len(index_matches):
                    end_search = index_matches[i + 1].start()
                else:
                    end_search = len(content)

                # Find the closing brace
                obj_text = content[start:end_search]

                # Count braces to find proper end
                brace_count = 0
                end_pos = 0
                for j, char in enumerate(obj_text):
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            end_pos = j + 1
                            break

                if end_pos > 0:
                    try:
                        obj = json.loads(obj_text[:end_pos])
                        # Validate required fields and data integrity
                        if self._validate_recovered_review(obj):
                            review_idx = obj["review_index"]
                            if review_idx not in found_indices:
                                parsed_reviews.append(obj)
                                found_indices.add(review_idx)
                    except json.JSONDecodeError:
                        continue

        # Determine missing indices
        missing_indices = [i for i in range(expected_count) if i not in found_indices]

        logger.info(
            f"Partial JSON recovery: {len(parsed_reviews)}/{expected_count} reviews recovered, "
            f"{len(missing_indices)} missing"
        )

        return parsed_reviews, missing_indices

    def _validate_recovered_review(self, obj: dict[str, Any]) -> bool:
        """
        Validate a recovered review has all required fields with valid data.

        Returns True only if the review is complete and usable.
        Rejects:
        - Missing review_index or spans
        - Empty spans array
        - Spans missing required fields (text, urt_primary, valence, intensity)
        - Empty field values
        """
        # Check required top-level fields
        if "review_index" not in obj:
            return False
        if not isinstance(obj.get("review_index"), int):
            return False

        if "spans" not in obj:
            return False
        if not isinstance(obj["spans"], list):
            return False
        if len(obj["spans"]) == 0:
            # Empty spans = no useful classification data
            return False

        # Validate each span has required fields with non-empty values
        required_span_fields = ["text", "urt_primary", "valence", "intensity"]
        for span in obj["spans"]:
            if not isinstance(span, dict):
                return False
            for field in required_span_fields:
                if field not in span:
                    return False
                if not span[field]:  # Empty string or None
                    return False

        # review_summary is optional but if present should be a dict
        if "review_summary" in obj and not isinstance(obj["review_summary"], dict):
            return False

        return True

    async def close(self) -> None:
        """Close the Anthropic client."""
        await self.client.close()


class LLMClient:
    """
    Factory class for LLM clients.

    Usage:
        client = LLMClient.create(config)
        result, metadata = await client.classify(review_text)
        await client.close()
    """

    @staticmethod
    def create(config: Config) -> LLMClientBase:
        """
        Create an LLM client based on configuration.

        Args:
            config: Pipeline configuration

        Returns:
            LLM client instance (OpenAI or Anthropic)
        """
        if config.llm_provider == "openai":
            return OpenAIClient(config)
        elif config.llm_provider == "anthropic":
            return AnthropicClient(config)
        else:
            raise ValueError(f"Unsupported LLM provider: {config.llm_provider}")


def create_fallback_response(review_text: str) -> LLMClassificationResponse:
    """
    Create a fallback classification response when LLM fails.

    Args:
        review_text: Original review text

    Returns:
        Minimal valid classification response
    """
    return {
        "spans": [
            {
                "span_index": 0,
                "span_text": review_text,
                "span_start": 0,
                "span_end": len(review_text),
                "urt_primary": "O1.01",
                "urt_secondary": [],
                "valence": "V0",
                "intensity": "I1",
                "specificity": "S1",
                "actionability": "A1",
                "temporal": "TC",
                "evidence": "ES",
                "comparative": "CR-N",
                "is_primary": True,
                "confidence": "low",
                "entity": None,
                "entity_type": None,
                "relation_type": None,
                "related_span_index": None,
                "usn": "URT:S:O1.01:01:11TC.ES.N",
            }
        ],
        "review_summary": {
            "dominant_valence": "V0",
            "dominant_domain": "O",
            "span_count": 1,
            "has_comparative": False,
            "has_entity": False,
        },
    }
