import json
import logging
import re
from dataclasses import dataclass
from typing import Annotated, Any, Literal, Optional, TypedDict

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.tools import tool
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages

from semantic_models.agent.generation_state import AutogenContext
from semantic_models.agent.autogen_tools import AutogenTools
from semantic_models.agentic_generation.prompts import (
    get_attribute_enrichment_prompt,
    get_system_prompt,
)
from semantic_models.agentic_generation.tools import AgenticTools
from semantic_models.services.llm_utils import extract_json

logger = logging.getLogger(__name__)

MAX_DATASET_DESC_LEN = 800
MAX_COL_DESC_LEN = 400
MAX_SAMPLE_VALUE_LENGTH = 100
MAX_SAMPLE_VALUES = 10
ATTR_ENRICH_BATCH_SIZE = 50
LOW_CARDINALITY_THRESHOLD = 50


@dataclass
class AgenticRuntimeContext:
    project_key: str
    dataset_name: str
    llm_id: str
    user_instructions: Optional[str]
    trace: Any = None


class AgenticState(TypedDict):
    ctx: AgenticRuntimeContext
    autogen_ctx: AutogenContext
    messages: Annotated[list[BaseMessage], add_messages]
    tool_call_count: int
    sampled_columns: list[str]
    stats_calls: int
    budgets: dict[str, int]
    final_groups: Optional[list[dict[str, Any]]]


def create_tools(state: AgenticState):
    agentic_tools = AgenticTools(state["autogen_ctx"])

    @tool
    def get_column_samples(column_name: str, limit: int = 10) -> str:
        """Get distinct sample values for a column."""
        result = agentic_tools.get_column_samples(column_name, limit=limit)
        if result.success:
            return json.dumps(result.data, indent=2)
        return json.dumps({"error": result.error}, indent=2)

    @tool
    def get_column_stats(column_name: str) -> str:
        """Get null_count and distinct_count for a column."""
        result = agentic_tools.get_column_stats(column_name)
        if result.success:
            return json.dumps(result.data, indent=2)
        return json.dumps({"error": result.error}, indent=2)

    return [get_column_samples, get_column_stats]


def call_model(state: AgenticState) -> dict:
    from dataiku.langchain.dku_llm import DKUChatModel

    llm = DKUChatModel(llm_id=state["ctx"].llm_id)
    tools = create_tools(state)
    model_with_tools = llm.bind_tools(tools)
    response = model_with_tools.invoke(state["messages"])
    return {"messages": [response]}


def execute_tools(state: AgenticState) -> dict:
    last_message = state["messages"][-1]
    tool_calls = last_message.tool_calls

    tools = create_tools(state)
    tools_by_name = {t.name: t for t in tools}

    tool_messages = []
    for tool_call in tool_calls:
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]

        state["tool_call_count"] += 1
        if state["tool_call_count"] > state["budgets"]["max_tool_calls"]:
            tool_messages.append(
                ToolMessage(
                    content=json.dumps({"error": "Tool call budget exceeded"}),
                    name=tool_name,
                    tool_call_id=tool_call["id"],
                )
            )
            continue

        if tool_name == "get_column_samples":
            col = tool_args.get("column_name")
            if col and col not in state["sampled_columns"]:
                state["sampled_columns"].append(col)
            if len(state["sampled_columns"]) > state["budgets"]["max_sampled_columns"]:
                tool_messages.append(
                    ToolMessage(
                        content=json.dumps({"error": "Sampled columns budget exceeded"}),
                        name=tool_name,
                        tool_call_id=tool_call["id"],
                    )
                )
                continue

        if tool_name == "get_column_stats":
            state["stats_calls"] += 1
            if state["stats_calls"] > state["budgets"]["max_stats_calls"]:
                tool_messages.append(
                    ToolMessage(
                        content=json.dumps({"error": "Stats calls budget exceeded"}),
                        name=tool_name,
                        tool_call_id=tool_call["id"],
                    )
                )
                continue

        if tool_name in tools_by_name:
            try:
                result = tools_by_name[tool_name].invoke(tool_args)
            except Exception as e:
                result = json.dumps({"error": str(e)})
        else:
            result = json.dumps({"error": f"Unknown tool: {tool_name}"})

        tool_messages.append(
            ToolMessage(
                content=result,
                name=tool_name,
                tool_call_id=tool_call["id"],
            )
        )

    return {"messages": tool_messages}


def should_continue(state: AgenticState) -> Literal["tools", "end"]:
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"
    return "end"


def build_agent_graph() -> StateGraph:
    graph = StateGraph(AgenticState)
    graph.add_node("agent", call_model)
    graph.add_node("tools", execute_tools)
    graph.set_entry_point("agent")
    graph.add_conditional_edges("agent", should_continue, {"tools": "tools", "end": END})
    graph.add_edge("tools", "agent")
    return graph.compile()


class AgenticRunner:
    def __init__(self, llm_id: str):
        self.llm_id = llm_id

    def generate(
        self,
        project_key: str,
        dataset_name: str,
        llm_id: Optional[str] = None,
        user_instructions: Optional[str] = None,
    ) -> dict[str, Any]:
        autogen_ctx = AutogenContext(
            project_key=project_key,
            dataset_name=dataset_name,
            llm_id=llm_id or self.llm_id,
            user_instructions=user_instructions,
            working_model={"entities": [], "relationships": []},
        )
        tools = AgenticTools(autogen_ctx)

        bootstrap = tools.get_dataset_profile(
            project_key=project_key,
            dataset_name=dataset_name,
            sample_value_limit=0,
        )
        if not bootstrap.success:
            autogen_ctx.status = "failed"
            return self._build_response(autogen_ctx, error=bootstrap.error)

        dataset_context = self._sanitize_dataset_context(bootstrap.data)
        autogen_ctx.dataset_profile = dataset_context.get("dataset", {})

        budgets = {
            "max_tool_calls": 12,
            "max_sampled_columns": 25,
            "max_stats_calls": 15,
        }

        question = f"Generate entities for dataset {project_key}.{dataset_name}."
        if user_instructions:
            question += f" Instructions: {user_instructions}"

        initial_state: AgenticState = {
            "ctx": AgenticRuntimeContext(
                project_key=project_key,
                dataset_name=dataset_name,
                llm_id=autogen_ctx.llm_id or self.llm_id,
                user_instructions=user_instructions,
            ),
            "autogen_ctx": autogen_ctx,
            "messages": [
                SystemMessage(content=get_system_prompt()),
                ToolMessage(
                    content=json.dumps(dataset_context),
                    name="get_dataset_profile",
                    tool_call_id="bootstrap",
                ),
                HumanMessage(content=question),
            ],
            "tool_call_count": 0,
            "sampled_columns": [],
            "stats_calls": 0,
            "budgets": budgets,
            "final_groups": None,
        }

        try:
            agent = build_agent_graph()
            final_state = agent.invoke(initial_state, config={"recursion_limit": 30})
            groups = self._extract_groups(final_state["messages"])
            if not groups:
                autogen_ctx.status = "failed"
                return self._build_response(autogen_ctx, error="Agentic generation returned no groups")

            autogen_ctx.entity_groups = groups
            self._generate_entities(autogen_ctx)
            self._suggest_relationships(autogen_ctx)
            self._validate_model(autogen_ctx)

            autogen_ctx.status = "completed"
            return self._build_response(autogen_ctx)
        except Exception as e:
            logger.exception("Agentic runner failed")
            autogen_ctx.status = "failed"
            return self._build_response(autogen_ctx, error=str(e))

    def _extract_groups(self, messages: list[BaseMessage]) -> list[dict[str, Any]]:
        for msg in reversed(messages):
            if isinstance(msg, AIMessage):
                data = extract_json(msg.content)
                if data and isinstance(data.get("groups"), list):
                    return data["groups"]
        return []

    def _truncate_text(self, value: Optional[str], max_len: int) -> Optional[str]:
        if not value:
            return value
        if len(value) > max_len:
            return value[: max_len - 3] + "..."
        return value

    def _sanitize_dataset_context(self, data: dict[str, Any]) -> dict[str, Any]:
        dataset = data.get("dataset", {}) or {}
        columns = data.get("columns", []) or []

        sanitized_dataset = {
            "name": dataset.get("name"),
            "label": self._truncate_text(dataset.get("label"), MAX_DATASET_DESC_LEN),
            "description": self._truncate_text(dataset.get("description"), MAX_DATASET_DESC_LEN),
        }

        sanitized_columns = []
        for col in columns:
            sample_values = col.get("sample_values") or []
            if sample_values:
                sample_values = [
                    (v[: MAX_SAMPLE_VALUE_LENGTH - 3] + "...") if isinstance(v, str) and len(v) > MAX_SAMPLE_VALUE_LENGTH else v
                    for v in sample_values
                ][:MAX_SAMPLE_VALUES]
            sanitized_columns.append(
                {
                    "name": col.get("name"),
                    "type": col.get("type"),
                    "description": self._truncate_text(col.get("description"), MAX_COL_DESC_LEN),
                    "sample_values": sample_values if sample_values else None,
                }
            )

        return {"dataset": sanitized_dataset, "columns": sanitized_columns}

    def _build_response(self, ctx: AutogenContext, error: Optional[str] = None) -> dict[str, Any]:
        tools = AutogenTools(ctx)
        coverage_result = tools.coverage_report()
        coverage = coverage_result.data if coverage_result.success else None
        response: dict[str, Any] = {
            "status": ctx.status,
            "model": ctx.working_model,
            "coverage": coverage,
            "warnings": list(set(ctx.warnings)),
        }
        if error:
            response["error"] = error
        return response

    def _generate_entities(self, ctx: AutogenContext) -> None:
        tools = AutogenTools(ctx)
        ctx.working_model = {"entities": [], "relationships": []}
        for group in ctx.entity_groups:
            entity_id = group.get("entity_id") or group.get("entity_name")
            if not entity_id:
                continue
            if "entity_id" not in group:
                group["entity_id"] = self._normalize_entity_id(entity_id)
            tools.create_entity(group)
            self._enrich_attributes_for_group(ctx, group)

            metrics, filters = self._llm_generate_metrics_filters(ctx, group)
            if metrics:
                tools.add_metrics(group["entity_id"], metrics)
            if filters:
                tools.add_filters(group["entity_id"], filters)

    def _normalize_entity_id(self, name: str) -> str:
        return re.sub(r"[^a-z0-9_]+", "_", name.strip().lower()).strip("_")

    def _chunk_list(self, items: list[Any], size: int) -> list[list[Any]]:
        return [items[i : i + size] for i in range(0, len(items), size)]

    def _enrich_attributes_for_group(self, ctx: AutogenContext, group: dict[str, Any]) -> None:
        entity_id = group.get("entity_id")
        entity_name = group.get("entity_name", entity_id)
        if not entity_id or not entity_name:
            return

        entity = next((e for e in ctx.working_model.get("entities", []) if e.get("name") == entity_id), None)
        if not entity:
            return

        columns = group.get("columns", []) or []
        if not columns:
            return

        dataset_info = getattr(ctx, "dataset_profile", None) or {}
        dataset_info = {
            "label": self._truncate_text(dataset_info.get("label"), MAX_DATASET_DESC_LEN),
            "description": self._truncate_text(dataset_info.get("description"), MAX_DATASET_DESC_LEN),
        }

        # Build per-column context with descriptions + sample values if available
        column_context = []
        schema_cols = {c.get("name"): c for c in (ctx.schema_cache or {}).get("columns", [])}
        for col_name in columns:
            schema = schema_cols.get(col_name, {})
            profile = ctx.profile_cache.get(col_name, {})
            sample_values = profile.get("sample_values") or []
            if sample_values:
                sample_values = [
                    (v[: MAX_SAMPLE_VALUE_LENGTH - 3] + "...")
                    if isinstance(v, str) and len(v) > MAX_SAMPLE_VALUE_LENGTH
                    else v
                    for v in sample_values
                ][:MAX_SAMPLE_VALUES]
            column_context.append(
                {
                    "name": col_name,
                    "type": schema.get("type"),
                    "description": schema.get("description"),
                    "sample_values": sample_values if sample_values else None,
                }
            )

        # Batch enrichment to avoid huge prompts
        for batch in self._chunk_list(column_context, ATTR_ENRICH_BATCH_SIZE):
            columns_block = "\n".join(
                [
                    f"- {c['name']} ({c.get('type','string')}): {self._truncate_text(c.get('description'), MAX_COL_DESC_LEN) or 'N/A'}"
                    + (f" | samples: {c.get('sample_values')}" if c.get("sample_values") else "")
                    for c in batch
                ]
            )

            prompt = get_attribute_enrichment_prompt(
                dataset_label=dataset_info.get("label"),
                dataset_description=dataset_info.get("description"),
                entity_name=entity_name,
                entity_description=group.get("description"),
                columns_block=columns_block,
            )

            try:
                data = self._call_llm_json(ctx, prompt)
                attrs = data.get("attributes", []) if data else []
                self._apply_attribute_enrichment(ctx, entity, attrs)
            except Exception as e:
                logger.warning("Attribute enrichment failed for %s: %s", entity_name, e)

        # Apply deterministic rules for resolvable + distinct values modes
        self._apply_attribute_rules(ctx, entity)

    def _call_llm_json(self, ctx: AutogenContext, prompt: str) -> Optional[dict[str, Any]]:
        import dataiku

        client = dataiku.api_client()
        project = client.get_default_project()
        llm_id = ctx.llm_id or self.llm_id
        llm = project.get_llm(llm_id)

        logger.info("[LLM Prompt] Attribute enrichment:\n%s", prompt)
        completion = llm.new_completion()
        completion.with_message(prompt)
        response = completion.execute()
        raw = response.text.strip()
        logger.info("[LLM Response] Attribute enrichment:\n%s", raw)
        return extract_json(raw)

    def _apply_attribute_enrichment(self, ctx: AutogenContext, entity: dict[str, Any], attrs: list[dict]) -> None:
        attr_map = {a.get("name"): a for a in entity.get("attributes", []) if a.get("name")}
        for attr in attrs:
            name = attr.get("name")
            if not name or name not in attr_map:
                continue
            if attr.get("description"):
                attr_map[name]["description"] = attr.get("description")
            if isinstance(attr.get("resolvable"), bool):
                attr_map[name]["resolveInUserRequests"] = attr.get("resolvable")
            if attr.get("preferredDistinctValuesHandlingMode"):
                attr_map[name]["_preferredDistinctValuesHandlingMode"] = attr.get(
                    "preferredDistinctValuesHandlingMode"
                )

    def _apply_attribute_rules(self, ctx: AutogenContext, entity: dict[str, Any]) -> None:
        from semantic_models.services.dss_utils import is_categorical_type

        for attr in entity.get("attributes", []):
            col_name = attr.get("name")
            if not col_name:
                continue
            profile = ctx.profile_cache.get(col_name, {})
            distinct_count = profile.get("distinct_count")
            sample_values = profile.get("sample_values") or []
            dss_type = attr.get("dssType")

            is_categorical = False
            if dss_type:
                if distinct_count is not None:
                    is_categorical = is_categorical_type(dss_type) and distinct_count <= LOW_CARDINALITY_THRESHOLD
                elif sample_values:
                    # Fallback when stats are unavailable but profiling returned sample values.
                    is_categorical = is_categorical_type(dss_type)

            if is_categorical:
                attr["manualValues"] = []
                attr["distinctValuesHandlingMode"] = "AUTO_INDEX"
                attr["indexDistinctValues"] = True
                attr["resolveInUserRequests"] = True
            else:
                attr["manualValues"] = []
                attr["distinctValuesHandlingMode"] = "NONE"
                attr["indexDistinctValues"] = False
                if attr.get("resolveInUserRequests") is True:
                    attr["resolveInUserRequests"] = False

    def _llm_generate_metrics_filters(
        self, ctx: AutogenContext, group: dict[str, Any]
    ) -> tuple[list[dict], list[dict]]:
        import dataiku

        columns = group.get("columns", [])
        if not columns:
            return [], []

        column_info = []
        for col_name in columns:
            col_type = "string"
            if ctx.schema_cache:
                for c in ctx.schema_cache.get("columns", []):
                    if c["name"] == col_name:
                        col_type = c.get("type", "string")
                        break
            column_info.append(f"- {col_name} ({col_type})")

        columns_text = "\n".join(column_info)
        entity_name = group.get("entity_name", group.get("entity_id"))
        entity_desc = group.get("description", "")

        prompt = f"""Generate metrics and filters for the "{entity_name}" entity.

DESCRIPTION: {entity_desc if entity_desc else "No description provided"}

COLUMNS:
{columns_text}

Return JSON with metrics and filters:
{{
  "metrics": [
    {{
      "name": "Metric Name",
      "description": "What this metric measures",
      "expression": "SUM(column_name)"
    }}
  ],
  "filters": [
    {{
      "name": "Filter Name",
      "description": "What this filter does",
      "expression": "is_active = true"
    }}
  ]
}}

Guidelines for METRICS:
- MUST use aggregate functions (SUM, COUNT, AVG, MIN, MAX)
- NEVER include GROUP BY
- Only create metrics for numeric columns or COUNT for any column
- Create 2-5 useful metrics

Guidelines for FILTERS:
- Create common business filters, NOT arbitrary value filters
- Only create 0-2 filters if there are obvious boolean/status/date columns
- If no obvious filter patterns exist, return empty filters array
"""

        try:
            client = dataiku.api_client()
            project = client.get_default_project()
            llm_id = ctx.llm_id or self.llm_id
            llm = project.get_llm(llm_id)

            logger.info("[LLM Prompt] Metrics/filters for %s:\n%s", entity_name, prompt)
            completion = llm.new_completion()
            completion.with_message(prompt)
            response = completion.execute()

            raw = response.text.strip()
            logger.info("[LLM Response] Metrics/filters for %s:\n%s", entity_name, raw)
            data = extract_json(raw)
            if not data:
                return [], []

            metrics = data.get("metrics", [])
            filters = data.get("filters", [])

            for m in metrics:
                if "expression" in m:
                    m["expression"] = re.sub(
                        r"\s+GROUP\s+BY\s+.*$", "", m["expression"], flags=re.IGNORECASE
                    ).strip()

            return metrics, filters
        except Exception as e:
            logger.warning("Failed to generate metrics/filters for %s: %s", group.get("entity_name"), e)
            return [], []

    def _suggest_relationships(self, ctx: AutogenContext) -> None:
        if len(ctx.entity_groups) < 2:
            return
        relationships = self._llm_suggest_relationships(ctx)
        tools = AutogenTools(ctx)
        existing = set()
        for rel in ctx.working_model.get("relationships", []):
            key = (rel.get("firstEntity"), rel.get("secondEntity"))
            existing.add(key)
            existing.add((key[1], key[0]))
        for rel in relationships:
            left_entity = rel.get("left_entity")
            right_entity = rel.get("right_entity")
            left_column = rel.get("left_column")
            right_column = rel.get("right_column")
            if not all([left_entity, right_entity, left_column, right_column]):
                continue
            rel_key = (left_entity, right_entity)
            if rel_key in existing:
                continue
            tools.add_relationship(left_entity, right_entity, left_column, right_column)
            existing.add(rel_key)
            existing.add((right_entity, left_entity))

    def _llm_suggest_relationships(self, ctx: AutogenContext) -> list[dict[str, Any]]:
        import dataiku

        entity_summaries = []
        for group in ctx.entity_groups:
            pk_list = group.get("primaryKey", [])
            pk_str = ", ".join(pk_list) if pk_list else "none identified"
            cols = group.get("columns", [])
            potential_fks = [
                c for c in cols if any(c.lower().endswith(suffix) for suffix in ["_id", "_key", "_code", "_ref"])
            ]
            fk_str = ", ".join(potential_fks[:10]) if potential_fks else "none"
            entity_summaries.append(
                f"- {group.get('entity_name')} (id: {group.get('entity_id')})\n"
                f"  Primary Key: {pk_str}\n"
                f"  Potential FK columns: {fk_str}\n"
                f"  All columns: {', '.join(cols[:15])}{' ...' if len(cols) > 15 else ''}"
            )

        prompt = f"""Analyze these entities and identify relationships between them.

ENTITIES:
{chr(10).join(entity_summaries)}

Return JSON with relationships:
{{
  "relationships": [
    {{
      "left_entity": "entity_id of the source entity (the one with the FK)",
      "right_entity": "entity_id of the target entity (the one being referenced)",
      "left_column": "FK column name in left entity",
      "right_column": "PK column name in right entity",
      "reason": "brief explanation"
    }}
  ]
}}

Guidelines:
- Look for foreign key patterns
- Only include high-confidence relationships
- If no clear relationships exist, return empty array
"""

        try:
            client = dataiku.api_client()
            project = client.get_default_project()
            llm_id = ctx.llm_id or self.llm_id
            llm = project.get_llm(llm_id)

            logger.info("[LLM Prompt] Relationships:\n%s", prompt)
            completion = llm.new_completion()
            completion.with_message(prompt)
            response = completion.execute()

            raw = response.text.strip()
            logger.info("[LLM Response] Relationships:\n%s", raw)
            data = extract_json(raw)
            return data.get("relationships", []) if data else []
        except Exception as e:
            logger.warning("Failed to get relationship suggestions: %s", e)
            return []

    def _validate_model(self, ctx: AutogenContext) -> None:
        tools = AutogenTools(ctx)
        result = tools.validate_model()
        if result.success:
            return
        errors = result.data.get("errors", []) if result.data else []
        has_metric_error = any("metric" in e.lower() for e in errors)
        if has_metric_error:
            self._repair_metrics(ctx)

    def _repair_metrics(self, ctx: AutogenContext) -> None:
        tools = AutogenTools(ctx)
        for entity in ctx.working_model.get("entities", []):
            entity_id = entity.get("name")
            entity["metrics"] = []
            entity["filters"] = []
            group = next((g for g in ctx.entity_groups if g.get("entity_id") == entity_id), None)
            if group:
                metrics, filters = self._llm_generate_metrics_filters(ctx, group)
                if metrics:
                    tools.add_metrics(entity_id, metrics)
                if filters:
                    tools.add_filters(entity_id, filters)
