"""
Stage 2: LLM Classification

Classify normalized reviews into URT codes with span-level extraction.

Responsibilities:
- Call LLM for span extraction and classification
- Generate embeddings
- Calculate trust scores
- Select primary span
- Write to reviews_enriched and review_spans tables
"""

from __future__ import annotations

import hashlib
import logging
import re
import uuid
from typing import TYPE_CHECKING, Any

from reviewiq_pipeline.contracts import (
    ClassifiedReview,
    ExtractedSpan,
    ReviewToClassify,
    Stage2Input,
    Stage2Output,
    Stage2Stats,
)
from reviewiq_pipeline.services.llm_client import LLMClient, create_fallback_response

if TYPE_CHECKING:
    from reviewiq_pipeline.config import Config
    from reviewiq_pipeline.contracts import LLMClassificationResponse, LLMSpanResponse
    from reviewiq_pipeline.db.connection import DatabasePool
    from reviewiq_pipeline.db.repositories import ReviewRepository, SpanRepository
    from reviewiq_pipeline.services.embeddings import EmbeddingService
    from reviewiq_pipeline.services.llm_client import LLMClientBase

logger = logging.getLogger(__name__)

# URT code validation pattern
URT_CODE_PATTERN = re.compile(r"^[OPJEAVR][1-4]\.[0-9]{2}$")

# Valence priority for primary span selection (lower = higher priority)
VALENCE_PRIORITY = {"V-": 0, "V±": 1, "V0": 2, "V+": 3}

# Intensity priority (lower = higher priority for I3)
INTENSITY_PRIORITY = {"I3": 0, "I2": 1, "I1": 2}


class Stage2Classifier:
    """
    Stage 2: Classify reviews using LLM and extract spans.

    This stage:
    1. Calls LLM to extract and classify spans
    2. Generates embeddings for each review
    3. Calculates trust scores
    4. Selects primary span
    5. Writes classification results to database
    """

    def __init__(
        self,
        config: Config,
        db: DatabasePool | None = None,
        review_repo: ReviewRepository | None = None,
        span_repo: SpanRepository | None = None,
        embedding_service: EmbeddingService | None = None,
    ):
        self.config = config
        self.db = db
        self.review_repo = review_repo
        self.span_repo = span_repo
        self.embedding_service = embedding_service
        self._llm_client: LLMClientBase | None = None

    async def _get_llm_client(self) -> LLMClientBase:
        """Get or create LLM client."""
        if self._llm_client is None:
            self._llm_client = LLMClient.create(self.config)
        return self._llm_client

    async def close(self) -> None:
        """Close resources."""
        if self._llm_client:
            await self._llm_client.close()
            self._llm_client = None

    async def process(self, input_data: Stage2Input) -> Stage2Output:
        """
        Process reviews through classification stage.

        Args:
            input_data: Stage 2 input with reviews and config

        Returns:
            Stage2Output with classified reviews and stats
        """
        batch_id = str(uuid.uuid4())[:8]
        logger.info(
            f"Stage 2: Classifying {len(input_data['reviews'])} reviews "
            f"(batch {batch_id})"
        )

        classified_reviews: list[ClassifiedReview] = []
        total_tokens = 0
        total_cost = 0.0
        total_spans = 0
        error_count = 0

        llm_client = await self._get_llm_client()

        for review in input_data["reviews"]:
            try:
                classified, metadata = await self._classify_review(
                    review,
                    input_data["config"]["profile"],
                    llm_client,
                    batch_id,
                )

                if classified:
                    classified_reviews.append(classified)
                    total_spans += len(classified.get("spans", []))
                    total_tokens += metadata.get("total_tokens", 0)
                    total_cost += metadata.get("cost_usd", 0.0)

                    # Persist to database if configured
                    if self.review_repo and self.span_repo:
                        await self._persist_classification(
                            classified,
                            review,
                            batch_id,
                            input_data["config"],
                        )

            except Exception as e:
                logger.error(
                    f"Error classifying review {review['review_id']}: {e}",
                    exc_info=True,
                )
                error_count += 1

        avg_spans = total_spans / len(classified_reviews) if classified_reviews else 0

        logger.info(
            f"Stage 2 complete: {len(classified_reviews)} classified, "
            f"{error_count} errors, {total_spans} spans total"
        )

        return Stage2Output(
            batch_id=batch_id,
            taxonomy_version=input_data["config"]["taxonomy_version"],
            model_version=self.config.llm_model,
            prompt_version="v1.0",
            reviews_classified=classified_reviews,
            stats=Stage2Stats(
                input_count=len(input_data["reviews"]),
                success_count=len(classified_reviews),
                error_count=error_count,
                total_spans=total_spans,
                avg_spans_per_review=avg_spans,
                llm_tokens_used=total_tokens,
                llm_cost_usd=total_cost,
            ),
        )

    async def _classify_review(
        self,
        review: ReviewToClassify,
        profile: str,
        llm_client: LLMClientBase,
        batch_id: str,
    ) -> tuple[ClassifiedReview | None, dict[str, Any]]:
        """
        Classify a single review.

        Args:
            review: Review to classify
            profile: Classification profile
            llm_client: LLM client instance
            batch_id: Batch identifier

        Returns:
            Tuple of (classified review, metadata)
        """
        metadata: dict[str, Any] = {}

        # Call LLM for classification
        try:
            llm_response, llm_metadata = await llm_client.classify(
                review["text"],
                profile,
            )
            metadata.update(llm_metadata)
        except Exception as e:
            logger.warning(
                f"LLM classification failed for {review['review_id']}, "
                f"using fallback: {e}"
            )
            llm_response = create_fallback_response(review["text"])
            metadata["fallback"] = True

        # Validate and fix response
        llm_response = self._validate_and_fix_response(llm_response, review["text"])

        # Convert spans to our format
        spans = self._convert_spans(
            llm_response["spans"],
            review,
            profile,
            batch_id,
        )

        # Ensure exactly one primary span
        spans = self._ensure_primary_span(spans)

        # Find the primary span for review-level classification
        primary_span = next((s for s in spans if s.get("is_primary")), spans[0] if spans else None)

        # Generate embedding
        embedding: list[float] = []
        if self.embedding_service:
            embedding = self.embedding_service.embed(review["text_normalized"])

        # Calculate trust score
        trust_score = self._calculate_trust_score(review, spans)

        # Extract staff mentions and quotes
        staff_mentions = self._extract_staff_mentions(spans)
        quotes = self._extract_quotes(spans)

        return ClassifiedReview(
            source=review["source"],
            review_id=review["review_id"],
            review_version=review["review_version"],
            urt_primary=primary_span["urt_primary"] if primary_span else "O1.01",
            urt_secondary=primary_span.get("urt_secondary", []) if primary_span else [],
            valence=primary_span["valence"] if primary_span else "V0",
            intensity=primary_span["intensity"] if primary_span else "I1",
            comparative=primary_span.get("comparative", "CR-N") if primary_span else "CR-N",
            staff_mentions=staff_mentions,
            quotes=quotes,
            trust_score=trust_score,
            embedding=embedding,
            spans=spans,
            classification_confidence={
                "overall": 0.8 if not metadata.get("fallback") else 0.3
            },
            processing_time_ms=metadata.get("latency_ms", 0),
        ), metadata

    def _validate_and_fix_response(
        self,
        response: LLMClassificationResponse,
        original_text: str,
    ) -> LLMClassificationResponse:
        """
        Validate LLM response and fix common issues.

        Args:
            response: Raw LLM response
            original_text: Original review text for offset validation

        Returns:
            Validated and fixed response
        """
        spans = response.get("spans", [])
        if not spans:
            # Create fallback if no spans
            return create_fallback_response(original_text)

        fixed_spans = []
        for i, span in enumerate(spans):
            # Ensure required fields
            span["span_index"] = i

            # Validate and fix offsets
            start = span.get("span_start", 0)
            end = span.get("span_end", len(original_text))

            if start < 0:
                start = 0
            if end > len(original_text):
                end = len(original_text)
            if end <= start:
                end = start + len(span.get("span_text", "")) or len(original_text)

            span["span_start"] = start
            span["span_end"] = end

            # Validate URT code
            urt_primary = span.get("urt_primary", "O1.01")
            if not URT_CODE_PATTERN.match(urt_primary):
                logger.warning(f"Invalid URT code '{urt_primary}', defaulting to O1.01")
                span["urt_primary"] = "O1.01"

            # Ensure valid enums
            if span.get("valence") not in ("V+", "V-", "V0", "V±"):
                span["valence"] = "V0"
            if span.get("intensity") not in ("I1", "I2", "I3"):
                span["intensity"] = "I1"

            fixed_spans.append(span)

        response["spans"] = fixed_spans
        return response

    def _convert_spans(
        self,
        llm_spans: list[LLMSpanResponse],
        review: ReviewToClassify,
        profile: str,
        batch_id: str,
    ) -> list[ExtractedSpan]:
        """
        Convert LLM spans to our ExtractedSpan format.

        Args:
            llm_spans: Spans from LLM response
            review: Source review
            profile: Classification profile
            batch_id: Batch identifier

        Returns:
            List of ExtractedSpan objects
        """
        spans = []

        for llm_span in llm_spans:
            # Generate deterministic span ID
            span_key = f"{review['review_id']}:{llm_span['span_index']}:{llm_span.get('span_text', '')[:50]}"
            span_hash = hashlib.sha256(span_key.encode()).hexdigest()[:16]
            span_id = f"SPN-{span_hash}"

            span = ExtractedSpan(
                span_id=span_id,
                span_index=llm_span["span_index"],
                span_text=llm_span.get("span_text", ""),
                span_start=llm_span.get("span_start", 0),
                span_end=llm_span.get("span_end", 0),
                profile=profile,  # type: ignore
                urt_primary=llm_span["urt_primary"],
                urt_secondary=llm_span.get("urt_secondary", []),
                valence=llm_span["valence"],
                intensity=llm_span["intensity"],
                comparative=llm_span.get("comparative", "CR-N"),
                specificity=llm_span.get("specificity"),
                actionability=llm_span.get("actionability"),
                temporal=llm_span.get("temporal"),
                evidence=llm_span.get("evidence"),
                entity=llm_span.get("entity"),
                entity_type=llm_span.get("entity_type"),
                entity_normalized=llm_span.get("entity", "").lower() if llm_span.get("entity") else None,
                relation_type=llm_span.get("relation_type"),
                related_span_index=llm_span.get("related_span_index"),
                confidence=llm_span.get("confidence", "medium"),
                usn=llm_span.get("usn", self._generate_usn(llm_span)),
                is_primary=llm_span.get("is_primary", False),
            )
            spans.append(span)

        return spans

    def _ensure_primary_span(self, spans: list[ExtractedSpan]) -> list[ExtractedSpan]:
        """
        Ensure exactly one span is marked as primary.

        Uses selection rules:
        1. Highest intensity (I3 > I2 > I1)
        2. Tie-break: negative over positive (V- > V± > V0 > V+)
        3. Tie-break: earliest span_index

        Args:
            spans: List of spans

        Returns:
            List of spans with exactly one primary
        """
        if not spans:
            return spans

        # Count current primaries
        primary_count = sum(1 for s in spans if s.get("is_primary"))

        if primary_count == 1:
            return spans

        # Clear all primaries and re-select
        for span in spans:
            span["is_primary"] = False

        # Sort by selection criteria
        def sort_key(s: ExtractedSpan) -> tuple[int, int, int]:
            return (
                INTENSITY_PRIORITY.get(s["intensity"], 2),
                VALENCE_PRIORITY.get(s["valence"], 3),
                s["span_index"],
            )

        sorted_spans = sorted(spans, key=sort_key)
        sorted_spans[0]["is_primary"] = True

        return spans

    def _calculate_trust_score(
        self,
        review: ReviewToClassify,
        spans: list[ExtractedSpan],
    ) -> float:
        """
        Calculate trust score for a review.

        Factors:
        - Text length (longer = more trust)
        - Specificity of spans
        - Confidence levels

        Args:
            review: Source review
            spans: Classified spans

        Returns:
            Trust score between 0.2 and 1.0
        """
        score = 0.5  # Base score

        # Length factor (up to +0.2)
        text_len = len(review["text"])
        if text_len > 200:
            score += 0.2
        elif text_len > 100:
            score += 0.1
        elif text_len > 50:
            score += 0.05

        # Specificity factor (up to +0.2)
        if spans:
            high_spec_count = sum(1 for s in spans if s.get("specificity") == "S3")
            if high_spec_count > 0:
                score += 0.1 + (0.1 * min(high_spec_count / len(spans), 1.0))

        # Confidence factor (up to +0.1)
        if spans:
            high_conf_count = sum(1 for s in spans if s.get("confidence") == "high")
            score += 0.1 * (high_conf_count / len(spans))

        # Ensure floor of 0.2 and ceiling of 1.0
        return max(self.config.trust_score_floor, min(1.0, score))

    def _extract_staff_mentions(self, spans: list[ExtractedSpan]) -> list[str]:
        """Extract staff names from spans."""
        staff = []
        for span in spans:
            if span.get("entity_type") == "staff" and span.get("entity"):
                staff.append(span["entity"])
        return list(set(staff))

    def _extract_quotes(self, spans: list[ExtractedSpan]) -> dict[str, str]:
        """Extract representative quotes by URT code."""
        quotes = {}
        for span in spans:
            code = span["urt_primary"]
            if code not in quotes:
                quotes[code] = span["span_text"][:100]
        return quotes

    def _generate_usn(self, span: LLMSpanResponse) -> str:
        """
        Generate USN (URT String Notation) for a span.

        Format: URT:S:{primary}[+{sec}]:{valence_sign}{intensity_num}:{S#}{A#}{temporal}.{evidence}.{CR_suffix}
        """
        primary = span.get("urt_primary", "O1.01")
        secondary = span.get("urt_secondary", [])
        valence = span.get("valence", "V0")
        intensity = span.get("intensity", "I1")
        specificity = span.get("specificity", "S1")
        actionability = span.get("actionability", "A1")
        temporal = span.get("temporal", "TC")
        evidence = span.get("evidence", "ES")
        comparative = span.get("comparative", "CR-N")

        # Build code portion
        code_part = primary
        for sec in secondary[:2]:
            code_part += f"+{sec}"

        # Valence encoding
        valence_map = {"V+": "+", "V-": "-", "V0": "0", "V±": "±"}
        valence_sign = valence_map.get(valence, "0")

        # Intensity number
        intensity_num = intensity[1] if intensity.startswith("I") else "1"

        # Dimensions
        spec_num = specificity[1] if specificity and specificity.startswith("S") else "1"
        act_num = actionability[1] if actionability and actionability.startswith("A") else "1"

        # CR suffix
        cr_map = {"CR-N": "N", "CR-B": "B", "CR-W": "W", "CR-S": "S"}
        cr_suffix = cr_map.get(comparative, "N")

        return f"URT:S:{code_part}:{valence_sign}{intensity_num}:{spec_num}{act_num}{temporal}.{evidence}.{cr_suffix}"

    async def _persist_classification(
        self,
        classified: ClassifiedReview,
        review: ReviewToClassify,
        batch_id: str,
        config: dict[str, Any],
    ) -> None:
        """Persist classification results to database."""
        if not self.review_repo or not self.span_repo:
            return

        # Update reviews_enriched
        await self.review_repo.update_enriched_with_classification(
            classified,
            self.config.llm_model,
            config["taxonomy_version"],
        )

        # Insert spans
        for span in classified.get("spans", []):
            await self.span_repo.insert_span(
                span,
                review["business_id"],
                review["place_id"],
                review["source"],
                review["review_id"],
                review["review_version"],
                review["review_time"],
                batch_id,
                self.config.llm_model,
                config["taxonomy_version"],
            )
