from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Optional, Union, Type, Dict, Any
import json

import polars as pl

from polar_llama.utils import parse_into_expr, register_plugin, parse_version

if TYPE_CHECKING:
    from polars.type_aliases import IntoExpr
    from pydantic import BaseModel

if parse_version(pl.__version__) < parse_version("0.20.16"):
    from polars.utils.udfs import _get_shared_lib_location

    lib: str | Path = _get_shared_lib_location(__file__)
else:
    lib = Path(__file__).parent

# Import Provider enum directly from the extension module
try:
    # First try relative import from the extension module in current directory
    from .polar_llama import Provider
except ImportError:
    # Fallback to try absolute import
    try:
        from polar_llama.polar_llama import Provider
    except ImportError:
        # Define a basic Provider class as fallback if neither import works
        class Provider:
            OPENAI = "openai"
            ANTHROPIC = "anthropic"
            GEMINI = "gemini"
            GROQ = "groq"
            BEDROCK = "bedrock"
            
            def __init__(self, provider_str):
                self.value = provider_str
                
            def __str__(self):
                return self.value

# Import and initialize the expressions helper to ensure expressions are registered
from polar_llama.expressions import ensure_expressions_registered, get_lib_path

# Ensure the expressions are registered
ensure_expressions_registered()
# Update the lib path to make sure we're using the actual library
lib = get_lib_path()

def _pydantic_to_json_schema(model: Type['BaseModel']) -> dict:
    """Convert a Pydantic model to JSON schema."""
    try:
        from pydantic import BaseModel
        if not issubclass(model, BaseModel):
            raise ValueError("response_model must be a Pydantic BaseModel subclass")

        # Get the JSON schema from the Pydantic model
        schema = model.model_json_schema()

        # Recursively add additionalProperties: false to all objects
        # This is required by some providers like Groq
        # However, we skip objects that already have additionalProperties defined
        # (like Dict[str, str] types which need dynamic keys)
        def add_additional_properties_false(obj):
            if isinstance(obj, dict):
                if obj.get("type") == "object" and "additionalProperties" not in obj:
                    # Only add if not already present (Dict types already have it defined)
                    obj["additionalProperties"] = False
                # Recursively process nested objects
                for key, value in obj.items():
                    if isinstance(value, dict):
                        add_additional_properties_false(value)
                    elif isinstance(value, list):
                        for item in value:
                            if isinstance(item, dict):
                                add_additional_properties_false(item)

        add_additional_properties_false(schema)
        return schema
    except ImportError:
        raise ImportError("Pydantic is required for structured outputs. Install with: pip install pydantic>=2.0.0")

def _json_schema_to_polars_dtype(schema: dict, root_schema: dict = None) -> pl.DataType:
    """Convert a JSON schema to a Polars DataType."""
    # Keep reference to root schema for resolving $ref
    if root_schema is None:
        root_schema = schema

    properties = schema.get("properties", {})

    fields = []
    for field_name, field_schema in properties.items():
        # Handle $ref references
        if "$ref" in field_schema:
            # Extract the reference path (e.g., "#/$defs/SentimentResult")
            ref_path = field_schema["$ref"]
            if ref_path.startswith("#/"):
                # Navigate to the referenced schema
                ref_parts = ref_path[2:].split("/")
                ref_schema = root_schema
                for part in ref_parts:
                    ref_schema = ref_schema.get(part, {})
                # Recursively convert the referenced schema
                pl_type = _json_schema_to_polars_dtype(ref_schema, root_schema)
            else:
                # Unknown reference format, default to string
                pl_type = pl.Utf8
        else:
            field_type = field_schema.get("type")

            # Map JSON schema types to Polars types
            if field_type == "string":
                pl_type = pl.Utf8
            elif field_type == "integer":
                pl_type = pl.Int64
            elif field_type == "number":
                pl_type = pl.Float64
            elif field_type == "boolean":
                pl_type = pl.Boolean
            elif field_type == "array":
                # Handle arrays - use List type
                items_schema = field_schema.get("items", {})
                items_type = items_schema.get("type", "string")
                if items_type == "string":
                    pl_type = pl.List(pl.Utf8)
                elif items_type == "integer":
                    pl_type = pl.List(pl.Int64)
                elif items_type == "number":
                    pl_type = pl.List(pl.Float64)
                else:
                    pl_type = pl.List(pl.Utf8)  # Default to string list
            elif field_type == "object":
                # Check if this is a Dict[str, str] type (has additionalProperties)
                if "additionalProperties" in field_schema and field_schema["additionalProperties"] != False:
                    # This is a dictionary type, just use Utf8 for now
                    # (Polars doesn't have a good way to represent arbitrary dicts in structs)
                    pl_type = pl.Utf8
                else:
                    # Nested objects - recursively convert
                    pl_type = _json_schema_to_polars_dtype(field_schema, root_schema)
            else:
                # Default to string for unknown types
                pl_type = pl.Utf8

        fields.append(pl.Field(field_name, pl_type))

    # Add error fields for error handling (only at the root level)
    if root_schema == schema:
        fields.append(pl.Field("_error", pl.Utf8))
        fields.append(pl.Field("_details", pl.Utf8))
        fields.append(pl.Field("_raw", pl.Utf8))

    return pl.Struct(fields)

def _parse_json_to_struct(json_str_series: pl.Series, dtype: pl.DataType) -> pl.Series:
    """Parse a JSON string series into a struct series."""
    # Use Polars' built-in JSON parsing
    try:
        # First, try to decode as JSON
        parsed = json_str_series.str.json_decode()

        # Cast to the target struct type to ensure field types match
        return parsed.cast(dtype, strict=False)
    except Exception:
        # If parsing fails, return the series as-is (will contain error objects)
        return json_str_series.str.json_decode()


# ============================================================================
# Taxonomy-based Tagging
# ============================================================================

def _create_taxonomy_pydantic_model(taxonomy: Dict[str, Dict[str, Any]]) -> Type['BaseModel']:
    """
    Create a Pydantic model from a taxonomy definition.

    The taxonomy should be a dict with the following structure:
    {
        "field_name": {
            "description": "Description of the field",
            "values": {
                "value1": "Definition of value1",
                "value2": "Definition of value2",
                ...
            }
        },
        ...
    }

    Each field in the output model will be a struct containing:
    - thinking: Dict[str, str] - reasoning for each possible value
    - reflection: str - overall reflection on the field analysis
    - value: str - the selected value
    - confidence: float - confidence in the selection (0.0 to 1.0)
    """
    try:
        from pydantic import BaseModel, Field, create_model
        from typing import Dict

        # Create a field result model for each taxonomy field
        field_models = {}

        for field_name, field_config in taxonomy.items():
            values = field_config.get("values", {})
            value_names = list(values.keys())

            # Create the field result model with dynamic thinking keys
            field_result_model = create_model(
                f"{field_name.title()}Result",
                thinking=(Dict[str, str], Field(..., description=f"Reasoning for each possible value: {', '.join(value_names)}")),
                reflection=(str, Field(..., description="Overall reflection on your analysis of this field")),
                value=(str, Field(..., description=f"Selected value from: {', '.join(value_names)}")),
                confidence=(float, Field(..., ge=0.0, le=1.0, description="Confidence in the selected value (0.0 to 1.0)"))
            )

            field_models[field_name] = (field_result_model, Field(..., description=field_config.get("description", "")))

        # Create the main taxonomy result model
        TaxonomyResult = create_model("TaxonomyResult", **field_models)

        return TaxonomyResult

    except ImportError:
        raise ImportError("Pydantic is required for taxonomy tagging. Install with: pip install pydantic>=2.0.0")


def _create_taxonomy_prompt(taxonomy: Dict[str, Dict[str, Any]], document_field_name: str = "document") -> str:
    """
    Create a system prompt for taxonomy-based tagging.

    This prompt instructs the model to analyze a document according to the
    provided taxonomy and return structured tags with reasoning.
    """
    prompt_parts = [
        f"You are an expert document analyst. Analyze the provided {document_field_name} and tag it according to the following taxonomy.",
        "",
        "# Taxonomy Fields",
        ""
    ]

    for field_name, field_config in taxonomy.items():
        description = field_config.get("description", "")
        values = field_config.get("values", {})

        prompt_parts.append(f"## {field_name}")
        if description:
            prompt_parts.append(f"{description}")
        prompt_parts.append("")
        prompt_parts.append("Possible values:")

        for value_name, value_definition in values.items():
            prompt_parts.append(f"- **{value_name}**: {value_definition}")

        prompt_parts.append("")

    prompt_parts.extend([
        "# Instructions",
        "",
        "For each field in the taxonomy:",
        "",
        "1. **Thinking**: Consider each possible value and write your reasoning for why it might or might not apply to the document. Provide one reasoning string for each possible value.",
        "",
        "2. **Reflection**: After thinking through all values, reflect on your analysis. Consider which value best fits the document and why.",
        "",
        "3. **Value**: Select the single best value from the possible values for this field.",
        "",
        "4. **Confidence**: Provide your confidence in this selection as a number between 0.0 (not confident) and 1.0 (very confident).",
        "",
        "Return your analysis in the structured format with all required fields."
    ])

    return "\n".join(prompt_parts)

def inference_async(
    expr: IntoExpr,
    *,
    provider: Optional[Union[str, Provider]] = None,
    model: Optional[str] = None,
    response_model: Optional[Type['BaseModel']] = None
) -> pl.Expr:
    """
    Asynchronously infer completions for the given text expressions using an LLM.

    Parameters
    ----------
    expr : polars.Expr
        The text expression to use for inference
    provider : str or Provider, optional
        The provider to use (OpenAI, Anthropic, Gemini, Groq, Bedrock)
    model : str, optional
        The model name to use
    response_model : Type[BaseModel], optional
        A Pydantic model class to define structured output schema.
        The LLM response will be validated against this schema.
        Returns a Struct with fields matching the Pydantic model.

    Returns
    -------
    polars.Expr
        Expression with inferred completions as a Struct (if response_model provided)
        or String (if no response_model)
    """
    expr = parse_into_expr(expr)
    kwargs = {}

    if provider is not None:
        # Convert Provider to string to make it picklable
        if isinstance(provider, Provider):
            provider = str(provider)
        kwargs["provider"] = provider

    if model is not None:
        kwargs["model"] = model

    struct_dtype = None
    if response_model is not None:
        schema = _pydantic_to_json_schema(response_model)
        # Pass the JSON schema as a JSON string to Rust
        kwargs["response_schema"] = json.dumps(schema)
        kwargs["response_model_name"] = response_model.__name__
        # Create the target struct dtype for later conversion
        struct_dtype = _json_schema_to_polars_dtype(schema)

    result_expr = register_plugin(
        args=[expr],
        symbol="inference_async",
        is_elementwise=True,
        lib=lib,
        kwargs=kwargs,
    )

    # If response_model was provided, convert JSON strings to structs
    if struct_dtype is not None:
        # Use map_batches to convert the JSON string series to struct series
        result_expr = result_expr.map_batches(
            lambda s: _parse_json_to_struct(s, struct_dtype),
            return_dtype=struct_dtype
        )

    return result_expr

def inference(
    expr: IntoExpr,
    *,
    provider: Optional[Union[str, Provider]] = None,
    model: Optional[str] = None,
    response_model: Optional[Type['BaseModel']] = None
) -> pl.Expr:
    """
    Synchronously infer completions for the given text expressions using an LLM.

    Parameters
    ----------
    expr : polars.Expr
        The text expression to use for inference
    provider : str or Provider, optional
        The provider to use (OpenAI, Anthropic, Gemini, Groq, Bedrock)
    model : str, optional
        The model name to use
    response_model : Type[BaseModel], optional
        A Pydantic model class to define structured output schema.
        The LLM response will be validated against this schema.
        Returns a Struct with fields matching the Pydantic model.

    Returns
    -------
    polars.Expr
        Expression with inferred completions as a Struct (if response_model provided)
        or String (if no response_model)
    """
    expr = parse_into_expr(expr)
    kwargs = {}

    if provider is not None:
        # Convert Provider to string to make it picklable
        if isinstance(provider, Provider):
            provider = str(provider)
        kwargs["provider"] = provider

    if model is not None:
        kwargs["model"] = model

    struct_dtype = None
    if response_model is not None:
        schema = _pydantic_to_json_schema(response_model)
        # Pass the JSON schema as a JSON string to Rust
        kwargs["response_schema"] = json.dumps(schema)
        kwargs["response_model_name"] = response_model.__name__
        # Create the target struct dtype for later conversion
        struct_dtype = _json_schema_to_polars_dtype(schema)

    result_expr = register_plugin(
        args=[expr],
        symbol="inference",
        is_elementwise=True,
        lib=lib,
        kwargs=kwargs,
    )

    # If response_model was provided, convert JSON strings to structs
    if struct_dtype is not None:
        # Use map_batches to convert the JSON string series to struct series
        result_expr = result_expr.map_batches(
            lambda s: _parse_json_to_struct(s, struct_dtype),
            return_dtype=struct_dtype
        )

    return result_expr

def inference_messages(
    expr: IntoExpr,
    *,
    provider: Optional[Union[str, Provider]] = None,
    model: Optional[str] = None,
    response_model: Optional[Type['BaseModel']] = None
) -> pl.Expr:
    """
    Process message arrays (conversations) for inference using LLMs.

    This function accepts properly formatted JSON message arrays and sends them
    to the LLM for inference while preserving conversation context.

    Parameters
    ----------
    expr : polars.Expr
        The expression containing JSON message arrays
    provider : str or Provider, optional
        The provider to use (OpenAI, Anthropic, Gemini, Groq, Bedrock)
    model : str, optional
        The model name to use
    response_model : Type[BaseModel], optional
        A Pydantic model class to define structured output schema.
        The LLM response will be validated against this schema.
        Returns a Struct with fields matching the Pydantic model.

    Returns
    -------
    polars.Expr
        Expression with inferred completions as a Struct (if response_model provided)
        or String (if no response_model)
    """
    expr = parse_into_expr(expr)
    kwargs = {}

    if provider is not None:
        # Convert Provider to string to make it picklable
        if hasattr(provider, 'as_str'):
            provider_str = provider.as_str()
        elif hasattr(provider, '__str__'):
            provider_str = str(provider)
        else:
            provider_str = provider

        kwargs["provider"] = provider_str

    if model is not None:
        kwargs["model"] = model

    struct_dtype = None
    if response_model is not None:
        schema = _pydantic_to_json_schema(response_model)
        # Pass the JSON schema as a JSON string to Rust
        kwargs["response_schema"] = json.dumps(schema)
        kwargs["response_model_name"] = response_model.__name__
        # Create the target struct dtype for later conversion
        struct_dtype = _json_schema_to_polars_dtype(schema)

    # Don't pass empty kwargs dictionary
    if not kwargs:
        result_expr = register_plugin(
            args=[expr],
            symbol="inference_messages",
            is_elementwise=True,
            lib=lib,
        )
    else:
        result_expr = register_plugin(
            args=[expr],
            symbol="inference_messages",
            is_elementwise=True,
            lib=lib,
            kwargs=kwargs,
        )

    # If response_model was provided, convert JSON strings to structs
    if struct_dtype is not None:
        # Use map_batches to convert the JSON string series to struct series
        result_expr = result_expr.map_batches(
            lambda s: _parse_json_to_struct(s, struct_dtype),
            return_dtype=struct_dtype
        )

    return result_expr

def string_to_message(expr: IntoExpr, *, message_type: str) -> pl.Expr:
    """
    Convert a string to a message with the specified type.
    
    Parameters
    ----------
    expr : polars.Expr
        The text expression to convert
    message_type : str
        The type of message to create ("user", "system", "assistant")
        
    Returns
    -------
    polars.Expr
        Expression with formatted messages
    """
    expr = parse_into_expr(expr)
    return register_plugin(
        args=[expr],
        symbol="string_to_message",
        is_elementwise=True,
        lib=lib,
        kwargs={"message_type": message_type},
    )

def combine_messages(*exprs: IntoExpr) -> pl.Expr:
    """
    Combine multiple message expressions into a single message array.

    This function takes multiple message expressions (strings containing JSON formatted messages)
    and combines them into a single JSON array of messages, preserving the order.

    Parameters
    ----------
    *exprs : polars.Expr
        One or more expressions containing messages to combine

    Returns
    -------
    polars.Expr
        Expression with combined message arrays
    """
    args = [parse_into_expr(expr) for expr in exprs]

    return register_plugin(
        args=args,
        symbol="combine_messages",
        is_elementwise=True,
        lib=lib,
    )


def tag_taxonomy(
    expr: IntoExpr,
    taxonomy: Dict[str, Dict[str, Any]],
    *,
    provider: Optional[Union[str, Provider]] = None,
    model: Optional[str] = None,
) -> pl.Expr:
    """
    Tag documents according to a taxonomy definition with detailed reasoning.

    This function analyzes documents and assigns tags based on a predefined taxonomy,
    providing detailed reasoning for each classification decision. The taxonomy allows
    you to define fields (categories), their possible values, and definitions for each value.

    For each taxonomy field, the model will:
    1. Think through each possible value with reasoning
    2. Reflect on the overall analysis
    3. Select the best value
    4. Provide a confidence score

    Parameters
    ----------
    expr : polars.Expr
        The document expression to analyze and tag
    taxonomy : Dict[str, Dict[str, Any]]
        A dictionary defining the taxonomy structure:
        {
            "field_name": {
                "description": "Description of what this field represents",
                "values": {
                    "value1": "Definition of value1",
                    "value2": "Definition of value2",
                    ...
                }
            },
            ...
        }
    provider : str or Provider, optional
        The LLM provider to use (OpenAI, Anthropic, Gemini, Groq, Bedrock)
    model : str, optional
        The specific model name to use

    Returns
    -------
    polars.Expr
        Expression with structured tags as a Struct. Each taxonomy field becomes
        a nested struct containing:
        - thinking: Dict[str, str] - reasoning for each possible value
        - reflection: str - overall reflection on the field analysis
        - value: str - the selected value
        - confidence: float - confidence score (0.0 to 1.0)

    Examples
    --------
    >>> import polars as pl
    >>> from polar_llama import tag_taxonomy, Provider
    >>>
    >>> # Define a taxonomy
    >>> taxonomy = {
    ...     "sentiment": {
    ...         "description": "The emotional tone of the text",
    ...         "values": {
    ...             "positive": "Text expresses positive emotions, optimism, or favorable opinions",
    ...             "negative": "Text expresses negative emotions, pessimism, or unfavorable opinions",
    ...             "neutral": "Text is factual and objective without clear emotional content"
    ...         }
    ...     },
    ...     "urgency": {
    ...         "description": "How urgent or time-sensitive the content is",
    ...         "values": {
    ...             "high": "Requires immediate attention or action",
    ...             "medium": "Should be addressed soon but not immediately critical",
    ...             "low": "Can be addressed at any convenient time"
    ...         }
    ...     }
    ... }
    >>>
    >>> # Create a dataframe with documents
    >>> df = pl.DataFrame({
    ...     "id": [1, 2],
    ...     "document": [
    ...         "URGENT: The server is down and customers can't access the site!",
    ...         "Our quarterly results exceeded expectations. Great work team!"
    ...     ]
    ... })
    >>>
    >>> # Apply taxonomy tagging
    >>> result = df.with_columns(
    ...     tags=tag_taxonomy(
    ...         pl.col("document"),
    ...         taxonomy,
    ...         provider=Provider.OPENAI,
    ...         model="gpt-4"
    ...     )
    ... )
    >>>
    >>> # Access specific fields and values
    >>> result.select([
    ...     "document",
    ...     pl.col("tags").struct.field("sentiment").struct.field("value").alias("sentiment"),
    ...     pl.col("tags").struct.field("sentiment").struct.field("confidence").alias("sentiment_conf"),
    ...     pl.col("tags").struct.field("urgency").struct.field("value").alias("urgency")
    ... ])
    """
    # Create the Pydantic model from the taxonomy
    response_model = _create_taxonomy_pydantic_model(taxonomy)

    # Create the system prompt with taxonomy instructions
    system_prompt = _create_taxonomy_prompt(taxonomy, document_field_name="document")

    # Parse the document expression
    doc_expr = parse_into_expr(expr)

    # Create a system message for each row by mapping over the document column
    # This ensures the system message is broadcast to match the number of rows
    system_message_expr = doc_expr.map_batches(
        lambda s: pl.Series([system_prompt] * len(s)),
        return_dtype=pl.Utf8
    ).pipe(string_to_message, message_type="system")

    # Create a user message with the document
    user_message_expr = doc_expr.pipe(
        string_to_message, message_type="user"
    )

    # Combine the messages
    messages_expr = combine_messages(system_message_expr, user_message_expr)

    # Call inference_messages with the structured output model
    return inference_messages(
        messages_expr,
        provider=provider,
        model=model,
        response_model=response_model
    )