"""Decomposition orchestration — turns a goal into a confirmed task graph."""

from __future__ import annotations

import asyncio
import json
import logging
import os
from pathlib import Path

import asyncpg
from pydantic import BaseModel, Field
from redis.asyncio import Redis

from loom.bus.events import EventType
from loom.complexity import COMPLEXITY_THRESHOLD
from loom.bus.publisher import publish_event
from loom.config import SkillsConfig
from loom.exceptions import DecompositionError
from loom.graph import cache, store
from loom.graph.task import Priority, Task, TaskStatus
from loom.ids import task_id as gen_task_id
from loom.scanner import CodebaseSummary, scan_codebase
from loom.skills.loader import load_skill
from loom.skills.runner import run_skill
from loom.validator import (
    ValidationStatus,
    classify_match,
    match_task_to_codebase,
    merge_by_file_affinity,
    merge_trivial_tasks,
    validate_task_graph,
)

log = logging.getLogger(__name__)

# ── Codebase context prompt template ─────────────────────────────────
# Sentinel marker for testability — search for this string in rendered prompts.
CODEBASE_CONTEXT_SENTINEL = "# --- CODEBASE CONTEXT ---"

# Default token budget for the codebase context section.
CODEBASE_CONTEXT_TOKEN_BUDGET = 2000

# Maximum characters for the rendered codebase context section.
# If the serialised summary exceeds this, it is truncated with a marker.
MAX_SCAN_CONTEXT_CHARS = 4000

CODEBASE_CONTEXT_TEMPLATE = """\
{sentinel}

The following symbols already exist in the codebase. Use this information
to avoid redundant work and produce a more accurate task graph.

RULES — read these carefully:
- DO NOT generate tasks whose sole purpose is to create files, classes,
  or functions that already appear below. They are already implemented.
- If an existing symbol is only a stub or partial implementation, treat it
  as partially complete. Generate a task to *finish* it, not to create it
  from scratch.
- Fold trivially small tasks (single-line changes, renaming, import adds)
  into the nearest parent task instead of listing them separately.
- When a task modifies an existing file, reference the file path in the
  task context so agents can find it quickly.

{summary}
"""


def build_codebase_context(
    summary: CodebaseSummary | None,
    max_tokens: int = CODEBASE_CONTEXT_TOKEN_BUDGET,
    max_chars: int = MAX_SCAN_CONTEXT_CHARS,
) -> str:
    """Render a codebase-context prompt section from a scanner summary.

    Returns an empty string when the summary is None or contains no files,
    so callers can unconditionally concatenate the result.

    The rendered text is hard-truncated to *max_chars* characters so that
    unexpectedly large codebases cannot blow out the prompt budget.
    """
    if summary is None or not summary.files:
        return ""

    rendered_summary = summary.to_prompt_str(max_tokens=max_tokens)
    result = CODEBASE_CONTEXT_TEMPLATE.format(
        sentinel=CODEBASE_CONTEXT_SENTINEL,
        summary=rendered_summary,
    )

    if len(result) > max_chars:
        truncation_marker = "\n[... truncated to fit context budget ...]"
        result = result[: max_chars - len(truncation_marker)] + truncation_marker

    return result


def _auto_scan(scan_root: Path | None) -> CodebaseSummary | None:
    """Run the codebase scanner, returning None on any failure.

    Respects the ``LOOM_SKIP_SCAN`` environment variable — when set to a
    truthy value (``1``, ``true``, ``yes``), scanning is skipped entirely.
    """
    if os.environ.get("LOOM_SKIP_SCAN", "").lower() in ("1", "true", "yes"):
        log.info("LOOM_SKIP_SCAN is set — skipping codebase scan")
        return None

    if scan_root is None:
        return None

    try:
        summary = scan_codebase(scan_root)
        log.info(
            "Scanned %d files in %s",
            len(summary.files),
            summary.root,
        )
        return summary
    except Exception:
        log.warning("Codebase scan failed — continuing without context", exc_info=True)
        return None


class TaskDefinition(BaseModel):
    """Raw LLM output for a single task — uses title-based deps."""

    title: str
    type: str = "task"  # "epic" or "task"
    parent_epic: str | None = None
    depends_on_titles: list[str] = Field(default_factory=list)
    priority: str = "p1"
    context: dict = Field(default_factory=dict)
    done_when: str | None = None


class DecompositionResult(BaseModel):
    """Result of a decomposition run."""

    tasks: list[dict] = Field(default_factory=list)
    epics: list[str] = Field(default_factory=list)
    total_count: int = 0
    dependency_edges: int = 0
    written_to_db: bool = False
    validation_summary: dict | None = None


def _validate_tasks(
    task_defs: list[TaskDefinition],
    summary: CodebaseSummary,
) -> tuple[list[TaskDefinition], dict]:
    """Run the validator against each task and annotate its context.

    Returns the (possibly mutated) task_defs and a summary dict with counts
    per ValidationStatus.
    """
    counts: dict[str, int] = {"pending": 0, "uncertain": 0, "pre_completed": 0}

    for td in task_defs:
        if td.type == "epic":
            continue

        task_dict = {"title": td.title, "context": td.context}
        result = match_task_to_codebase(task_dict, summary)
        status = classify_match(result)

        td.context["validation_status"] = status.value
        td.context["matched_artifact"] = result.matched_artifact
        td.context["match_rationale"] = result.rationale

        counts[status.value] += 1

    return task_defs, counts


async def decompose(
    goal: str,
    project_id: str,
    config: SkillsConfig,
    pool: asyncpg.Pool,
    redis: Redis,
    context: str | None = None,
    confirm: bool = True,
    project_dir: str | None = None,
    depth: int = 3,
    enrich: bool = True,
    codebase_summary: CodebaseSummary | None = None,
    scan_root: Path | None = None,
    min_complexity: float | None = None,
    merge_trivial: bool = True,
    optimize_parallelism: bool = False,
    enrichment_concurrency: int | None = None,
    parent_epic_id: str | None = None,
    force_leaf: bool = False,
    parent_epic_deps: list[str] | None = None,
    auto_fix_paths: bool = True,
    merge_file_affinity: bool = True,
) -> DecompositionResult:
    """Decompose a goal into a task graph.

    If confirm=True, returns proposed tasks without writing to DB.
    If confirm=False, writes tasks to DB and publishes events.

    When *codebase_summary* is provided, its symbols are injected into
    the prompt so the LLM avoids generating tasks for existing code.

    When *scan_root* is provided (and no *codebase_summary* was given),
    the codebase scanner runs automatically to produce the summary.
    Set ``LOOM_SKIP_SCAN=1`` to disable auto-scanning globally.

    Parameters
    ----------
    goal:
        The high-level project goal to decompose into tasks.
    project_id:
        UUID of the Loom project that will own the generated tasks.
    config:
        ``SkillsConfig`` with API key and model settings for LLM calls.
    pool:
        asyncpg connection pool for Postgres operations.
    redis:
        Redis connection for cache sync and event publishing.
    context:
        Optional free-text context appended to the LLM prompt.
    confirm:
        When ``True`` (default), the result is returned without writing
        to the database, allowing the caller to inspect the proposed
        graph first.  Set to ``False`` to write immediately.
    project_dir:
        Path to the project directory, used for skill and workflow
        discovery.
    depth:
        Maximum hierarchy depth for the decomposed task tree
        (default: 3).
    enrich:
        When ``True`` (default), each non-epic task is enriched with
        context, acceptance criteria, and complexity estimates via
        additional LLM calls.  Set to ``False`` to skip enrichment.
    codebase_summary:
        Pre-built ``CodebaseSummary`` to inject into the prompt and
        use for post-decompose validation.  When provided, *scan_root*
        is ignored.
    scan_root:
        Directory to scan for existing code.  Triggers the codebase
        scanner automatically when *codebase_summary* is not provided.
        The resulting summary is used both for prompt injection and
        post-decompose validation/pruning.
    min_complexity:
        Override the default ``COMPLEXITY_THRESHOLD`` for the trivial-task
        merge step.  Must be >= 1.  When ``None``, the default from
        ``loom.complexity.COMPLEXITY_THRESHOLD`` is used.
    merge_trivial:
        When ``True`` (default), trivial tasks are merged into their
        parents or dependents after graph validation.  Set to ``False``
        to skip the merge step entirely.
    """
    # Step 0a: Auto-scan if caller provided scan_root but not an explicit summary
    if codebase_summary is None and scan_root is not None:
        codebase_summary = _auto_scan(scan_root)

    # Step 0b: Inject codebase context before the user-supplied context
    codebase_ctx = build_codebase_context(codebase_summary)
    merged_context = (codebase_ctx + "\n" + (context or "")).strip() if codebase_ctx else (context or "")

    # Step 1: Run the decompose_project skill
    skill = load_skill("decompose_project", project_dir)
    inputs = {"goal": goal, "context": merged_context, "depth": depth}
    raw_output = await run_skill(skill, inputs, config, pool, project_id)

    # Step 2: Parse task definitions from LLM output
    raw_tasks = raw_output.get("tasks", [])
    if not raw_tasks:
        raise DecompositionError("LLM returned no tasks")

    task_defs = [TaskDefinition(**t) for t in raw_tasks]

    # Step 3: Optional enrichment (best-effort, failures silently skipped)
    if enrich:
        max_concurrent = enrichment_concurrency if enrichment_concurrency is not None else config.enrichment_concurrency
        task_defs = await _enrich_tasks(task_defs, goal, config, pool, project_id, project_dir, max_concurrent=max_concurrent)

    # Step 3b: Optional validation against existing codebase
    validation_summary = None
    if codebase_summary is not None:
        try:
            task_defs, validation_summary = _validate_tasks(task_defs, codebase_summary)
        except Exception:
            log.warning("Task validation failed — continuing without validation", exc_info=True)

    # Step 3c: Auto-fix stale file paths against real project tree
    path_corrections: list[dict] = []
    if auto_fix_paths and scan_root is not None:
        try:
            from loom.pathcheck import fix_paths_in_place
            path_corrections = fix_paths_in_place(task_defs, scan_root)
            if path_corrections:
                log.info("Auto-fixed %d stale file paths in task definitions", len(path_corrections))
        except Exception:
            log.warning("Auto-fix paths failed — continuing without path correction", exc_info=True)

    # Step 4: Build task graph with real IDs
    tasks = _build_task_graph(task_defs, project_id, parent_epic_id=parent_epic_id, force_leaf=force_leaf, parent_epic_deps=parent_epic_deps)

    # Step 5: Validate — cycle detection
    _validate_graph(tasks)

    # Step 5b: Post-decompose pruning — remove tasks satisfied by existing code
    pruning_summary = None
    if codebase_summary is not None:
        try:
            task_dicts = [t.model_dump(mode="json") for t in tasks]
            pruning_result = validate_task_graph(
                task_dicts, codebase_summary, mode="mark",
            )
            pruning_summary = pruning_result.stats
            # Rebuild tasks from the pruned dicts (mark mode preserves all)
            tasks = [
                Task(**td)
                for td in pruning_result.tasks
            ]
            log.info(
                "Post-decompose pruning: %d pre-completed, %d folded, %d remaining",
                pruning_result.stats.get("pre_completed_count", 0),
                pruning_result.stats.get("folded_count", 0),
                pruning_result.stats.get("remaining_count", 0),
            )
        except Exception:
            log.warning("Post-decompose pruning failed — continuing without pruning", exc_info=True)

    # Step 5c: Merge trivial tasks into parents/dependents
    if merge_trivial:
        try:
            threshold = min_complexity if min_complexity is not None else COMPLEXITY_THRESHOLD
            task_dicts = [t.model_dump(mode="json") for t in tasks]
            merged_tasks, merge_report = merge_trivial_tasks(task_dicts, threshold=threshold, dep_key="depends_on")
            if merge_report:
                log.debug("Merged %d trivial tasks", len(merge_report))
            tasks = [Task(**td) for td in merged_tasks]
        except Exception:
            log.warning("Trivial task merge failed — continuing without merge", exc_info=True)

    # Step 5e: Merge tasks with high file overlap
    if merge_file_affinity:
        try:
            task_dicts = [t.model_dump(mode="json") for t in tasks]
            affinity_merged, affinity_report = merge_by_file_affinity(task_dicts, dep_key="depends_on")
            if affinity_report:
                log.info("File-affinity merged %d task pairs", len(affinity_report))
            tasks = [Task(**td) for td in affinity_merged]
        except Exception:
            log.warning("File-affinity merge failed — continuing without", exc_info=True)

    # Step 5d: Optional parallelism optimization
    if optimize_parallelism:
        try:
            parallelism_skill = load_skill("identify_parallelism", project_dir)
            task_dicts = [t.model_dump(mode="json") for t in tasks]
            parallelism_result = await run_skill(
                parallelism_skill,
                {"tasks": task_dicts},
                config, pool, project_id,
            )
            if isinstance(parallelism_result, dict) and "tasks" in parallelism_result:
                tasks = [Task(**td) for td in parallelism_result["tasks"]]
                log.info("Parallelism optimization applied")
        except Exception:
            log.warning("Parallelism optimization failed — continuing without", exc_info=True)

    # Build result
    epic_titles = [t.title for t in tasks if t.status == TaskStatus.EPIC]
    dep_count = sum(len(t.depends_on) for t in tasks)
    result = DecompositionResult(
        tasks=[t.model_dump(mode="json") for t in tasks],
        epics=epic_titles,
        total_count=len(tasks),
        dependency_edges=dep_count,
        written_to_db=False,
        validation_summary=validation_summary,
    )
    if pruning_summary is not None:
        if result.validation_summary is None:
            result.validation_summary = {}
        result.validation_summary["pruning"] = pruning_summary
    if path_corrections:
        if result.validation_summary is None:
            result.validation_summary = {}
        result.validation_summary["path_corrections"] = path_corrections

    # Step 6: Write to DB if not in confirm mode
    if not confirm:
        await _write_to_db(tasks, pool, redis, project_id)
        result.written_to_db = True

    return result


async def _enrich_single_task(
    td: TaskDefinition,
    goal: str,
    config: SkillsConfig,
    pool: asyncpg.Pool | None,
    project_id: str | None,
    project_dir: str | None,
) -> None:
    """Enrich one task: context first, then done+complexity in parallel."""
    ctx_str = json.dumps(td.context) if td.context else ""

    # Step 1: Enrich context (must run first — later steps read td.context)
    try:
        skill = load_skill("write_task_context", project_dir)
        enriched = await run_skill(
            skill,
            {"task_title": td.title, "task_context": ctx_str, "project_goal": goal},
            config, pool, project_id,
        )
        if isinstance(enriched, dict):
            td.context = enriched
    except Exception:
        log.debug("Context enrichment skipped for %s", td.title, exc_info=True)

    # Step 2: define_done + estimate_complexity in parallel (both read td.context)
    async def _define_done() -> None:
        if td.done_when:
            return
        try:
            skill = load_skill("define_done", project_dir)
            done_result = await run_skill(
                skill,
                {"task_title": td.title, "task_context": json.dumps(td.context)},
                config, pool, project_id,
            )
            if isinstance(done_result, dict) and "done_when" in done_result:
                td.done_when = done_result["done_when"]
        except Exception:
            log.debug("Define-done skipped for %s", td.title, exc_info=True)

    async def _estimate_complexity() -> None:
        try:
            skill = load_skill("estimate_complexity", project_dir)
            est = await run_skill(
                skill,
                {"task_title": td.title, "task_context": json.dumps(td.context)},
                config, pool, project_id,
            )
            if isinstance(est, dict) and "complexity" in est:
                td.context["complexity"] = est["complexity"]
                if "reasoning" in est:
                    td.context["complexity_reasoning"] = est["reasoning"]
        except Exception:
            log.debug("Complexity estimation skipped for %s", td.title, exc_info=True)

    await asyncio.gather(_define_done(), _estimate_complexity())


async def _enrich_tasks(
    task_defs: list[TaskDefinition],
    goal: str,
    config: SkillsConfig,
    pool: asyncpg.Pool | None,
    project_id: str | None,
    project_dir: str | None,
    max_concurrent: int = 5,
) -> list[TaskDefinition]:
    """Best-effort enrichment: context, acceptance criteria, complexity.

    Runs up to *max_concurrent* tasks in parallel. Within each task,
    context enrichment runs first, then done-criteria and complexity
    estimation run concurrently.
    """
    sem = asyncio.Semaphore(max_concurrent)

    async def _bounded(td: TaskDefinition) -> None:
        async with sem:
            await _enrich_single_task(td, goal, config, pool, project_id, project_dir)

    non_epics = [td for td in task_defs if td.type != "epic"]
    await asyncio.gather(*[_bounded(td) for td in non_epics])
    return task_defs


def _build_task_graph(
    task_defs: list[TaskDefinition], project_id: str,
    parent_epic_id: str | None = None,
    force_leaf: bool = False,
    parent_epic_deps: list[str] | None = None,
) -> list[Task]:
    """Convert title-based TaskDefinitions into Task objects with real IDs.

    When parent_epic_id is provided, top-level epics (those without a
    parent_epic from the LLM) are linked to it, and tasks whose
    parent_epic title doesn't match any generated epic also fall back to it.

    When force_leaf is True, childless epics are converted to pending tasks
    so they can be claimed immediately.
    """
    # Generate IDs and build title -> ID mapping
    title_to_id: dict[str, str] = {}
    for td in task_defs:
        title_to_id[td.title] = gen_task_id()

    # Build epic title -> ID mapping for parent resolution
    epic_ids: dict[str, str] = {}
    for td in task_defs:
        if td.type == "epic":
            epic_ids[td.title] = title_to_id[td.title]

    tasks: list[Task] = []
    for td in task_defs:
        tid = title_to_id[td.title]

        # Resolve parent epic — fall back to parent_epic_id for unmatched/missing
        parent_id = None
        if td.parent_epic and td.parent_epic in epic_ids:
            parent_id = epic_ids[td.parent_epic]
        elif td.parent_epic and parent_epic_id:
            # LLM referenced a parent epic title that doesn't match — use fallback
            parent_id = parent_epic_id
        elif not td.parent_epic and td.type == "epic" and parent_epic_id:
            # Top-level epic with no parent — link to the provided parent
            parent_id = parent_epic_id

        # Resolve title-based dependencies to IDs (unknown titles silently dropped)
        depends_on = []
        for dep_title in td.depends_on_titles:
            if dep_title in title_to_id:
                depends_on.append(title_to_id[dep_title])

        # Inherit parent epic's cross-epic dependencies for non-epic children
        if parent_epic_deps and td.type != "epic":
            for dep_id in parent_epic_deps:
                if dep_id not in depends_on:
                    depends_on.append(dep_id)

        # Determine status
        if td.type == "epic":
            status = TaskStatus.EPIC
        elif depends_on:
            status = TaskStatus.BLOCKED
        else:
            status = TaskStatus.PENDING

        # Normalize priority
        priority = td.priority if td.priority in ("p0", "p1", "p2") else "p1"

        tasks.append(Task(
            id=tid,
            project_id=project_id,
            title=td.title,
            status=status,
            priority=Priority(priority),
            parent_id=parent_id,
            context=td.context,
            done_when=td.done_when,
            depends_on=depends_on,
        ))

    # Handle epics based on force_leaf setting
    epic_ids_set = {t.id for t in tasks if t.status == TaskStatus.EPIC}
    if force_leaf and epic_ids_set:
        # Convert ALL epics to pending tasks so they can be claimed immediately
        for t in tasks:
            if t.status == TaskStatus.EPIC:
                t.status = TaskStatus.PENDING if not t.depends_on else TaskStatus.BLOCKED
                log.info("Converted epic '%s' (%s) to %s task (force_leaf=True)", t.title, t.id, t.status.value)
    else:
        # Warn about childless epics that probably need further decomposition
        parented_epics = {t.parent_id for t in tasks if t.parent_id in epic_ids_set}
        childless = epic_ids_set - parented_epics
        for t in tasks:
            if t.id in childless:
                log.warning("Epic '%s' (%s) has no subtasks — consider decomposing it further with loom_decompose(epic_id='%s')", t.title, t.id, t.id)
                t.context = {**t.context, "_warning": "empty_epic"}

    return tasks


def _validate_graph(tasks: list[Task]) -> None:
    """In-memory DFS cycle detection on the proposed task graph."""
    adj: dict[str, list[str]] = {t.id: t.depends_on for t in tasks}
    all_ids = {t.id for t in tasks}

    WHITE, GRAY, BLACK = 0, 1, 2
    color: dict[str, int] = {tid: WHITE for tid in all_ids}

    def dfs(node: str, path: list[str]) -> None:
        color[node] = GRAY
        path.append(node)
        for neighbor in adj.get(node, []):
            if neighbor not in all_ids:
                continue
            if color[neighbor] == GRAY:
                cycle_start = path.index(neighbor)
                raise DecompositionError(
                    f"Dependency cycle: {' -> '.join(path[cycle_start:] + [neighbor])}"
                )
            if color[neighbor] == WHITE:
                dfs(neighbor, path)
        path.pop()
        color[node] = BLACK

    for tid in all_ids:
        if color[tid] == WHITE:
            dfs(tid, [])


async def _write_to_db(
    tasks: list[Task],
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
) -> None:
    """Write the full task graph to Postgres + Redis."""
    created = await store.create_tasks_batch(pool, tasks)

    for task in created:
        await cache.sync_task(redis, task)
        if task.status == TaskStatus.PENDING:
            await cache.add_to_ready_queue(redis, task)

    await publish_event(
        redis, project_id, EventType.PROJECT_DECOMPOSED,
        payload={"task_count": len(created)},
    )
    await store.record_event(
        pool, project_id, EventType.PROJECT_DECOMPOSED,
        payload={"task_count": len(created)},
    )
