"""get_customer_metadata tool -- pulls customer code and config from the org.

This tool:
1. Generates package.xml manifest(s) covering all relevant metadata types
2. Runs sf project retrieve start to pull metadata into the SFDX project
3. Uses adaptive auto-scaling: if LIMIT_EXCEEDED, splits into smaller batches
   via recursive binary split until all batches succeed
4. Records successful batch strategy in state for idempotent resume
5. Reports retrieval counts per metadata type
"""

import json
import logging
from pathlib import Path
from typing import Any
from xml.etree.ElementTree import Element, SubElement, tostring
from xml.dom.minidom import parseString

from ..utils import sf_cli, state, paths

logger = logging.getLogger(__name__)

API_VERSION = "62.0"

METADATA_BATCHES: list[dict[str, Any]] = [
    {
        "name": "code",
        "label": "Code (Apex, LWC, Aura, VF)",
        "types": [
            "ApexClass", "ApexTrigger", "ApexComponent", "ApexPage",
            "ApexTestSuite", "AuraDefinitionBundle",
            "LightningComponentBundle",
        ],
    },
    {
        "name": "automation",
        "label": "Automation (Flows, Workflows, Validation Rules)",
        "types": [
            "Flow", "WorkflowRule", "ValidationRule",
            "CustomLabel", "CustomPermission",
        ],
    },
    {
        "name": "objects",
        "label": "Objects (CustomObject wildcard)",
        "types": ["CustomObject"],
    },
    {
        "name": "ui_resources",
        "label": "UI & Resources (Static Resources, Layouts, Tabs)",
        "types": [
            "StaticResource", "FlexiPage", "Layout",
            "CustomTab", "CustomApplication",
        ],
    },
    {
        "name": "security",
        "label": "Security & Integration",
        "types": [
            "PermissionSet", "RemoteSiteSetting", "NamedCredential",
            "ConnectedApp", "CustomMetadata",
        ],
    },
]

TOOL_DEFINITION = {
    "name": "get_customer_metadata",
    "description": (
        "Pull customer metadata (Apex, LWC, Aura, VF, objects, flows, "
        "validation rules, triggers, static resources, permissions, etc.) "
        "from the Salesforce org into the local SFDX project. "
        "Uses adaptive batching to handle the 10,000 file retrieval limit. "
        "Must be called AFTER get_installed_packages completes."
    ),
    "inputSchema": {
        "type": "object",
        "properties": {
            "working_directory": {
                "type": "string",
                "description": (
                    "Absolute path to the assessment working directory "
                    "(same value used in init_assessment)."
                ),
            },
            "org_alias": {
                "type": "string",
                "description": "Salesforce org alias to retrieve from.",
            },
            "customer_name": {
                "type": "string",
                "description": (
                    "Customer name (same value used in init_assessment). "
                    "Used to locate the SFDX project directory."
                ),
            },
        },
        "required": ["working_directory", "org_alias", "customer_name"],
    },
}


def _generate_package_xml(
    metadata_types: list[str],
    output_path: Path,
    api_version: str = API_VERSION,
) -> Path:
    """Generate a package.xml file with wildcard members for given types."""
    root = Element("Package")
    root.set("xmlns", "http://soap.sforce.com/2006/04/metadata")

    for md_type in metadata_types:
        types_elem = SubElement(root, "types")
        members = SubElement(types_elem, "members")
        members.text = "*"
        name = SubElement(types_elem, "name")
        name.text = md_type

    version = SubElement(root, "version")
    version.text = api_version

    raw_xml = tostring(root, encoding="unicode")
    pretty = parseString(raw_xml).toprettyxml(indent="    ", encoding="UTF-8")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(pretty)
    return output_path


def _is_limit_exceeded(error: sf_cli.SfCliError) -> bool:
    """Check if the error is a Salesforce retrieval limit exceeded."""
    msg = str(error).upper()
    return "LIMIT_EXCEEDED" in msg or "TOO MANY FILES" in msg


def _retrieve_batch(
    batch_name: str,
    metadata_types: list[str],
    manifest_dir: Path,
    project_dir: Path,
    org_alias: str,
    report: list[str],
) -> dict[str, Any]:
    """Attempt to retrieve a batch. Returns result dict with status and counts.

    If LIMIT_EXCEEDED, performs recursive binary split.
    """
    result = {
        "batch": batch_name,
        "types": metadata_types,
        "status": "pending",
        "file_count": 0,
        "sub_batches": [],
    }

    manifest_path = manifest_dir / f"package-{batch_name}.xml"
    _generate_package_xml(metadata_types, manifest_path)

    try:
        data = sf_cli.retrieve_metadata(
            manifest_path, org_alias, cwd=project_dir
        )
        files = data.get("result", {}).get("files", [])
        file_count = len(files) if isinstance(files, list) else 0
        result["status"] = "success"
        result["file_count"] = file_count
        report.append(
            f"  - **{batch_name}**: {', '.join(metadata_types)} "
            f"-> {file_count} files retrieved"
        )
        return result

    except sf_cli.SfCliError as e:
        if _is_limit_exceeded(e) and len(metadata_types) > 1:
            report.append(
                f"  - **{batch_name}**: LIMIT_EXCEEDED with "
                f"{len(metadata_types)} types, splitting..."
            )
            mid = len(metadata_types) // 2
            left_types = metadata_types[:mid]
            right_types = metadata_types[mid:]

            left_result = _retrieve_batch(
                f"{batch_name}-a", left_types,
                manifest_dir, project_dir, org_alias, report,
            )
            right_result = _retrieve_batch(
                f"{batch_name}-b", right_types,
                manifest_dir, project_dir, org_alias, report,
            )

            result["status"] = "split"
            result["sub_batches"] = [left_result, right_result]
            result["file_count"] = (
                left_result["file_count"] + right_result["file_count"]
            )
            return result

        elif _is_limit_exceeded(e) and len(metadata_types) == 1:
            report.append(
                f"  - **{batch_name}**: WARNING -- single type "
                f"'{metadata_types[0]}' exceeds 10K files. "
                f"Retrieved what was possible."
            )
            result["status"] = "partial"
            result["error"] = str(e)
            return result

        else:
            report.append(
                f"  - **{batch_name}**: FAILED -- {e}"
            )
            result["status"] = "failed"
            result["error"] = str(e)
            return result


async def run(arguments: dict[str, Any]) -> str:
    """Execute the get_customer_metadata tool."""
    working_dir = Path(arguments["working_directory"])
    org_alias = arguments["org_alias"]
    customer = arguments["customer_name"]

    report: list[str] = []
    warnings: list[str] = []

    # ── 0. Verify prerequisites ───────────────────────────────────
    if not state.is_step_completed(working_dir, "get_packages"):
        return (
            "## BLOCKER: get_installed_packages not completed\n\n"
            "You must call get_installed_packages first.\n"
        )

    # ── 1. Resume check ──────────────────────────────────────────
    if state.is_step_completed(working_dir, "get_metadata"):
        current_state = state.load_state(working_dir)
        metadata_info = current_state.get("metadata_retrieved", {}) if current_state else {}
        total_files = metadata_info.get("total_files", "unknown")
        return (
            "## Step already completed\n\n"
            f"Metadata retrieval was already run. "
            f"Total files retrieved: {total_files}.\n\n"
            "To re-run, manually remove 'get_metadata' from completed_steps "
            "in the state file and call this tool again.\n\n"
            "Next step: call get_managed_package_objects to retrieve "
            "CS managed package objects (not included in wildcard retrieval)."
        )

    state.mark_step_started(working_dir, "get_metadata")

    # ── 2. Setup paths ────────────────────────────────────────────
    project_dir = working_dir / customer
    manifest_dir = paths.ensure_dir(project_dir / "manifest")

    if not (project_dir / "sfdx-project.json").exists():
        state.mark_step_failed(
            working_dir, "get_metadata",
            f"SFDX project not found at {project_dir}"
        )
        return (
            f"## BLOCKER: No SFDX project at {project_dir}\n\n"
            "Run init_assessment first to create the project.\n"
        )

    # ── 3. Retrieve metadata in batches ───────────────────────────
    report.append("## Metadata Retrieval\n")
    report.append(
        f"Retrieving customer metadata from org **{org_alias}** "
        f"in {len(METADATA_BATCHES)} batches.\n"
    )

    all_batch_results = []
    total_files = 0
    failed_batches = []

    for batch_def in METADATA_BATCHES:
        report.append(f"\n### Batch: {batch_def['label']}\n")
        result = _retrieve_batch(
            batch_def["name"],
            batch_def["types"],
            manifest_dir,
            project_dir,
            org_alias,
            report,
        )
        all_batch_results.append(result)
        total_files += result["file_count"]
        if result["status"] == "failed":
            failed_batches.append(result)

    # ── 4. Summary ────────────────────────────────────────────────
    report.append(f"\n### Retrieval Summary\n")
    report.append(f"- **Total files retrieved:** {total_files}")
    report.append(f"- **Batches:** {len(METADATA_BATCHES)}")
    report.append(
        f"- **Successful:** "
        f"{sum(1 for r in all_batch_results if r['status'] in ('success', 'split'))}"
    )
    if failed_batches:
        report.append(f"- **Failed:** {len(failed_batches)}")
        for fb in failed_batches:
            warnings.append(
                f"Batch '{fb['batch']}' failed: {fb.get('error', 'unknown')}. "
                f"Types: {', '.join(fb['types'])}"
            )

    # ── 5. Write batch results ────────────────────────────────────
    output_dir = paths.ensure_dir(working_dir / "analysis")
    results_file = output_dir / "metadata-retrieval-results.json"
    results_file.write_text(json.dumps(all_batch_results, indent=2))
    report.append(f"\n- Batch results: `{results_file}`")

    # ── 6. Update state ───────────────────────────────────────────
    if not failed_batches:
        state.mark_step_completed(working_dir, "get_metadata")
        state.update_state(
            working_dir,
            metadata_retrieved={
                "total_files": total_files,
                "batches": len(all_batch_results),
                "batch_strategy": [
                    {
                        "name": r["batch"],
                        "types": r["types"],
                        "status": r["status"],
                        "file_count": r["file_count"],
                    }
                    for r in all_batch_results
                ],
            },
        )
    else:
        state.mark_step_failed(
            working_dir, "get_metadata",
            f"{len(failed_batches)} batch(es) failed"
        )

    # ── 7. Status ─────────────────────────────────────────────────
    report.append("\n---\n")
    if failed_batches:
        report.append("## STATUS: PARTIAL - Some batches failed\n")
        report.append(
            "Review the failures above. You can re-run this tool "
            "to retry failed batches.\n"
        )
    else:
        report.append("## STATUS: READY\n")
        report.append(
            f"Successfully retrieved {total_files} metadata files "
            f"from {org_alias}.\n\n"
            "**IMPORTANT:** Wildcard CustomObject retrieval does NOT include "
            "managed package objects (CS namespace objects). "
            "You MUST call `get_managed_package_objects` next to retrieve those.\n\n"
            "**Next step:** call `get_managed_package_objects`."
        )

    if warnings:
        report.append("\n## Warnings\n")
        for w in warnings:
            report.append(f"- {w}")

    return "\n".join(report)
