"""MODE A / MODE B scene setup for blender_render.py.

MODE A — factory settings (no template): auto-camera + auto-lights
MODE B — template file: load .blend, import into named collection
"""
import time as _time
from typing import Callable

import bpy  # type: ignore[import]

from _blender_camera    import setup_auto_camera, setup_auto_lights
from _blender_import    import import_glb, apply_rotation, import_usd_file
from _blender_materials import (
    assign_failed_material,
    build_mat_map_lower,
    apply_material_library,
    apply_material_library_direct,
)
from _blender_scene     import (
    ensure_collection,
    apply_smooth_batch,
    apply_sharp_edges_from_occ,
    setup_shadow_catcher,
)


def setup_scene(args, lap_fn: Callable[[str], None]) -> None:
    """Set up the Blender scene according to args (MODE A or B).

    Handles import, rotation, smooth shading, material assignment, shadow
    catcher, and auto-camera/lights. lap_fn is called with a label string
    at each timing checkpoint.
    """
    if args.use_template:
        _setup_mode_b(args, lap_fn)
    else:
        _setup_mode_a(args)


def _setup_mode_b(args, lap_fn: Callable[[str], None]) -> None:
    """MODE B: Template-based render — load .blend, import into collection."""
    print(f"[blender_render] Opening template: {args.template_path}")
    bpy.ops.wm.open_mainfile(filepath=args.template_path)
    lap_fn("template_load")

    target_col = ensure_collection(args.target_collection)
    usd_material_lookup: dict = {}
    if args.usd_path:
        parts, usd_material_lookup = import_usd_file(args.usd_path)
    else:
        parts = import_glb(args.glb_path)
    lap_fn("glb_import")
    apply_rotation(parts, args.rotation_x, args.rotation_y, args.rotation_z)
    lap_fn("rotation")

    for part in parts:
        for col in list(part.users_collection):
            col.objects.unlink(part)
        target_col.objects.link(part)

    apply_smooth_batch(parts, args.smooth_angle)
    if not args.usd_path:
        _occ_pairs = args.mesh_attributes.get("sharp_edge_pairs") or []
        if _occ_pairs:
            apply_sharp_edges_from_occ(parts, _occ_pairs)
    lap_fn("smooth_shading")

    # Apply material override: replace all material lookups with a single material
    if getattr(args, 'material_override', None):
        print(f"[blender_render] material_override active: all parts → {args.material_override}", flush=True)
        if usd_material_lookup:
            usd_material_lookup = {k: args.material_override for k in usd_material_lookup}
        if args.material_map:
            args.material_map = {k: args.material_override for k in args.material_map}

    if args.material_library_path and usd_material_lookup:
        # USD primvar path: direct material assignment (no name-matching needed)
        apply_material_library_direct(
            parts, args.material_library_path, usd_material_lookup,
        )
        # Fall back to name-matching for any parts missing primvars
        if args.material_map:
            _unassigned = [p for p in parts if not p.data.materials or
                           (len(p.data.materials) == 1 and
                            p.data.materials[0] and
                            p.data.materials[0].name == "HARTOMAT_059999_FailedMaterial")]
            if _unassigned:
                print(f"[blender_render] {len(_unassigned)} parts without USD primvar — "
                      f"falling back to name-matching", flush=True)
                apply_material_library(
                    _unassigned, args.material_library_path,
                    build_mat_map_lower(args.material_map), args.part_names_ordered,
                )
    elif args.material_library_path and args.material_map:
        apply_material_library(
            parts, args.material_library_path,
            build_mat_map_lower(args.material_map), args.part_names_ordered,
        )
    else:
        for part in parts:
            assign_failed_material(part)
    lap_fn("material_assign")

    if args.shadow_catcher:
        setup_shadow_catcher(parts)

    needs_auto_camera = (
        (args.lighting_only and not args.shadow_catcher)
        or not bpy.context.scene.camera
    )
    if args.lighting_only and not args.shadow_catcher:
        print("[blender_render] lighting_only mode: using template World/HDRI, forcing auto-camera")
    elif needs_auto_camera:
        print("[blender_render] WARNING: template has no camera — will create auto-camera")

    if not needs_auto_camera and bpy.context.scene.camera:
        bpy.context.scene.camera.data.clip_start = 0.001

    print(f"[blender_render] template mode: {len(parts)} parts imported into collection '{args.target_collection}'")

    if needs_auto_camera:
        setup_auto_camera(parts, args.width, args.height,
                          lens_mm=args.focal_length_mm,
                          sensor_width_mm=args.sensor_width_mm)


def _setup_mode_a(args) -> None:
    """MODE A: Factory settings — auto-camera + auto-lights."""
    bpy.ops.wm.read_factory_settings(use_empty=True)
    usd_material_lookup: dict = {}
    if args.usd_path:
        parts, usd_material_lookup = import_usd_file(args.usd_path)
    else:
        parts = import_glb(args.glb_path)
    apply_rotation(parts, args.rotation_x, args.rotation_y, args.rotation_z)

    _t = _time.time()
    apply_smooth_batch(parts, args.smooth_angle)
    if not args.usd_path:
        _occ_pairs = args.mesh_attributes.get("sharp_edge_pairs") or []
        if _occ_pairs:
            apply_sharp_edges_from_occ(parts, _occ_pairs)
    for part in parts:
        assign_failed_material(part)
    print(f"[blender_render] smooth+fallback-material: {len(parts)} parts ({_time.time()-_t:.2f}s)", flush=True)

    # Apply material override: replace all material lookups with a single material
    if getattr(args, 'material_override', None):
        print(f"[blender_render] material_override active (Mode A): all parts → {args.material_override}", flush=True)
        if usd_material_lookup:
            usd_material_lookup = {k: args.material_override for k in usd_material_lookup}
        if args.material_map:
            args.material_map = {k: args.material_override for k in args.material_map}

    if args.material_library_path and usd_material_lookup:
        # USD primvar path: direct material assignment
        apply_material_library_direct(
            parts, args.material_library_path, usd_material_lookup,
        )
        # Fall back to name-matching for parts without primvars
        if args.material_map:
            _unassigned = [p for p in parts if not p.data.materials or
                           (len(p.data.materials) == 1 and
                            p.data.materials[0] and
                            p.data.materials[0].name == "HARTOMAT_059999_FailedMaterial")]
            if _unassigned:
                apply_material_library(
                    _unassigned, args.material_library_path,
                    build_mat_map_lower(args.material_map), args.part_names_ordered,
                )
    elif args.material_library_path and args.material_map:
        apply_material_library(
            parts, args.material_library_path,
            build_mat_map_lower(args.material_map), args.part_names_ordered,
        )

    bbox_center, bsphere_radius = setup_auto_camera(parts, args.width, args.height,
                                                     lens_mm=args.focal_length_mm,
                                                     sensor_width_mm=args.sensor_width_mm)
    setup_auto_lights(bbox_center, bsphere_radius)
    world = bpy.data.worlds.new("World")
    bpy.context.scene.world = world
    world.use_nodes = True
    bg = world.node_tree.nodes["Background"]
    bg.inputs["Color"].default_value    = (0.96, 0.96, 0.97, 1.0)
    bg.inputs["Strength"].default_value = 0.15