"""One composition pass per codegen run.

The last completed GenerationPlan is the installed baseline. This run's
prompt is not appended to discovery. Newly selected blocks are unioned on.
A sentence that names an installed block for removal drops that block.
A newly admitted privileged block stays out until the owner approves it.
"""

from __future__ import annotations

from uuid import UUID

import structlog
from sqlalchemy.orm.attributes import flag_modified

from src.domain.app_blocks.catalogs import default_known_packs, default_primitives
from src.domain.app_blocks.compiler import compile_install_plan
from src.domain.app_blocks.job_packs import known_job_pack_ids
from src.domain.app_blocks.planner import (
    blocks_explicitly_removed,
    source_prompt_from_discovery,
)
from src.domain.app_blocks.registry import load_default_registry
from src.domain.codegen.app_block_plan import _latest_generation_plan
from src.domain.codegen.application_experience import compile_application_experience
from src.domain.codegen.generation_plan import (
    GenerationPlanV1,
    _plan_hash,
    compile_generation_plan,
    validate_generation_plan_catalog,
)
from src.domain.forms.canonical import FormDefinition
from src.domain.projects.discovery_catalog import CLOUD_CAPABILITIES
from src.domain.projects.discovery_schemas import DiscoveryResult
from src.domain.starter_library.archetypes import RECIPE_STARTER_IDS
from src.domain.starter_library.catalog import known_contract_ids

logger = structlog.get_logger()


def union_generation_plans(
    fresh: GenerationPlanV1,
    baseline: dict | None,
) -> GenerationPlanV1:
    """Add baseline capabilities the fresh compile did not select.

    A later prompt that only changes copy must not drop an installed block.
    Removal is applied by the caller after this union, and only for blocks
    the later sentence names.
    """
    if not isinstance(baseline, dict):
        return fresh
    base_composition = baseline.get("composition")
    if not isinstance(base_composition, dict):
        return fresh

    registry = load_default_registry()
    known_packs = default_known_packs() | known_job_pack_ids()
    known_primitives = default_primitives()
    known_contracts = known_contract_ids()
    known_capabilities = set(CLOUD_CAPABILITIES)
    known_capabilities.update(
        capability.id for block in registry.all() for capability in block.provides
    )

    def _block_ok(block_id: str) -> bool:
        block = registry.get(block_id)
        return block is not None and block.implemented

    def _pack_ok(pack_id: str) -> bool:
        from src.domain.app_blocks.catalogs import canonical_feature_pack

        return canonical_feature_pack(pack_id) in known_packs

    blocks = _union(
        fresh.composition.app_block_ids,
        base_composition.get("app_block_ids"),
        _block_ok,
    )
    packs = _union(
        fresh.composition.feature_pack_ids,
        base_composition.get("feature_pack_ids"),
        _pack_ok,
    )
    primitives = _union(
        fresh.composition.primitive_ids,
        base_composition.get("primitive_ids"),
        lambda item: item in known_primitives,
    )
    cloud = _union(
        fresh.composition.cloud_capability_ids,
        base_composition.get("cloud_capability_ids"),
        lambda item: item in known_capabilities,
    )
    recipes = _union(
        fresh.composition.recipe_ids,
        base_composition.get("recipe_ids"),
        lambda item: item in RECIPE_STARTER_IDS or item in fresh.composition.recipe_versions,
    )
    base_runtime = baseline.get("runtime") if isinstance(baseline.get("runtime"), dict) else {}
    contracts = _union(
        fresh.runtime.starter_contract_ids,
        (base_runtime or {}).get("starter_contract_ids"),
        lambda item: item in known_contracts,
    )
    forms = _merge_form_definitions(
        base_composition.get("form_definitions"),
        fresh.composition.form_definitions,
    )
    fresh_blocks = set(fresh.composition.app_block_ids)
    baseline_blocks = {
        str(item) for item in (base_composition.get("app_block_ids") or []) if isinstance(item, str)
    }
    experience = compile_application_experience(
        domain=fresh.intent.domain,
        archetype_id=fresh.composition.archetype_id,
        purpose=fresh.intent.purpose,
        app_block_ids=blocks,
        feature_pack_ids=packs,
        authentication_required=fresh.runtime.authentication_required,
        fresh_block_ids=fresh_blocks,
        baseline_block_ids=baseline_blocks,
    )
    required = sorted(
        set(blocks) | set(packs) | set(primitives) | set(cloud) | set(contracts)
    )
    merged = fresh.model_copy(
        update={
            "composition": fresh.composition.model_copy(
                update={
                    "app_block_ids": blocks,
                    "feature_pack_ids": packs,
                    "primitive_ids": primitives,
                    "cloud_capability_ids": cloud,
                    "recipe_ids": recipes,
                    "form_definitions": forms,
                }
            ),
            "experience": experience,
            "runtime": fresh.runtime.model_copy(
                update={"starter_contract_ids": contracts}
            ),
            "constraints": fresh.constraints.model_copy(
                update={"required_ids": required}
            ),
        }
    )
    validate_generation_plan_catalog(merged)
    merged.provenance.plan_hash = _plan_hash(merged)
    return merged


# Paid or provider-backed blocks. A later sentence may ask for them, but they
# stay out of the installed plan until a capability approval is approved.
PRIVILEGED_COMPOSITION_BLOCKS = frozenset({
    "payments.checkout",
    "billing.portal",
    "billing.tax",
    "booking.appointment",
})


def _feature_packs_for(block_id: str) -> set[str]:
    block = load_default_registry().get(block_id)
    if block is None:
        return set()
    return {item.id for item in block.requires if item.kind == "feature_pack"}


def without_blocks(plan: GenerationPlanV1, removed: set[str]) -> GenerationPlanV1:
    """Drop named blocks and feature packs that only those blocks required."""
    if not removed:
        return plan
    remaining = [block_id for block_id in plan.composition.app_block_ids if block_id not in removed]
    exclusive: set[str] = set()
    for block_id in removed:
        exclusive |= _feature_packs_for(block_id)
    for block_id in remaining:
        exclusive -= _feature_packs_for(block_id)
    packs = [pack_id for pack_id in plan.composition.feature_pack_ids if pack_id not in exclusive]
    required = [
        item
        for item in plan.constraints.required_ids
        if item not in removed and item not in exclusive
    ]
    updated = plan.model_copy(
        update={
            "composition": plan.composition.model_copy(
                update={"app_block_ids": remaining, "feature_pack_ids": packs}
            ),
            "constraints": plan.constraints.model_copy(update={"required_ids": required}),
        }
    )
    validate_generation_plan_catalog(updated)
    updated.provenance.plan_hash = _plan_hash(updated)
    return updated


def apply_privileged_hold(plan: dict, admitted: list[str], approved: list[str]) -> dict:
    """Keep a newly admitted privileged block out of the plan until it is approved."""
    approved_ids = set(approved)
    held = [
        block_id
        for block_id in admitted
        if block_id in PRIVILEGED_COMPOSITION_BLOCKS and block_id not in approved_ids
    ]
    stored = {key: value for key, value in plan.items() if not str(key).startswith("_")}
    if not held:
        dumped = dict(plan)
        dumped["_explicit_blocks_admitted"] = list(admitted)
        dumped["_privileged_blocks_pending"] = []
        return dumped
    updated = without_blocks(GenerationPlanV1.model_validate(stored), set(held))
    dumped = updated.model_dump(mode="json")
    for key, value in plan.items():
        if str(key).startswith("_"):
            dumped[key] = value
    dumped["_privileged_blocks_pending"] = held
    dumped["_explicit_blocks_admitted"] = [
        block_id for block_id in admitted if block_id not in set(held)
    ]
    return dumped


def compile_reconciled_plan(
    discovery: DiscoveryResult,
    *,
    source_prompt: str,
    baseline: dict | None,
    project_id: UUID | None,
    source_discovery_revision: int | None,
    composition_recipe=None,
    explicit_prompt: str = "",
) -> dict:
    """Compile from the pinned discovery closure, then keep the installed baseline.

    ``explicit_prompt`` is the later customer sentence. It is not appended to
    the discovery source. Blocks it admits are unioned onto the baseline.
    Copy-only sentences add nothing. ``_explicit_blocks_admitted`` is stripped
    by the caller before the plan is stored.
    """
    app_block_plan = compile_install_plan(
        discovery.spec,
        source_prompt=source_prompt,
        composition_recipe=composition_recipe,
    )
    fresh = compile_generation_plan(
        discovery,
        app_block_plan,
        project_id=project_id,
        source_discovery_revision=source_discovery_revision,
        catalog_digests={"discovery": "2"},
    )
    merged = union_generation_plans(fresh, baseline)
    before = set(merged.composition.app_block_ids)
    sentence = (explicit_prompt or "").strip()
    if sentence and sentence != (source_prompt or "").strip():
        delta = compile_generation_plan(
            discovery,
            compile_install_plan(
                discovery.spec,
                source_prompt=sentence,
                composition_recipe=composition_recipe,
            ),
            project_id=project_id,
            source_discovery_revision=source_discovery_revision,
            catalog_digests={"discovery": "2"},
        )
        merged = union_generation_plans(delta, merged.model_dump(mode="json"))
    removed = blocks_explicitly_removed(sentence, merged.composition.app_block_ids) if sentence else []
    if removed:
        merged = without_blocks(merged, set(removed))
    dumped = merged.model_dump(mode="json")
    dumped["_explicit_blocks_removed"] = removed
    dumped["_explicit_blocks_admitted"] = sorted(set(merged.composition.app_block_ids) - before)
    return dumped


async def reconcile_composition(
    db,
    project_id: UUID,
    user_prompt: str,
    settings,
    *,
    run_id: UUID | None = None,
) -> dict | None:
    """Sync structured requirements, recompile, and store this run's plan."""
    from src.domain.projects.entity_requirement_sync import (
        sync_entity_field_requirements_from_prompt,
    )
    from src.domain.projects.form_requirement_sync import sync_form_requirements_from_prompt

    try:
        await sync_form_requirements_from_prompt(db, project_id, user_prompt, settings)
    except Exception:
        logger.warning(
            "composition.form_requirement_sync_failed",
            project_id=str(project_id),
            exc_info=True,
        )
    try:
        await sync_entity_field_requirements_from_prompt(
            db, project_id, user_prompt, settings
        )
    except Exception:
        logger.warning(
            "composition.entity_requirement_sync_failed",
            project_id=str(project_id),
            exc_info=True,
        )

    from src.domain.projects.discovery_repository import DiscoveryRepository

    project = await DiscoveryRepository(db).get(project_id)
    if project is None:
        return None
    discovery_json = project.discovery_json or {}
    raw_result = discovery_json.get("result") or {}
    if not isinstance(raw_result, dict) or not isinstance(raw_result.get("spec"), dict):
        return None
    try:
        discovery = DiscoveryResult.model_validate(raw_result)
    except Exception:
        logger.warning(
            "composition.discovery_unreadable",
            project_id=str(project_id),
            exc_info=True,
        )
        return None

    source_prompt = source_prompt_from_discovery(discovery_json)
    from src.domain.composition_recipes.compile import pinned_recipe_for_state
    from src.domain.composition_recipes.installed_store import (
        overlay_installed_baseline,
        refresh_snapshot_for_explicit_change,
    )

    recipe = await pinned_recipe_for_state(
        db,
        discovery_json,
        source_prompt=source_prompt,
        project_type=discovery.project_type,
    )
    baseline = await overlay_installed_baseline(
        db, project_id, await _latest_generation_plan(db, project_id)
    )
    try:
        generation_plan = compile_reconciled_plan(
            discovery,
            source_prompt=source_prompt,
            baseline=baseline,
            project_id=project_id,
            source_discovery_revision=discovery_json.get("version"),
            composition_recipe=recipe,
            explicit_prompt=user_prompt,
        )
    except Exception:
        logger.warning(
            "composition.reconcile_failed",
            project_id=str(project_id),
            exc_info=True,
        )
        return None

    admitted = list(generation_plan.pop("_explicit_blocks_admitted", []) or [])
    removed = list(generation_plan.pop("_explicit_blocks_removed", []) or [])
    privileged = [block_id for block_id in admitted if block_id in PRIVILEGED_COMPOSITION_BLOCKS]
    if privileged:
        approved = await _approved_composition_blocks(db, project_id, run_id, privileged)
        generation_plan = apply_privileged_hold(generation_plan, privileged, approved)
        admitted = list(generation_plan.pop("_explicit_blocks_admitted", []) or [])
        generation_plan.pop("_privileged_blocks_pending", None)
    await refresh_snapshot_for_explicit_change(
        db,
        project_id,
        generation_plan,
        explicit_block_ids=sorted(set(admitted) | set(removed)),
    )
    if run_id is not None:
        await _persist_run_plan(db, run_id, generation_plan)
    logger.info(
        "composition.reconciled",
        project_id=str(project_id),
        run_id=str(run_id) if run_id else None,
        app_blocks=(generation_plan.get("composition") or {}).get("app_block_ids") or [],
        feature_packs=(generation_plan.get("composition") or {}).get("feature_pack_ids") or [],
    )
    return generation_plan


async def _approved_composition_blocks(db, project_id: UUID, run_id: UUID | None, block_ids: list[str]) -> list[str]:
    """Privileged blocks the owner has already approved. A store failure holds every block."""
    from src.domain.codegen.capability_approval import (
        Sensitive,
        SqlApprovalStore,
        _fingerprint,
        evaluate,
    )

    items = [
        Sensitive(
            "compose",
            block_id,
            _fingerprint("compose", block_id, [{"block": block_id}]),
            [{"block": block_id}],
            (
                f"Adds {block_id} to the installed application. "
                "This capability is paid or privileged and stays out of the plan until approved."
            ),
        )
        for block_id in block_ids
    ]
    try:
        decision = await evaluate(SqlApprovalStore(db), project_id, run_id, items)
    except Exception:
        logger.warning(
            "composition.privileged_approval_unavailable",
            project_id=str(project_id),
            exc_info=True,
        )
        return []
    return [str(getattr(row, "name", "")) for row in decision.approved if getattr(row, "name", "")]


async def _persist_run_plan(db, run_id: UUID, generation_plan: dict) -> None:
    from src.domain.codegen.repository import CodegenRunRepository

    run = await CodegenRunRepository(db).get_by_id(run_id)
    if run is None:
        return
    plan = dict(run.plan or {})
    stored = plan.get("generation_plan") if isinstance(plan.get("generation_plan"), dict) else {}
    if _stored_plan_is_newer(stored, generation_plan):
        logger.info(
            "composition.stale_reconcile_skipped",
            run_id=str(run_id),
        )
        return
    plan["schema"] = plan.get("schema") or "generation_plan.v1"
    plan["generation_plan"] = generation_plan
    plan["composition_reconciled"] = True
    run.plan = plan
    flag_modified(run, "plan")
    await db.commit()


def _stored_plan_is_newer(stored: dict, incoming: dict) -> bool:
    """A later reconcile on this run must not be replaced by an older one."""
    stored_at = ((stored.get("provenance") or {}) if isinstance(stored, dict) else {}).get(
        "compiled_at"
    )
    incoming_at = ((incoming.get("provenance") or {}) if isinstance(incoming, dict) else {}).get(
        "compiled_at"
    )
    if not stored_at or not incoming_at:
        return False
    return str(stored_at) > str(incoming_at)


def _union(current: list[str], baseline, admit) -> list[str]:
    values = {str(item) for item in current}
    if isinstance(baseline, list):
        values.update(str(item) for item in baseline if admit(str(item)))
    return sorted(values)


def _merge_form_definitions(
    baseline, fresh: dict[str, FormDefinition],
) -> dict[str, FormDefinition]:
    merged: dict[str, FormDefinition] = {}
    if isinstance(baseline, dict):
        for key, value in baseline.items():
            try:
                definition = FormDefinition.model_validate(value)
            except Exception:
                continue
            if definition.key == key:
                merged[str(key)] = definition
    for key, definition in fresh.items():
        merged[key] = definition
    return merged