"""
Constraint analyzer for extracting constraints from queries.
"""

import re
from typing import List

from langchain_core.language_models import BaseChatModel
from loguru import logger

from .base_constraint import Constraint, ConstraintType


class ConstraintAnalyzer:
    """Analyzes queries to extract constraints."""

    def __init__(self, model: BaseChatModel):
        """Initialize the constraint analyzer."""
        self.model = model

    def extract_constraints(self, query: str) -> List[Constraint]:
        """Extract constraints from a query."""
        prompt = f"""
Generate constraints to verify if an answer candidate correctly answers this question.

Question: {query}

Create constraints that would help verify if a proposed answer is correct. Focus on the RELATIONSHIP between the question and answer, not just query analysis.

Examples:
- "Which university did Alice study at?" → "Alice studied at this university"
- "What year was the company founded?" → "The company was founded in this year"
- "Who invented the device?" → "This person invented the device"
- "Where is the building located?" → "The building is located at this place"

For each constraint, identify:
1. Type: property, name_pattern, event, statistic, temporal, location, comparison, existence
2. Description: What relationship must hold between question and answer
3. Value: The specific relationship to verify
4. Weight: How critical this constraint is (0.0-1.0)

Format your response as:
CONSTRAINT_1:
Type: [type]
Description: [description]
Value: [value]
Weight: [0.0-1.0]

CONSTRAINT_2:
Type: [type]
Description: [description]
Value: [value]
Weight: [0.0-1.0]

Focus on answer verification, not query parsing.
"""

        response = self.model.invoke(prompt)
        content = response.content

        constraints = []
        current_constraint = {}
        constraint_id = 1

        for line in content.strip().split("\n"):
            line = line.strip()

            if line.startswith("CONSTRAINT_"):
                if current_constraint and all(
                    k in current_constraint
                    for k in ["type", "description", "value"]
                ):
                    constraint = Constraint(
                        id=f"c{constraint_id}",
                        type=self._parse_constraint_type(
                            current_constraint["type"]
                        ),
                        description=current_constraint["description"],
                        value=current_constraint["value"],
                        weight=self._parse_weight(
                            current_constraint.get("weight", 1.0)
                        ),
                    )
                    constraints.append(constraint)
                    constraint_id += 1
                current_constraint = {}
            elif ":" in line:
                key, value = line.split(":", 1)
                key = key.strip().lower()
                value = value.strip()
                if key in ["type", "description", "value", "weight"]:
                    current_constraint[key] = value

        # Don't forget the last constraint
        if current_constraint and all(
            k in current_constraint for k in ["type", "description", "value"]
        ):
            constraint = Constraint(
                id=f"c{constraint_id}",
                type=self._parse_constraint_type(current_constraint["type"]),
                description=current_constraint["description"],
                value=current_constraint["value"],
                weight=self._parse_weight(
                    current_constraint.get("weight", 1.0)
                ),
            )
            constraints.append(constraint)

        logger.info(f"Extracted {len(constraints)} constraints from query")
        return constraints

    def _parse_constraint_type(self, type_str: str) -> ConstraintType:
        """Parse constraint type from string."""
        type_map = {
            "property": ConstraintType.PROPERTY,
            "name_pattern": ConstraintType.NAME_PATTERN,
            "event": ConstraintType.EVENT,
            "statistic": ConstraintType.STATISTIC,
            "temporal": ConstraintType.TEMPORAL,
            "location": ConstraintType.LOCATION,
            "comparison": ConstraintType.COMPARISON,
            "existence": ConstraintType.EXISTENCE,
        }
        return type_map.get(type_str.lower(), ConstraintType.PROPERTY)

    def _parse_weight(self, weight_value) -> float:
        """Parse weight value to float, handling text annotations.

        Args:
            weight_value: String or numeric weight value, possibly with text annotations

        Returns:
            float: Parsed weight value
        """
        if isinstance(weight_value, (int, float)):
            return float(weight_value)
        if isinstance(weight_value, str):
            # Extract the first number from the string
            match = re.search(r"(\d+(\.\d+)?)", weight_value)
            if match:
                return float(match.group(1))
        return 1.0  # Default weight
