import json
import logging
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import dataiku
from semantic_models.schemas import SemanticModelSpec

from semantic_models.config import get_app_config
from semantic_models.exceptions import ExecutionError
from semantic_models.resolution.index import (
    IndexStatus,
    SourceIndex,
    _is_simple_column,
    load_indices_for_spec,
)
from semantic_models.services.llm_utils import extract_json_block
from semantic_models.utils.spec_utils import (
    get_distinct_values_mode,
    get_sample_limit,
    is_auto_values_enabled,
    should_use_distinct_values,
)

logger = logging.getLogger(__name__)


@dataclass
class TermCorrection:
    """A correction to apply to a user query term."""

    original: str  # Original substring from user query
    resolved: str  # Resolved value from index (canonical)
    entity_id: str  # Entity the value belongs to
    attribute: str  # Attribute name
    score: float  # Match confidence (0-100)
    match_type: str  # "exact" | "fuzzy" | "semantic"
    source: str = "index"  # "index" | "glossary"
    glossary_term_id: Optional[str] = None  # ID of glossary term if source="glossary"
    llm_inferred: bool = False  # True if LLM translated the term (candidate_value != text_in_question)
    # Glossary-specific fields for UI display
    mapping_type: Optional[str] = None  # "attribute" | "metric" | "filter"
    mapping_name: Optional[str] = None  # Name of the mapped element


@dataclass
class ResolutionResult:
    """Result of resolving terms in a query."""

    original_question: str
    resolved_question: str
    corrections: List[TermCorrection]


@dataclass
class ExtractedTerm:
    """Intermediate representation of a term extracted by the LLM."""

    text_in_question: str
    entity_id: str
    attribute: str
    candidate_value: Optional[str] = None


def _get_term_llm_id(explicit_llm_id: Optional[str]) -> str:
    # Return explicit LLM ID immediately if provided (agent tool path)
    if explicit_llm_id:
        return explicit_llm_id

    # Fall back to config (webapp path)
    config = get_app_config()

    if config.term_extraction_llm_id:
        return config.term_extraction_llm_id

    raise ExecutionError("No LLM configured for term extraction")


def _build_term_extraction_prompt(
    question: str,
    spec: SemanticModelSpec,
    attribute_examples: dict[str, dict[str, list[str]]] | None = None,
) -> str:
    entity_blocks: List[str] = []
    examples_by_entity = attribute_examples or {}

    for entity in spec.entities:
        # Only include attributes that are resolvable AND simple columns
        resolvable_attrs = [
            a
            for a in entity.attributes
            if a.resolveInUserRequests
            and should_use_distinct_values(a)
            and _is_simple_column(
                {"name": a.name, "type": a.type, "column": a.column, "pseudoSQLExpression": a.pseudoSQLExpression}
            )
        ]
        if not resolvable_attrs:
            continue

        attr_lines = []
        for attr in resolvable_attrs:
            samples_list = (examples_by_entity.get(entity.name, {}) or {}).get(attr.name, [])
            samples = ", ".join(samples_list) if samples_list else ""
            line = f"  - {attr.name}: {attr.description or 'N/A'}"
            if samples:
                line += f" (examples: {samples})"
            attr_lines.append(line)

        block = f"""Entity: {entity.name} (ID: {entity.name})
Resolvable attributes:
{chr(10).join(attr_lines)}
"""
        entity_blocks.append(block)

    model_description = "\n".join(entity_blocks) if entity_blocks else "No resolvable attributes defined."

    prompt = f"""You are helping interpret a natural language question against a semantic data model.

SEMANTIC MODEL (resolvable attributes only):
{model_description}

USER QUESTION:
{question}

TASK:
Identify spans of text in the question that represent **specific data values** that should be looked up in the resolvable attributes listed above.

For each span found, determine:
1. text_in_question: the exact substring as it appears in the question
2. entity_id: which entity it belongs to (use the entity name)
3. attribute: which attribute name it matches
4. candidate_value: the actual data value to search for (NOT the attribute name)

Return ONLY valid JSON with this schema:
{{
  "items": [
    {{
      "text_in_question": "exact substring from question",
      "entity_id": "entity_id_here",
      "attribute": "attribute_name_here",
      "candidate_value": "actual data value to search for"
    }}
  ]
}}

Rules:
- Only extract terms that represent **actual data values** the user is filtering or searching for
- Do NOT extract terms that refer to the attribute/column itself (e.g., "by channel" means group by channel, not a value)
- candidate_value must be a **data value** (like "website", "US", "Consumer"), never an attribute name
- If the user says "by X" or "per X", they want to group by X, not filter for a specific value - do NOT include these
- Only include terms that clearly correspond to resolvable attributes listed above
- If uncertain about a mapping, do NOT include it
- Do not include explanations or text outside the JSON
- If no value terms match, return {{"items": []}}
"""

    return prompt


def _call_llm_for_extraction(
    question: str,
    spec: SemanticModelSpec,
    llm_id: Optional[str],
    attribute_examples: dict[str, dict[str, list[str]]] | None = None,
    trace_span: Any | None = None,
) -> List[ExtractedTerm]:
    chosen_llm_id = _get_term_llm_id(llm_id)
    logger.info("Calling LLM '%s' for term extraction", chosen_llm_id)

    prompt = _build_term_extraction_prompt(question, spec, attribute_examples)
    logger.info("Term extraction prompt prepared: chars=%d", len(prompt))
    logger.info("Term extraction prompt:\n%s", prompt)

    span = trace_span.subspan("term_extraction") if trace_span is not None else None
    if span:
        span.__enter__()
        span.attributes["llm_id"] = chosen_llm_id
        span.inputs["question"] = question
        span.inputs["prompt"] = prompt

    try:
        client = dataiku.api_client()
        project = client.get_default_project()
        llm = project.get_llm(chosen_llm_id)
        completion = llm.new_completion()
        completion.with_message(prompt)

        response = completion.execute()
        raw_text = response.text.strip()
        response_trace = getattr(response, "trace", None)
        if span and response_trace:
            span.append_trace(response_trace)
        if span:
            span.outputs["llm_response"] = raw_text
        logger.info("Term extraction LLM response received: chars=%d", len(raw_text))
        logger.info("Term extraction LLM response:\n%s", raw_text)

    except Exception as e:
        if span:
            span.attributes["status"] = "error"
            span.outputs["error"] = str(e)
        raise ExecutionError(f"LLM call for term extraction failed: {e}")
    finally:
        if span:
            span.__exit__(None, None, None)

    # Parse JSON response
    json_text = extract_json_block(raw_text)
    try:
        payload = json.loads(json_text)
    except json.JSONDecodeError as e:
        logger.warning("Failed to parse LLM JSON response: %s", e)
        return []

    items = payload.get("items") or []
    extracted: List[ExtractedTerm] = []

    for item in items:
        try:
            text_in_question = item["text_in_question"]
            entity_id = item["entity_id"]
            attribute = item["attribute"]
            candidate_value = item.get("candidate_value") or text_in_question
        except (KeyError, TypeError):
            # Skip malformed items
            continue

        extracted.append(
            ExtractedTerm(
                text_in_question=text_in_question,
                entity_id=entity_id,
                attribute=attribute,
                candidate_value=candidate_value,
            )
        )

    logger.info("Term extraction complete: %d candidate terms extracted", len(extracted))
    for term in extracted:
        logger.info(
            "  - Extracted term: '%s' -> entity=%s, attr=%s, candidate='%s'",
            term.text_in_question,
            term.entity_id,
            term.attribute,
            term.candidate_value,
        )
    return extracted


def _build_attribute_examples(
    spec: SemanticModelSpec,
    semantic_model_id: str,
    version_id: str,
    private_editor_data: Optional[dict],
) -> dict[str, dict[str, list[str]]]:
    examples: dict[str, dict[str, list[str]]] = {}

    # TODO: add in-memory cache for sample values per (model, version, entity, attribute)
    for entity in spec.entities:
        if entity.type != "DATASET":
            continue
        # Only include resolvable simple columns
        resolvable_attrs = [
            a
            for a in entity.attributes
            if a.resolveInUserRequests
            and should_use_distinct_values(a)
            and _is_simple_column(
                {"name": a.name, "type": a.type, "column": a.column, "pseudoSQLExpression": a.pseudoSQLExpression}
            )
        ]
        if not resolvable_attrs:
            continue

        for attr in resolvable_attrs:
            values: list[str] = []

            if attr.manualValues:
                values = list(attr.manualValues)
                logger.info(
                    "Using MANUAL sample values for term extraction: %s.%s (%d values)",
                    entity.name,
                    attr.name,
                    len(values),
                )
            elif is_auto_values_enabled(attr):
                limit = get_sample_limit(attr)
                if limit:
                    try:
                        from semantic_models.resolution import SourceKey
                        from semantic_models.services.column_values import fetch_column_values

                        source = SourceKey(
                            semantic_model_id=semantic_model_id,
                            version_id=version_id,
                            entity_name=entity.name,
                        )
                        fetch_limit = limit
                        values = fetch_column_values(source, attr.name, fetch_limit)
                        if values:
                            logger.info(
                                "Using AUTO sample values for term extraction: %s.%s (%d values)",
                                entity.name,
                                attr.name,
                                len(values),
                            )
                    except Exception as e:
                        logger.warning(
                            "Failed to fetch sample values for %s.%s: %s", entity.name, attr.name, e
                        )
            elif get_distinct_values_mode(attr) == "MANUAL":
                logger.info(
                    "MANUAL mode with empty sample values for term extraction: %s.%s",
                    entity.name,
                    attr.name,
                )

            if values:
                examples.setdefault(entity.name, {})[attr.name] = values

    return examples


def resolve_terms(
    question: str,
    spec: SemanticModelSpec,
    spec_dict: dict,
    semantic_model_id: str,
    version_id: str,
    llm_id: Optional[str] = None,
    embedding_llm_id: Optional[str] = None,
    private_editor_data: Optional[dict] = None,
    trace_span: Any | None = None,
) -> ResolutionResult:
    # Load only the indices needed for this spec's entities
    entity_indices: Dict[str, SourceIndex] = load_indices_for_spec(
        spec_dict,
        semantic_model_id=semantic_model_id,
        version_id=version_id,
        embedding_llm_id=embedding_llm_id,
    )

    if not entity_indices:
        logger.info("No source indices available, skipping term resolution")
        return ResolutionResult(
            original_question=question,
            resolved_question=question,
            corrections=[],
        )

    # Check if there are any resolvable simple-column attributes
    has_resolvable = any(
        attr.resolveInUserRequests
        and should_use_distinct_values(attr)
        and _is_simple_column(
            {"name": attr.name, "type": attr.type, "column": attr.column, "pseudoSQLExpression": attr.pseudoSQLExpression}
        )
        for entity in spec.entities
        for attr in entity.attributes
    )
    if not has_resolvable:
        logger.debug("No resolvable simple-column attributes in spec, skipping term extraction")
        return ResolutionResult(
            original_question=question,
            resolved_question=question,
            corrections=[],
        )

    # Extract terms using LLM
    try:
        examples = _build_attribute_examples(spec, semantic_model_id, version_id, private_editor_data)
        extracted_terms = _call_llm_for_extraction(question, spec, llm_id, examples, trace_span=trace_span)
    except Exception as e:
        logger.warning("Term extraction failed: %s", e)
        return ResolutionResult(
            original_question=question,
            resolved_question=question,
            corrections=[],
        )

    if not extracted_terms:
        return ResolutionResult(
            original_question=question,
            resolved_question=question,
            corrections=[],
        )

    corrections: List[TermCorrection] = []
    used_originals: Dict[str, TermCorrection] = {}

    for term in extracted_terms:
        # Validate entity exists
        entity = next((e for e in spec.entities if e.name == term.entity_id), None)
        if entity is None:
            logger.debug("Ignoring term for unknown entity_id '%s'", term.entity_id)
            continue

        # Validate attribute exists, is resolvable, and is a simple column
        attr = next(
            (
                a
                for a in entity.attributes
                if a.name == term.attribute
                and a.resolveInUserRequests
                and should_use_distinct_values(a)
                and _is_simple_column(
                    {"name": a.name, "type": a.type, "column": a.column, "pseudoSQLExpression": a.pseudoSQLExpression}
                )
            ),
            None,
        )
        if attr is None:
            logger.debug(
                "Ignoring term for non-resolvable/computed attribute '%s' on entity '%s'",
                term.attribute,
                term.entity_id,
            )
            continue

        # Skip if we already resolved this text
        if term.text_in_question in used_originals:
            continue

        # Skip if candidate_value is the attribute name itself (not an actual value)
        candidate = term.candidate_value or term.text_in_question
        if candidate.lower() == term.attribute.lower():
            logger.debug(
                "Skipping term '%s' - candidate_value '%s' equals attribute name", term.text_in_question, candidate
            )
            continue

        # Get the source index for this entity
        source_index = entity_indices.get(term.entity_id)
        if source_index is None:
            logger.debug("No source index loaded for entity '%s'", term.entity_id)
            continue

        # Determine search term: use LLM's candidate if provided, else user's text
        search_term = term.candidate_value or term.text_in_question
        llm_inferred = (
            term.candidate_value is not None and term.candidate_value.lower() != term.text_in_question.lower()
        )

        # Search the source index using the column name (same as attribute name for simple columns)
        match = source_index.search_best_value(
            term.attribute,  # column name = attribute name for simple columns
            search_term,
        )
        if match is None:
            continue

        resolved_value, score, match_type = match

        # Skip only if resolved value is exactly identical (including case)
        # Case differences matter for case-sensitive databases (e.g., "us" → "US")
        if resolved_value == term.text_in_question:
            continue

        correction = TermCorrection(
            original=term.text_in_question,
            resolved=resolved_value,
            entity_id=term.entity_id,
            attribute=term.attribute,
            score=score,
            match_type=match_type,
            llm_inferred=llm_inferred,
        )
        corrections.append(correction)
        used_originals[term.text_in_question] = correction

        if llm_inferred:
            logger.info(
                "Resolved '%s' -> '%s' (entity=%s, attr=%s, LLM inferred '%s', then %s match, score=%.1f)",
                correction.original,
                correction.resolved,
                correction.entity_id,
                correction.attribute,
                search_term,
                correction.match_type,
                correction.score,
            )
        else:
            logger.info(
                "Resolved '%s' -> '%s' (entity=%s, attr=%s, %s match, score=%.1f)",
                correction.original,
                correction.resolved,
                correction.entity_id,
                correction.attribute,
                correction.match_type,
                correction.score,
            )

    # Apply corrections to the question
    # Sort by length (longest first) to avoid overlapping replacements interfering
    sorted_corrections = sorted(corrections, key=lambda c: len(c.original), reverse=True)

    resolved_question = question
    for correction in sorted_corrections:
        escaped = re.escape(correction.original)
        # Use word boundaries for word-like terms to avoid substring matches
        # e.g., "US" shouldn't match inside "business"
        if re.match(r"^\w+$", correction.original):
            pattern = re.compile(rf"\b{escaped}\b", re.IGNORECASE)
        else:
            pattern = re.compile(escaped, re.IGNORECASE)
        resolved_question = pattern.sub(correction.resolved, resolved_question)

    # Log resolution summary
    if corrections:
        correction_details = [
            {
                "original": c.original,
                "resolved": c.resolved,
                "entity_id": c.entity_id,
                "attribute": c.attribute,
                "score": c.score,
                "match_type": c.match_type,
                "source": c.source,
            }
            for c in corrections
        ]
        logger.info(
            "Term resolution summary: original=%s resolved=%s corrections=%d",
            question,
            resolved_question,
            len(corrections),
        )
        logger.info("Term resolution details:\n%s", json.dumps(correction_details, indent=2, ensure_ascii=False))
    else:
        logger.info("Term resolution: no corrections needed for question=%s", question)

    return ResolutionResult(
        original_question=question,
        resolved_question=resolved_question,
        corrections=corrections,
    )


def resolve_glossary_terms(
    question: str,
    semantic_model_id: str,
    version_id: str,
    spec: SemanticModelSpec,
) -> List[TermCorrection]:
    from semantic_models.services.glossary_utils import load_glossary_terms_for_prompt

    bound_terms = load_glossary_terms_for_prompt(
        semantic_model_id=semantic_model_id,
        spec=spec,
        version_id=version_id,
        limit=0,
    )

    if not bound_terms:
        logger.debug(
            "No bound glossary terms for model %s version %s",
            semantic_model_id,
            version_id,
        )
        return []

    # Build phrase map: normalized phrase -> list of bound terms
    # (multiple bindings for same phrase possible due to M:N cardinality)
    phrase_map: Dict[str, list] = {}
    for bound_term in bound_terms:
        phrases = [bound_term["term"]] + (bound_term.get("synonyms") or [])
        for phrase in phrases:
            normalized = phrase.lower().strip()
            if normalized not in phrase_map:
                phrase_map[normalized] = []
            phrase_map[normalized].append(bound_term)

    if not phrase_map:
        logger.debug("No valid glossary phrases to match against")
        return []

    logger.debug("Glossary phrase map has %d entries", len(phrase_map))

    # Search for matches in question
    corrections: List[TermCorrection] = []
    matched_spans: set = set()  # Track matched character positions

    # Sort phrases by length (longest first) to prefer longer matches
    sorted_phrases = sorted(phrase_map.keys(), key=len, reverse=True)

    for phrase in sorted_phrases:
        bound_term_list = phrase_map[phrase]

        # Disambiguation: prefer primary binding
        primary_terms = [t for t in bound_term_list if t.get("is_primary")]
        chosen_term = primary_terms[0] if primary_terms else bound_term_list[0]

        # Build regex pattern with word boundaries
        escaped = re.escape(phrase)
        pattern = re.compile(rf"\b{escaped}\b", re.IGNORECASE)

        for match in pattern.finditer(question):
            start, end = match.span()

            # Skip if this span overlaps with already matched text
            if any(start <= pos < end for pos in matched_spans):
                continue

            original_text = match.group(0)

            # Build mapping info from binding structure
            target_type = chosen_term.get("target_type")
            target_name = chosen_term.get("target_name")
            entity_id = chosen_term.get("entity_id", "")

            # Determine attr_name for backward compatibility
            attr_name = ""
            if target_type == "attribute" and target_name:
                attr_name = target_name

            correction = TermCorrection(
                original=original_text,
                resolved=chosen_term["term"],  # Canonical term
                entity_id=entity_id,
                attribute=attr_name,
                score=100.0,  # Exact glossary match
                match_type="exact",
                source="glossary",
                glossary_term_id=chosen_term.get("term_id"),
                mapping_type=target_type,
                mapping_name=target_name,
            )
            corrections.append(correction)

            # Mark this span as matched
            matched_spans.update(range(start, end))

            logger.info(
                "Glossary match: '%s' -> '%s' (entity=%s, %s=%s)",
                original_text,
                chosen_term["term"],
                entity_id,
                target_type or "entity",
                target_name or entity_id,
            )

    return corrections
