#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from common import DeployError, db_exec, db_query_scalar, load_env_file
from salesforce_file_projection_audit import (
    DEFAULT_SOURCE_TYPE,
    DEFAULT_STORAGE_BUCKET,
    DEFAULT_STORAGE_KEY_PREFIX,
    DEFAULT_STORAGE_PROVIDER,
    normalize_logical_storage_key_prefix,
    normalize_storage_key_prefix,
    require_test_slot_env,
    sanitize_identifier,
    sql_string_literal,
    unique_ordered,
)


def query_batch_state(envv: dict[str, str], sql: str) -> dict[str, int]:
    raw = db_query_scalar(envv, sql, database=envv["TIDB_DATABASE"])
    first = raw.splitlines()[0].strip() if raw.splitlines() else ""
    fields = first.split("\t")
    if len(fields) != 2:
        raise DeployError(f"invalid compatibility URL batch state: {first!r}")
    try:
        return {
            "existing_target_count": int(fields[0]),
            "absolute_url_count": int(fields[1]),
        }
    except ValueError as exc:
        raise DeployError(f"invalid compatibility URL batch state: {first!r}") from exc


def absolute_url_predicate(alias: str, field: str) -> str:
    return (
        f"TRIM(CAST({alias}.`{sanitize_identifier(field)}` AS CHAR)) "
        "REGEXP '^https?://'"
    )


def compatibility_url_file_id_expr(alias: str, field: str) -> str:
    value = f"TRIM(CAST({alias}.`{sanitize_identifier(field)}` AS CHAR))"
    return (
        "SUBSTRING_INDEX("
        f"SUBSTRING_INDEX({value}, '/files/', -1), "
        "'/download', 1)"
    )


def eligible_file_rows_sql(
    *,
    system_file_table: str,
    parent_object_api_name: str,
    attachment_kind: str,
    source_type: str,
    storage_provider: str,
    storage_bucket: str,
    storage_key_prefix: str,
    logical_storage_key_prefix: str,
) -> str:
    physical_prefix = normalize_storage_key_prefix(storage_key_prefix)
    logical_prefix = normalize_logical_storage_key_prefix(
        logical_storage_key_prefix
    )
    expected_storage_key = (
        f"CONCAT({sql_string_literal(logical_prefix + '/')}, SourceContentVersionId)"
        if logical_prefix
        else "CAST(SourceContentVersionId AS CHAR)"
    )
    expected_storage_path = (
        "CONCAT('s3://', "
        f"{sql_string_literal(storage_bucket)}, '/', "
        f"{sql_string_literal(physical_prefix + '/')}, SourceContentVersionId)"
    )
    return f"""
SELECT ParentId, Id
FROM `{sanitize_identifier(system_file_table)}`
WHERE ParentObjectApiName = {sql_string_literal(parent_object_api_name)}
  AND AttachmentKind = {sql_string_literal(attachment_kind)}
  AND SourceType = {sql_string_literal(source_type)}
  AND SourceContentVersionId IS NOT NULL
  AND TRIM(CAST(SourceContentVersionId AS CHAR)) <> ''
  AND StorageProvider = {sql_string_literal(storage_provider)}
  AND StorageBucket = {sql_string_literal(storage_bucket)}
  AND StorageKey = {expected_storage_key}
  AND StoragePath = {expected_storage_path}
  AND Status = 'Active'
  AND ParentId IS NOT NULL
  AND TRIM(CAST(ParentId AS CHAR)) <> ''
""".strip()


def build_reconcile_update_sql(
    *,
    target_table: str,
    compatibility_url_field: str,
    system_file_table: str,
    parent_object_api_name: str,
    attachment_kind: str,
    target_parent_ids: list[str],
    source_type: str = DEFAULT_SOURCE_TYPE,
    storage_provider: str = DEFAULT_STORAGE_PROVIDER,
    storage_bucket: str = DEFAULT_STORAGE_BUCKET,
    storage_key_prefix: str = DEFAULT_STORAGE_KEY_PREFIX,
    logical_storage_key_prefix: str = "",
) -> str:
    files = eligible_file_rows_sql(
        system_file_table=system_file_table,
        parent_object_api_name=parent_object_api_name,
        attachment_kind=attachment_kind,
        source_type=source_type,
        storage_provider=storage_provider,
        storage_bucket=storage_bucket,
        storage_key_prefix=storage_key_prefix,
        logical_storage_key_prefix=logical_storage_key_prefix,
    )
    parent_ids = unique_ordered(target_parent_ids)
    target_scope = target_scope_predicate(
        alias="target_record",
        target_parent_ids=parent_ids,
    )
    guard_scope = target_scope_predicate(
        alias="guard_target",
        target_parent_ids=parent_ids,
    )
    expected_target_count = len(parent_ids)
    return f"""
UPDATE `{sanitize_identifier(target_table)}` target_record
JOIN (
{files}
) selected_file
  ON selected_file.ParentId = target_record.Id
 AND CAST(selected_file.Id AS CHAR) = {compatibility_url_file_id_expr('target_record', compatibility_url_field)}
JOIN (
  SELECT COUNT(DISTINCT CAST(guard_target.Id AS CHAR)) AS existing_target_count,
         SUM(CASE WHEN {absolute_url_predicate('guard_target', compatibility_url_field)} THEN 1 ELSE 0 END) AS scoped_absolute_count,
         SUM(CASE WHEN {absolute_url_predicate('guard_target', compatibility_url_field)} AND guard_file.Id IS NOT NULL THEN 1 ELSE 0 END) AS resolvable_count
  FROM `{sanitize_identifier(target_table)}` guard_target
  LEFT JOIN (
{files}
  ) guard_file
    ON guard_file.ParentId = guard_target.Id
   AND CAST(guard_file.Id AS CHAR) = {compatibility_url_file_id_expr('guard_target', compatibility_url_field)}
  WHERE {guard_scope}
) reconciliation_guard
  ON reconciliation_guard.existing_target_count = {expected_target_count}
 AND reconciliation_guard.scoped_absolute_count = reconciliation_guard.resolvable_count
SET target_record.`{sanitize_identifier(compatibility_url_field)}` =
    CONCAT('/services/data/v63.0/files/', selected_file.Id, '/download')
WHERE {absolute_url_predicate('target_record', compatibility_url_field)}
  AND ({target_scope})
""".strip()


def target_scope_predicate(
    *,
    alias: str,
    target_parent_ids: list[str],
) -> str:
    parent_ids = unique_ordered(target_parent_ids)
    if not parent_ids:
        raise DeployError("at least one --target-parent-id is required")
    id_literals = ", ".join(sql_string_literal(parent_id) for parent_id in parent_ids)
    return f"CAST({alias}.`Id` AS CHAR) IN ({id_literals})"


def build_batch_state_sql(
    *,
    target_table: str,
    compatibility_url_field: str,
    target_parent_ids: list[str],
) -> str:
    target_scope = target_scope_predicate(
        alias="target_record",
        target_parent_ids=target_parent_ids,
    )
    return f"""
SELECT COUNT(*) AS existing_target_count,
       COALESCE(SUM(CASE WHEN {absolute_url_predicate('target_record', compatibility_url_field)} THEN 1 ELSE 0 END), 0) AS absolute_url_count
FROM `{sanitize_identifier(target_table)}` target_record
WHERE {target_scope}
""".strip()


def reconcile_file_compatibility_urls(
    envv: dict[str, str],
    *,
    apply: bool,
    parent_object_api_name: str,
    target_table: str,
    attachment_kind: str,
    compatibility_url_field: str,
    target_parent_ids: list[str],
    system_file_table: str = "_System_File",
    source_type: str = DEFAULT_SOURCE_TYPE,
    storage_provider: str = DEFAULT_STORAGE_PROVIDER,
    storage_bucket: str = DEFAULT_STORAGE_BUCKET,
    storage_key_prefix: str = DEFAULT_STORAGE_KEY_PREFIX,
    logical_storage_key_prefix: str = "",
    allow_non_test_env: bool = False,
) -> dict[str, Any]:
    require_test_slot_env(envv, allow_non_test_env=allow_non_test_env)
    parent_ids = unique_ordered(target_parent_ids)
    state_sql = build_batch_state_sql(
        target_table=target_table,
        compatibility_url_field=compatibility_url_field,
        target_parent_ids=parent_ids,
    )
    expected_target_count = len(parent_ids)
    before = query_batch_state(envv, state_sql)
    complete_before = before["existing_target_count"] == expected_target_count
    if apply and before["absolute_url_count"] and complete_before:
        db_exec(
            envv,
            build_reconcile_update_sql(
                target_table=target_table,
                compatibility_url_field=compatibility_url_field,
                system_file_table=system_file_table,
                parent_object_api_name=parent_object_api_name,
                attachment_kind=attachment_kind,
                target_parent_ids=parent_ids,
                source_type=source_type,
                storage_provider=storage_provider,
                storage_bucket=storage_bucket,
                storage_key_prefix=storage_key_prefix,
                logical_storage_key_prefix=logical_storage_key_prefix,
            ),
            database=envv["TIDB_DATABASE"],
        )
    after = query_batch_state(envv, state_sql) if apply else before
    complete_after = after["existing_target_count"] == expected_target_count
    passed = complete_before and complete_after and after["absolute_url_count"] == 0
    rewritten = (
        max(0, before["absolute_url_count"] - after["absolute_url_count"])
        if complete_before and complete_after
        else 0
    )
    return {
        "status": "passed" if passed else "failed",
        "mode": "apply" if apply else "audit-only",
        "parent_object_api_name": parent_object_api_name,
        "target_table": target_table,
        "attachment_kind": attachment_kind,
        "source_type": source_type,
        "storage_provider": storage_provider,
        "storage_bucket": storage_bucket,
        "storage_key_prefix": normalize_storage_key_prefix(storage_key_prefix),
        "logical_storage_key_prefix": normalize_logical_storage_key_prefix(
            logical_storage_key_prefix
        ),
        "compatibility_url_field": compatibility_url_field,
        "target_parent_ids": parent_ids,
        "expected_target_count": expected_target_count,
        "existing_targets_before": before["existing_target_count"],
        "existing_targets_after": after["existing_target_count"],
        "missing_targets_before": max(
            0, expected_target_count - before["existing_target_count"]
        ),
        "missing_targets_after": max(
            0, expected_target_count - after["existing_target_count"]
        ),
        "absolute_urls_before": before["absolute_url_count"],
        "absolute_urls_after": after["absolute_url_count"],
        "rewritten": rewritten,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Reconcile copied absolute compatibility URLs to current-environment Files paths."
    )
    parser.add_argument("--env-file", required=True)
    parser.add_argument("--parent-object-api-name", required=True)
    parser.add_argument("--target-table", required=True)
    parser.add_argument("--attachment-kind", required=True)
    parser.add_argument("--source-type", default=DEFAULT_SOURCE_TYPE)
    parser.add_argument("--storage-provider", default=DEFAULT_STORAGE_PROVIDER)
    parser.add_argument("--storage-bucket", default=DEFAULT_STORAGE_BUCKET)
    parser.add_argument("--storage-key-prefix", default=DEFAULT_STORAGE_KEY_PREFIX)
    parser.add_argument("--logical-storage-key-prefix", default="")
    parser.add_argument("--compatibility-url-field", required=True)
    parser.add_argument("--system-file-table", default="_System_File")
    parser.add_argument(
        "--target-parent-id",
        action="append",
        default=[],
        help="Audited parent record ID to reconcile. May be repeated.",
    )
    parser.add_argument("--apply", action="store_true")
    parser.add_argument("--allow-non-test-env", action="store_true")
    parser.add_argument("--summary-json")
    parser.add_argument("--fail-on-problems", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    envv = load_env_file(Path(args.env_file).expanduser().resolve())
    summary = reconcile_file_compatibility_urls(
        envv,
        apply=args.apply,
        parent_object_api_name=args.parent_object_api_name,
        target_table=args.target_table,
        attachment_kind=args.attachment_kind,
        source_type=args.source_type,
        storage_provider=args.storage_provider,
        storage_bucket=args.storage_bucket,
        storage_key_prefix=args.storage_key_prefix,
        logical_storage_key_prefix=args.logical_storage_key_prefix,
        compatibility_url_field=args.compatibility_url_field,
        target_parent_ids=args.target_parent_id,
        system_file_table=args.system_file_table,
        allow_non_test_env=args.allow_non_test_env,
    )
    output = json.dumps(summary, indent=2, ensure_ascii=False)
    if args.summary_json:
        Path(args.summary_json).write_text(output + "\n", encoding="utf-8")
    print(output)
    return 1 if args.fail_on_problems and summary["status"] != "passed" else 0


if __name__ == "__main__":
    raise SystemExit(main())