"""
AYON Representation Registration - DIAGNOSTIC VERSION
- Logs to stderr regardless of what fails
- Wraps ayon_api import early so we see import-time failures
- Tolerates missing work_item output attribs (logs but doesn't crash)
"""

import os
import sys
import uuid
import hashlib
import platform
import traceback
import datetime


# ---------------------------------------------------------------------------
# Diagnostic logging - writes to stderr, never throws
# ---------------------------------------------------------------------------

def _log(msg):
    line = "[AYON-REG %s] %s\n" % (datetime.datetime.now().strftime("%H:%M:%S"), msg)
    try:
        sys.stderr.write(line)
        sys.stderr.flush()
    except Exception:
        pass


_log("=" * 60)
_log("AYON REG SCRIPT STARTING")
_log("Python: %s" % sys.version.split()[0])
_log("Platform: %s" % platform.system())


# ---------------------------------------------------------------------------
# Import ayon modules early so we see import errors clearly
# ---------------------------------------------------------------------------

try:
    import ayon_api
    _log("imported ayon_api OK")
except Exception as e:
    _log("FAILED to import ayon_api: %s" % e)
    _log(traceback.format_exc())
    raise

try:
    from ayon_api.operations import OperationsSession
    _log("imported OperationsSession OK")
except Exception as e:
    _log("FAILED to import OperationsSession: %s" % e)
    raise

try:
    from ayon_core.pipeline import get_current_project_name
    _log("imported get_current_project_name OK")
except Exception as e:
    _log("FAILED to import get_current_project_name: %s" % e)
    raise

try:
    from ayon_core.pipeline.anatomy import Anatomy
    _log("imported Anatomy OK")
except Exception as e:
    _log("FAILED to import Anatomy: %s" % e)
    raise


# ---------------------------------------------------------------------------
# Tolerant work_item attrib helpers
# ---------------------------------------------------------------------------

def safe_set_string(name, value):
    """Set a work_item string attrib if possible; never throws."""
    try:
        work_item.setStringAttrib(name, str(value))
    except Exception as e:
        _log("setStringAttrib(%s) failed (probably attrib not declared on TOP node): %s" % (name, e))


def safe_get_string(name):
    try:
        return work_item.stringAttribValue(name)
    except Exception as e:
        _log("stringAttribValue(%s) failed: %s" % (name, e))
        return ""


def safe_get_int(name, default=0):
    try:
        return work_item.intAttribValue(name)
    except Exception as e:
        _log("intAttribValue(%s) failed: %s" % (name, e))
        return default


# ---------------------------------------------------------------------------
# File / path helpers
# ---------------------------------------------------------------------------

def compute_file_hash(filepath, algorithm="md5"):
    h = hashlib.new(algorithm)
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def make_rootless_path(abs_path, root_value, root_name="work"):
    abs_path = abs_path.replace("\\", "/")
    root_value = root_value.replace("\\", "/").rstrip("/")
    if abs_path.startswith(root_value):
        return "{root[%s]}%s" % (root_name, abs_path[len(root_value):])
    if platform.system().lower() == "windows":
        if abs_path.lower().startswith(root_value.lower()):
            return "{root[%s]}%s" % (root_name, abs_path[len(root_value):])
    return abs_path


def build_file_entry(abs_path, root_value, root_name="work"):
    return {
        "id": uuid.uuid4().hex,
        "name": os.path.basename(abs_path),
        "path": make_rootless_path(abs_path, root_value, root_name),
        "size": os.path.getsize(abs_path),
        "hash": compute_file_hash(abs_path, "md5"),
        "hashType": "md5",
    }


def discover_usd_files(publish_dir):
    if not os.path.isdir(publish_dir):
        _log("publish_dir does NOT exist on disk: %s" % publish_dir)
        return {"main": [], "secondary": [], "all": []}

    usd_extensions = {".usd", ".usda", ".usdc", ".usdz"}
    all_files, main_files, secondary_files = [], [], []
    for filename in sorted(os.listdir(publish_dir)):
        filepath = os.path.join(publish_dir, filename).replace("\\", "/")
        _, ext = os.path.splitext(filename)
        if ext.lower() in usd_extensions and os.path.isfile(filepath):
            all_files.append(filepath)
            if ext.lower() == ".usd":
                main_files.append(filepath)
            else:
                secondary_files.append(filepath)
    if not main_files and all_files:
        main_files.append(all_files[0])
        secondary_files = all_files[1:]
    return {"main": main_files, "secondary": secondary_files, "all": all_files}


# ---------------------------------------------------------------------------
# Anatomy-based publish directory resolver
# ---------------------------------------------------------------------------

def compute_publish_dir(project_name, folder_path, product_name, product_type, version_num):
    anatomy = Anatomy(project_name)

    folder_parts = folder_path.strip("/").split("/")
    folder_name = folder_parts[-1]
    hierarchy = "/".join(folder_parts[:-1]) if len(folder_parts) > 1 else ""

    project_data = ayon_api.get_project(project_name)
    project_code = project_data.get("code", project_name)

    template_data = {
        "root": anatomy.roots,
        "project": {"name": project_name, "code": project_code},
        "folder": {"name": folder_name, "path": folder_path},
        "hierarchy": hierarchy,
        "product": {"name": product_name, "type": product_type},
        "version": version_num,
        "representation": "usd",
        "ext": "usd",
        "user": "",
        "comment": "",
    }

    publish_template = anatomy.get_template_item("publish", "default", "directory")
    result = publish_template.format_strict(template_data)
    publish_dir = result.normalized()

    root_name = "work"
    root_value = str(anatomy.roots[root_name]).replace("\\", "/")

    return publish_dir, root_value, root_name


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

try:
    # --- Check work_item is even available ---
    try:
        wi_name = work_item.name if hasattr(work_item, "name") else "<unknown>"
        _log("work_item present: %s" % wi_name)
    except NameError:
        _log("work_item is NOT defined - is this running in a TOPs Python Script context?")
        raise

    # --- Read input attribs ---
    asset_name   = safe_get_string("asset_name")
    ayon_status  = safe_get_string("ayon_status")
    folder_path  = safe_get_string("ayon_folder_path")
    folder_id    = safe_get_string("ayon_folder_id")
    product_name = safe_get_string("ayon_product_name")
    version_id   = safe_get_string("ayon_version_id")
    version_num  = safe_get_int("ayon_version", 0)

    _log("inputs read:")
    _log("  asset_name        = %r" % asset_name)
    _log("  ayon_status       = %r" % ayon_status)
    _log("  ayon_folder_path  = %r" % folder_path)
    _log("  ayon_product_name = %r" % product_name)
    _log("  ayon_version_id   = %r" % version_id)
    _log("  ayon_version      = %r" % version_num)

    if ayon_status != "success":
        _log("ayon_status is not 'success', skipping: %r" % ayon_status)
        safe_set_string("ayon_reg_status", "skipped")
        safe_set_string("ayon_reg_error", "ayon_status=%s" % ayon_status)
        raise SystemExit(0)

    if not version_id:
        raise RuntimeError("Missing ayon_version_id on work item")

    project_name = get_current_project_name()
    _log("project_name = %r" % project_name)

    con = ayon_api.get_server_api_connection()
    _log("got AYON server connection")

    # --- Resolve publish dir via Anatomy ---
    product_type = "usd"
    _log("calling compute_publish_dir(...)")
    publish_dir_clean, root_path, root_name = compute_publish_dir(
        project_name, folder_path, product_name, product_type, version_num
    )
    publish_dir_clean = publish_dir_clean.replace("\\", "/")
    root_path = root_path.rstrip("/")
    _log("publish_dir resolved: %s" % publish_dir_clean)
    _log("root_path = %s   root_name = %s" % (root_path, root_name))

    # --- Find USDs on disk ---
    usd_files = discover_usd_files(publish_dir_clean)
    _log("usd_files: main=%d secondary=%d total=%d" % (
        len(usd_files["main"]), len(usd_files["secondary"]), len(usd_files["all"])))
    for f in usd_files["all"]:
        _log("  found: %s" % f)

    if not usd_files["all"]:
        raise RuntimeError("No USD files found in: %s" % publish_dir_clean)

    main_file = usd_files["main"][0] if usd_files["main"] else usd_files["all"][0]
    main_rootless = make_rootless_path(main_file, root_path, root_name)
    _log("main_file     = %s" % main_file)
    _log("main_rootless = %s" % main_rootless)

    files_list = [build_file_entry(fp, root_path, root_name) for fp in usd_files["all"]]

    # --- Build context_data for the representation ---
    folder_parts = folder_path.strip("/").split("/")
    folder_name = folder_parts[-1] if folder_parts else asset_name
    hierarchy = "/".join(folder_parts[:-1]) if len(folder_parts) > 1 else ""

    project_data = ayon_api.get_project(project_name)
    project_code = project_data.get("code", project_name)

    context_data = {
        "project": {"name": project_name, "code": project_code},
        "folder": {"name": folder_name, "path": folder_path},
        "product": {"name": product_name, "type": "usd"},
        "version": version_num,
        "representation": "usd",
        "hierarchy": hierarchy,
        "root": {root_name: root_path},
        "ext": "usd",
    }

    # --- Find existing representation ---
    _log("looking up existing representation 'usd' for version_id=%s" % version_id)
    repre = ayon_api.get_representation_by_name(
        project_name, "usd", version_id,
        fields={"id", "attrib", "context", "data", "active", "files", "name"},
    )
    _log("existing repre found: %s" % bool(repre))

    desired_attrib = {
        "template": main_rootless,
        "path": main_rootless,
    }

    session = OperationsSession()

    if repre:
        repre_id = repre["id"]
        cur_attrib = repre.get("attrib") or {}
        cur_template = cur_attrib.get("template")
        needs_template = (not isinstance(cur_template, str)) or (not cur_template.strip())
        needs_context = not isinstance(repre.get("context"), dict)

        migrated_context = (repre.get("data") or {}).get("context") if needs_context else None
        final_context = context_data if not migrated_context else migrated_context

        _log("repair path: needs_template=%s needs_context=%s" % (needs_template, needs_context))

        if needs_template or needs_context:
            try:
                session.update_representation(
                    project_name, repre_id,
                    attrib=desired_attrib,
                    context=final_context,
                    files=files_list,
                    active=True,
                )
                _log("session.update_representation called")
            except TypeError as te:
                _log("update_representation TypeError, falling back to PATCH: %s" % te)
                con.patch(
                    "projects/%s/representations/%s" % (project_name, repre_id),
                    attrib=desired_attrib,
                    context=final_context,
                    files=files_list,
                    active=True,
                )
            session.commit()
            _log("session.commit done")

        safe_set_string("ayon_repre_id", repre_id)
        safe_set_string("ayon_reg_status", "repaired_or_ok")
        safe_set_string("ayon_reg_error", "")
        _log("DONE (repaired_or_ok) repre_id=%s" % repre_id)

    else:
        _log("creating new representation")
        try:
            repre_id = ayon_api.create_representation(
                project_name,
                name="usd",
                version_id=version_id,
                files=files_list,
                attrib=desired_attrib,
                context=context_data,
                data={"context": context_data},
                tags=[],
                status=None,
                active=True,
            )
            _log("ayon_api.create_representation returned id=%s" % repre_id)
        except TypeError as te:
            _log("create_representation TypeError, falling back to POST: %s" % te)
            repre_id = uuid.uuid4().hex
            con.post(
                "projects/%s/representations" % project_name,
                id=repre_id,
                name="usd",
                versionId=version_id,
                files=files_list,
                attrib=desired_attrib,
                context=context_data,
                data={"context": context_data},
                active=True,
            )

        safe_set_string("ayon_repre_id", repre_id)
        safe_set_string("ayon_reg_status", "created")
        safe_set_string("ayon_reg_error", "")
        _log("DONE (created) repre_id=%s" % repre_id)


except SystemExit:
    pass
except Exception as e:
    tb = traceback.format_exc()
    _log("=" * 60)
    _log("SCRIPT FAILED: %s" % e)
    _log(tb)
    _log("=" * 60)
    safe_set_string("ayon_reg_status", "failed")
    safe_set_string("ayon_reg_error", str(e)[:500])
    safe_set_string("ayon_repre_id", "")
    raise