import json
import re
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field, model_validator


class FeatureOutput(BaseModel):
    title: str = Field(..., description="A concise, descriptive name for the feature.")

    columns_involved: List[str] = Field(
        ..., description="List of all source columns used for this feature."
    )

    operation: str = Field(
        ..., description="The specific operation name used from the allowed list."
    )

    created_new_feature: Union[str, List[str]] = Field(
        ..., description="The name(s) of the new feature column(s) created by this SQL query."
    )

    formula_logic: Union[str, List[str]] = Field(
        ...,
        description="A human-readable description or mathematical formula of how each feature is calculated.",
    )

    reasoning: str = Field(
        ..., description="Explain WHY this feature is useful for the model, 1-2 lines."
    )

    group_by: Optional[List[str]] = Field(
        default_factory=list, description="Columns for aggregation/partition."
    )

    order_by: Optional[List[str]] = Field(
        default_factory=list, description="Columns to sort by for window functions."
    )

    window: Optional[int] = Field(
        None, description="Window size for rolling or lag operations."
    )

    sql: str = Field(
        ...,
        description="SQL query to generate this feature. Use column names as provided.",
    )


class FeatureOutputSqlQueryValidation(BaseModel):
    title: str = Field(..., description="A concise, descriptive name for the feature.")
    original_sql: str = Field(..., description="SQL query to generate this feature.")
    corrected_sql: Optional[str] = Field(
        ..., description="Corrected SQL query if the original had issues."
    )
    is_valid: bool = Field(..., description="Indicates if the original SQL was valid.")
    validation_message: str = Field(
        ..., description="Details about the validation result."
    )


class FeaturesListModel(BaseModel):
    features: List[FeatureOutput]

    @model_validator(mode="before")
    @classmethod
    def parse_llm_output(cls, raw_llm_output: Any) -> Any:
        """
        Accepts raw LLM output (str or list of dicts) and ensures we get a list of dicts.
        Removes markdown fences if needed and parses JSON automatically.
        """
        if isinstance(raw_llm_output, str):
            cleaned = raw_llm_output.strip()
            # Remove code fences
            if cleaned.startswith("```"):
                cleaned = re.sub(
                    r"^```(?:json)?", "", cleaned, flags=re.IGNORECASE
                ).strip()
                cleaned = re.sub(r"```$", "", cleaned).strip()
            try:
                raw_llm_output = json.loads(cleaned)
            except json.JSONDecodeError as e:
                raise ValueError(f"Invalid JSON from LLM output: {e}")

        if not isinstance(raw_llm_output, list):
            raise ValueError("LLM output must be a list of feature dicts")

        return {"features": raw_llm_output}


class FeatureInput(BaseModel):
    usecase: str = Field(
        ..., description="The ML problem or use case."
    )
    domain_info: Union[str, Dict[str, Any]] = Field(
        ..., description="The business context and domain information."
    )
    schema_name: str = Field(
        ..., description="Name of the database schema."
    )
    table_name: str = Field(
        ..., description="Name of the database table."
    )
    column_metadata: Dict[str, List[Dict[str, Any]]] = Field(
        ...,
        description="A mapping of table names to lists of column metadata dictionaries "
        "(including column_name, description, type, statistical data like min, max, avg, std, unique_count, null_fraction, etc.).",
    )
    target_logic: str = Field(
        ..., description="How the label/target is defined."
    )
    target_column: Optional[List[str]] = Field(
        None, description="The name(s) of the target column(s). Optional for unsupervised ML."
    )
    modeling_approach: str = Field(
        ...,
        description="The type of Machine Learning task "
        "(e.g., regression, binary_classification, ts_classification, clustering, anomaly_detection, etc.).",
    )
    data_source: Optional[str] = Field(
        "snowflake", description="Data store/database type for SQL compatibility (e.g., snowflake, postgres, mysql, bigquery)."
    )
    sql_query_validation: bool = Field(
        False, description="Whether to validate generated SQL queries against the database."
    )
    session: Optional[Any] = Field(
        None, description="Database session for SQL validation (e.g., Snowflake connection)."
    )
    sql_dialect: Optional[str] = Field(
        "snowflake", description="SQL dialect to use (e.g., snowflake, postgres)."
    )

    # Fields exposed to LLM prompts (excludes internal fields like session, sql_dialect, etc.)
    PROMPT_FIELDS: List[str] = Field(
        default=[
            "usecase", "domain_info", "schema_name", "table_name",
            "column_metadata", "target_logic", "target_column", "modeling_approach",
        ],
        exclude=True,
    )

    @classmethod
    def get_input_descriptions(cls, fields: Optional[List[str]] = None) -> str:
        """Generate a formatted string of input field descriptions for LLM system prompts."""
        prompt_fields = fields or [
            "usecase", "domain_info", "schema_name", "table_name",
            "column_metadata", "target_logic", "target_column", "modeling_approach",
        ]
        lines = []
        for name in prompt_fields:
            field_info = cls.model_fields.get(name)
            if field_info:
                desc = field_info.description or name
                lines.append(f"            -{name}: {desc}")
        return "\n".join(lines)

    def get_non_target_columns(self) -> List[str]:
        """Extract non-target column names from column_metadata, supporting both 'name' and 'column_name' keys."""
        target_columns = set(self.target_column) if self.target_column else set()
        all_columns = []
        for cols in self.column_metadata.values():
            for c in cols:
                col_name = c.get("column_name") or c.get("name", "")
                if col_name and col_name not in target_columns:
                    all_columns.append(col_name)
        return all_columns

    def format_prompt_inputs(self) -> str:
        """Format this instance's values as a string block for LLM user prompts."""
        # Format domain_info: extract domain_name and domain_description if dict
        if isinstance(self.domain_info, dict):
            domain_name = self.domain_info.get("domain_name", "")
            domain_desc = self.domain_info.get("domain_description", "")
            domain_block = (
                f"- Domain Name: {domain_name}\n"
                f"        - Domain Description: {domain_desc}"
            )
        else:
            domain_block = f"- Domain Info: {self.domain_info}"

        return (
            f"- Usecase: {self.usecase}\n"
            f"        {domain_block}\n"
            f"        - Schema_name: {self.schema_name}\n"
            f"        - Table_name: {self.table_name}\n"
            f"        - Column_metadata: {self.column_metadata}\n"
            f"        - Target Logic: {self.target_logic}\n"
            f"        - Target Column: {self.target_column}\n"
            f"        - Modeling Approach: {self.modeling_approach}\n"
            f"        - Data Source: {self.data_source}"
        )
