import logging
from collections import defaultdict
from typing import Any, Optional

from semantic_models.agent.autogen_tools import AutogenTools
from semantic_models.agent.generation_state import AutogenContext
from semantic_models.services.llm_utils import extract_json

logger = logging.getLogger(__name__)


class AutogenRunner:
    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]:
        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": []},
        )

        try:
            logger.info("[Runner] Starting autogen pipeline")

            logger.info("[Runner] Step: load_schema")
            ctx = self._load_schema(ctx)
            logger.info("[Runner] load_schema status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at load_schema")
                return self._build_response(ctx)

            logger.info("[Runner] Step: profile_columns")
            ctx = self._profile_columns(ctx)
            logger.info("[Runner] profile_columns status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at profile_columns")
                return self._build_response(ctx)

            logger.info("[Runner] Step: propose_groups")
            ctx = self._propose_groups(ctx)
            logger.info("[Runner] propose_groups status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at propose_groups")
                return self._build_response(ctx)

            logger.info("[Runner] Step: generate_entities")
            ctx = self._generate_entities(ctx)
            logger.info("[Runner] generate_entities status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at generate_entities")
                return self._build_response(ctx)

            logger.info("[Runner] Step: suggest_relationships")
            ctx = self._suggest_relationships(ctx)
            logger.info("[Runner] suggest_relationships status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at suggest_relationships")
                return self._build_response(ctx)

            logger.info("[Runner] Step: validate")
            ctx = self._validate_model(ctx)
            logger.info("[Runner] validate status=%s", ctx.status)
            if ctx.status == "failed":
                logger.info("[Runner] Pipeline failed at validate")
                return self._build_response(ctx)

            logger.info("[Runner] Step: finish")
            ctx = self._finish(ctx)
            logger.info("[Runner] finish status=%s", ctx.status)
            return self._build_response(ctx)

        except Exception as e:
            logger.exception("Autogen runner failed")
            ctx.status = "failed"
            return {
                "status": "failed",
                "model": ctx.working_model,
                "coverage": None,
                "warnings": list(set(ctx.warnings)),
                "error": str(e),
            }

    def _build_response(self, ctx: AutogenContext) -> dict[str, Any]:
        """Build API response from context."""
        tools = AutogenTools(ctx)
        coverage_result = tools.coverage_report()
        coverage = coverage_result.data if coverage_result.success else None

        # Dedupe warnings before returning
        unique_warnings = list(set(ctx.warnings))

        response: dict[str, Any] = {
            "status": ctx.status,
            "model": ctx.working_model,
            "coverage": coverage,
            "warnings": unique_warnings,
        }

        return response

    # =========================================================================
    # Pipeline Steps
    # =========================================================================

    def _load_schema(self, ctx: AutogenContext) -> AutogenContext:
        if ctx.schema_cache:
            logger.info("Schema already loaded, skipping")
            return ctx

        tools = AutogenTools(ctx)
        result = tools.get_dataset_schema()

        if not result.success:
            logger.error(f"Failed to load schema: {result.error}")
            ctx.status = "failed"
            return ctx

        logger.info(f"Loaded schema: {len(result.data.get('columns', []))} columns")
        return ctx

    def _profile_columns(self, ctx: AutogenContext) -> AutogenContext:
        # Check if already profiled
        if ctx.profile_cache and ctx.schema_cache:
            schema_cols = {c["name"] for c in ctx.schema_cache.get("columns", [])}
            cached_cols = set(ctx.profile_cache.keys())
            if schema_cols == cached_cols:
                logger.info("Columns already profiled, skipping")
                return ctx

        from semantic_models.services.dataset_profiling import profile_all_columns

        if not ctx.schema_cache:
            logger.error("Cannot profile columns: schema not loaded")
            ctx.status = "failed"
            return ctx

        try:
            profiles = profile_all_columns(
                connection=ctx.schema_cache["connection"],
                full_table=ctx.schema_cache["full_table"],
                columns=ctx.schema_cache.get("columns", []),
                dialect=ctx.schema_cache.get("dialect", "postgres"),
            )
            ctx.profile_cache = profiles
            logger.info(f"Profiled {len(profiles)} columns")
        except Exception as e:
            # Warn but continue - profiling failure is not fatal
            logger.warning(f"Column profiling failed: {e}")
            ctx.warnings.append(f"Column profiling failed: {e}")

        return ctx

    def _propose_groups(self, ctx: AutogenContext) -> AutogenContext:
        if not ctx.schema_cache:
            logger.error("Cannot propose groups: schema not loaded")
            ctx.status = "failed"
            return ctx

        columns = ctx.schema_cache.get("columns", [])
        logger.info(f"[propose_groups] Starting with {len(columns)} columns")

        # Batch columns to avoid token limits
        BATCH_SIZE = 80
        batches = [columns[i : i + BATCH_SIZE] for i in range(0, len(columns), BATCH_SIZE)]
        logger.info(f"[propose_groups] Split into {len(batches)} batch(es)")

        all_proposals = []
        existing_groups: list[dict[str, Any]] = []
        MAX_RETRIES = 2

        for i, batch in enumerate(batches):
            proposals = []
            for attempt in range(MAX_RETRIES):
                try:
                    logger.info(
                        f"[propose_groups] Processing batch {i + 1}/{len(batches)} ({len(batch)} columns)"
                        + (f" (retry {attempt})" if attempt > 0 else "")
                    )
                    proposals = self._llm_propose_groups_for_batch(ctx, batch, existing_groups)
                    if proposals:
                        logger.info(f"[propose_groups] Batch {i + 1} returned {len(proposals)} groups")
                        break
                    else:
                        logger.warning(f"[propose_groups] Batch {i + 1} returned 0 groups, retrying...")
                except Exception as e:
                    logger.warning(f"[propose_groups] Batch {i + 1} failed: {e}")
                    if attempt == MAX_RETRIES - 1:
                        logger.error(f"[propose_groups] Batch {i + 1} failed after {MAX_RETRIES} attempts")
            all_proposals.extend(proposals)
            existing_groups.extend(proposals)

        # Merge and normalize
        merged = self._merge_group_proposals(all_proposals)

        # Hard failure if no groups
        if not merged:
            logger.error("[propose_groups] LLM returned zero groups after retries")
            ctx.status = "failed"
            return ctx

        # Log unassigned columns
        all_column_names = {c["name"] for c in columns}
        assigned_columns = set()
        for group in merged:
            assigned_columns.update(group.get("columns", []))

        unassigned_columns = sorted(all_column_names - assigned_columns)
        if unassigned_columns:
            logger.info(f"[propose_groups] {len(unassigned_columns)} columns not assigned to any entity")

        ctx.entity_groups = merged
        logger.info(f"Proposed {len(merged)} entity groups from {len(all_column_names)} columns")
        return ctx

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

        # Build prompt with profile stats
        column_lines = []
        for c in columns:
            profile = ctx.profile_cache.get(c["name"], {})
            line = f"- {c['name']} ({c.get('type', 'string')})"
            if c.get("description"):
                line += f": {c['description']}"

            stats = []
            if profile.get("null_count", 0) > 0:
                stats.append("has_nulls")
            else:
                stats.append("no_nulls")
            if profile.get("distinct_count"):
                stats.append(f"distinct:{profile['distinct_count']}")
            if stats:
                line += f" [{', '.join(stats)}]"
            column_lines.append(line)

        columns_text = "\n".join(column_lines)

        existing_groups_text = ""
        if existing_groups:
            group_summaries = [
                f"- {g['entity_name']} ({g['entity_id']}): {g.get('description', '')[:80]}" for g in existing_groups
            ]
            existing_groups_text = f"""
EXISTING ENTITIES (from previous columns - assign to these if columns fit):
{chr(10).join(group_summaries)}

"""

        prompt = f"""Analyze these columns and group them into logical entities.
{existing_groups_text}
COLUMNS:
{columns_text}

{f"USER INSTRUCTIONS: {ctx.user_instructions}" if ctx.user_instructions else ""}

Return JSON with entity groupings.
IMPORTANT: If columns belong to an existing entity listed above, use the SAME entity_id.
Format:
{{
  "groups": [
    {{
      "entity_id": "snake_case_id",
      "entity_name": "Human Name",
      "description": "What this entity represents",
      "columns": ["col1", "col2", ...],
      "primaryKey": ["pk_column"]
    }}
  ]
}}

Guidelines:
- PREFER SINGLE ENTITY: Most tables represent ONE business concept. Only split if columns clearly belong to DIFFERENT concepts.
- Signs you should NOT split: all columns describe the same subject (e.g., a person, an order, a product)
- Signs you SHOULD split: columns from different domains are denormalized together (e.g., order columns + shipping address columns + product columns in one table)
- primaryKey: Identify the column(s) that uniquely identify a row. Look for columns with no_nulls and high distinct count. Usually ends with "_id" or named "id". Can be composite.
- Use business-friendly entity names
- Use snake_case for entity_id
- Identifier columns (IDs, keys, foreign key-like fields) may be shared across entities to enable relationships.
- Regular attribute columns should belong to a single entity.
- It is acceptable for some columns to remain unassigned if they don't fit any logical entity.
"""

        import time

        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(f"[LLM Prompt] Entity grouping:\n{prompt}")

        logger.info(f"[LLM] Calling {llm_id} for entity grouping...")
        start_time = time.time()

        completion = llm.new_completion()
        completion.with_message(prompt)
        response = completion.execute()

        elapsed = time.time() - start_time
        logger.info(f"[LLM] Response received in {elapsed:.1f}s")

        text = response.text.strip()
        logger.info(f"[LLM Response] Entity grouping:\n{text}")

        data = extract_json(text)
        if data:
            return data.get("groups", [])
        else:
            logger.warning("Failed to extract JSON from LLM response")
            return []

    def _merge_group_proposals(self, proposals: list[dict[str, Any]]) -> list[dict[str, Any]]:
        name_map: dict[str, list[dict]] = defaultdict(list)

        for p in proposals:
            if not p.get("entity_name"):
                continue
            normalized = p["entity_name"].lower().replace(" ", "_").replace("-", "_")
            name_map[normalized].append(p)

        merged = []
        for name, group_proposals in name_map.items():
            all_columns = []
            for gp in group_proposals:
                all_columns.extend(gp.get("columns", []))

            merged.append(
                {
                    "entity_id": name,
                    "entity_name": group_proposals[0].get("entity_name", name),
                    "description": group_proposals[0].get("description"),
                    "columns": list(set(all_columns)),
                    "primaryKey": group_proposals[0].get("primaryKey", []),
                }
            )

        return merged

    def _generate_entities(self, ctx: AutogenContext) -> AutogenContext:
        logger.info(f"[generate_entities] Starting for {len(ctx.entity_groups)} groups")
        tools = AutogenTools(ctx)

        # Reset working model
        ctx.working_model = {"entities": [], "relationships": []}

        for i, group in enumerate(ctx.entity_groups):
            entity_name = group.get("entity_name", group.get("entity_id"))
            logger.info(f"[generate_entities] Creating entity {i + 1}/{len(ctx.entity_groups)}: {entity_name}")

            tools.create_entity(group)

            logger.info(f"[generate_entities] Generating metrics/filters for {entity_name}...")
            metrics, filters = self._llm_generate_metrics_filters(ctx, group)

            if metrics:
                tools.add_metrics(group["entity_id"], metrics)
                logger.info(f"[generate_entities] Added {len(metrics)} metrics to {entity_name}")
            if filters:
                tools.add_filters(group["entity_id"], filters)
                logger.info(f"[generate_entities] Added {len(filters)} filters to {entity_name}")

        logger.info(f"[generate_entities] Completed: {len(ctx.entity_groups)} entities created")
        return ctx

    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 - it's added automatically
- Expression should be ONLY the aggregate: "COUNT(order_id)" NOT "COUNT(order_id) GROUP BY x"
- 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
- GOOD filters: "is_active = true", "deleted_at IS NULL", "status = 'active'", "created_at >= CURRENT_DATE - INTERVAL '30 days'"
- BAD filters: "country_code = 'US'" or "region = 'East'" (arbitrary hardcoded values)
- Only create 0-2 filters if there are obvious boolean/status/date columns
- If no obvious filter patterns exist, return empty filters array
"""

        try:
            import time

            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(f"[LLM Prompt] Metrics/filters for {group.get('entity_name')}:\n{prompt}")

            start_time = time.time()
            completion = llm.new_completion()
            completion.with_message(prompt)
            response = completion.execute()
            elapsed = time.time() - start_time
            logger.info(f"[LLM] Metrics/filters response in {elapsed:.1f}s")

            text = response.text.strip()
            logger.info(f"[LLM Response] Metrics/filters:\n{text}")

            data = extract_json(text)
            if not data:
                logger.warning(f"Failed to extract JSON from metrics/filters response")
                return [], []

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

            # Post-process: strip GROUP BY
            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(f"Failed to generate metrics/filters for {group.get('entity_name')}: {e}")
            return [], []

    def _suggest_relationships(self, ctx: AutogenContext) -> AutogenContext:
        """Use LLM to suggest relationships between entities."""
        tools = AutogenTools(ctx)

        # Need at least 2 entities for relationships
        if len(ctx.entity_groups) < 2:
            logger.info("Less than 2 entities, skipping relationship detection")
            return ctx

        # Use LLM to suggest relationships
        relationships = self._llm_suggest_relationships(ctx)

        # Track added relationships for deduplication
        existing_relationships = set()
        if ctx.working_model:
            for rel in ctx.working_model.get("relationships", []):
                key = (rel.get("left_entity"), rel.get("right_entity"))
                existing_relationships.add(key)
                # Also add reverse to prevent duplicates
                existing_relationships.add((rel.get("right_entity"), rel.get("left_entity")))

        # Add relationships
        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

            # Dedupe
            rel_key = (left_entity, right_entity)
            if rel_key in existing_relationships:
                continue
            existing_relationships.add(rel_key)
            existing_relationships.add((right_entity, left_entity))

            tools.add_relationship(
                left_entity=left_entity,
                right_entity=right_entity,
                left_column=left_column,
                right_column=right_column,
            )

        rel_count = len(ctx.working_model.get("relationships", [])) if ctx.working_model else 0
        logger.info(f"Added {rel_count} relationships")
        return ctx

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

        import dataiku

        # Build entity summaries with PKs and columns
        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", [])
            # Show columns that look like FKs (end with _id, _key, _code)
            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['entity_name']} (id: {group['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 ''}"
            )

        entities_text = "\n".join(entity_summaries)

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

ENTITIES:
{entities_text}

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: columns ending with _ID, _KEY, _CODE, _REF that reference another entity's primary key
- The left_entity contains the FK, right_entity contains the PK being referenced
- Column names are CASE-SENSITIVE - use exact column names as shown
- 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(f"[LLM Prompt] Relationships:\n{prompt}")
            logger.info(f"[LLM] Calling {llm_id} for relationship suggestions...")
            start_time = time.time()

            completion = llm.new_completion()
            completion.with_message(prompt)
            response = completion.execute()

            elapsed = time.time() - start_time
            logger.info(f"[LLM] Relationship response in {elapsed:.1f}s")

            text = response.text.strip()
            logger.info(f"[LLM Response] Relationships:\n{text}")

            data = extract_json(text)
            if data:
                relationships = data.get("relationships", [])
                logger.info(f"[LLM] Suggested {len(relationships)} relationships")
                return relationships
            else:
                logger.warning("Failed to extract JSON from relationship response")
                return []

        except Exception as e:
            logger.warning(f"Failed to get LLM relationship suggestions: {e}")
            return []

    def _validate_model(self, ctx: AutogenContext) -> AutogenContext:
        tools = AutogenTools(ctx)
        result = tools.validate_model()

        if result.success:
            logger.info("Validation passed")
            return ctx

        ctx.validation_attempts += 1
        errors = result.data.get("errors", []) if result.data else []
        logger.info(f"Validation failed (attempt {ctx.validation_attempts}): {errors}")

        # If under attempt limit and has metric errors, try repair
        if ctx.validation_attempts < 3:
            has_metric_error = any("metric" in e.lower() for e in errors)
            if has_metric_error:
                logger.info("Attempting metric repair...")
                ctx = self._repair_metrics(ctx)
                # Re-validate after repair
                return self._validate_model(ctx)

        # After 3 attempts, add errors to warnings and continue (non-failing finish)
        if ctx.validation_attempts >= 3:
            for error in errors:
                ctx.warnings.append(f"Validation error (not fixed): {error}")
            logger.info("Validation failed after 3 attempts, continuing with warnings")

        return ctx

    def _repair_metrics(self, ctx: AutogenContext) -> AutogenContext:
        tools = AutogenTools(ctx)

        if not ctx.working_model:
            return 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["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)

        logger.info("Repaired metrics for all entities")
        return ctx

    def _finish(self, ctx: AutogenContext) -> AutogenContext:
        ctx.status = "completed"
        logger.info("Autogen completed successfully")
        return ctx
