"""retrieve_metadata -- Retrieve metadata components from a customer org.

Retrieves ONE batch of metadata per call. The AI controls chunking,
parallelism, and retry. Updates metadata-manifest.json for resumability.
"""

import json
import logging
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from ..utils import sf_cli, state

logger = logging.getLogger(__name__)

API_VERSION = "61.0"

TOOL_DEFINITION = {
    "name": "retrieve_metadata",
    "description": (
        "Retrieve metadata components from a connected customer sandbox. "
        "Accepts explicit member names per type OR wildcard. AI controls "
        "chunking -- call in parallel for independent batches. Updates "
        "metadata-manifest.json for resumability. Requires "
        "connect_customer_org first."
    ),
    "inputSchema": {
        "type": "object",
        "properties": {
            "account_name": {
                "type": "string",
                "description": "Customer account name (e.g. 'NBN Co Ltd').",
            },
            "metadata_types": {
                "type": "array",
                "items": {"type": "string"},
                "description": (
                    "Types to retrieve. e.g. ['ApexClass'] or "
                    "['ApexTrigger', 'Flow']. Multiple small types can "
                    "be batched in one call."
                ),
            },
            "members": {
                "type": "object",
                "description": (
                    "Specific members per type. Keys are type names, values "
                    "are arrays of member names OR the string '*' for wildcard. "
                    "If omitted entirely, wildcard is used for all types. "
                    "Example: {\"ApexClass\": [\"Foo\", \"Bar\"], "
                    "\"ApexTrigger\": \"*\"}"
                ),
            },
            "working_directory": {
                "type": "string",
                "description": "Workspace directory. The AI passes this automatically.",
            },
        },
        "required": ["account_name", "metadata_types"],
    },
}


def _make_slug(name: str) -> str:
    slug = name.lower().strip()
    slug = re.sub(r"[^a-z0-9]+", "-", slug)
    return slug.strip("-") or "unknown"


def _load_org_connection(working_dir: Path, slug: str) -> dict[str, Any] | None:
    conn_path = working_dir / "customers" / slug / "org-connection.json"
    if not conn_path.exists():
        return None
    try:
        return json.loads(conn_path.read_text())
    except (json.JSONDecodeError, OSError):
        return None


def _build_package_xml(
    metadata_types: list[str],
    members: dict[str, Any] | None,
) -> str:
    """Build a package.xml string from types and optional member lists."""
    lines = [
        '<?xml version="1.0" encoding="UTF-8"?>',
        '<Package xmlns="http://soap.sforce.com/2006/04/metadata">',
    ]

    for mtype in metadata_types:
        lines.append("  <types>")

        type_members = None
        if members:
            type_members = members.get(mtype)

        if type_members and isinstance(type_members, list):
            for m in type_members:
                safe_name = m.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
                lines.append(f"    <members>{safe_name}</members>")
        else:
            lines.append("    <members>*</members>")

        lines.append(f"    <name>{mtype}</name>")
        lines.append("  </types>")

    lines.append(f"  <version>{API_VERSION}</version>")
    lines.append("</Package>")

    return "\n".join(lines)


def _count_retrieved_files(metadata_dir: Path, metadata_types: list[str]) -> dict[str, int]:
    """Count files retrieved per metadata type in force-app."""
    force_app = metadata_dir / "force-app" / "main" / "default"
    counts: dict[str, int] = {}

    type_to_dir = {
        "ApexClass": "classes",
        "ApexTrigger": "triggers",
        "Flow": "flows",
        "CustomObject": "objects",
        "ApexPage": "pages",
        "AuraDefinitionBundle": "aura",
        "LightningComponentBundle": "lwc",
        "CustomField": "objects",
        "ApexComponent": "components",
        "StaticResource": "staticresources",
        "PermissionSet": "permissionsets",
        "Profile": "profiles",
        "Layout": "layouts",
        "FlexiPage": "flexipages",
        "CustomLabel": "labels",
        "CustomMetadata": "customMetadata",
        "ValidationRule": "objects",
        "RecordType": "objects",
    }

    for mtype in metadata_types:
        dir_name = type_to_dir.get(mtype, mtype.lower() + "s")
        type_dir = force_app / dir_name
        if type_dir.exists():
            file_count = sum(1 for f in type_dir.rglob("*") if f.is_file())
            counts[mtype] = file_count
        else:
            counts[mtype] = 0

    return counts


def _update_manifest(
    manifest_path: Path,
    metadata_types: list[str],
    members_requested: dict[str, Any] | None,
    counts: dict[str, int],
) -> dict[str, Any]:
    """Update the metadata-manifest.json with retrieval results."""
    manifest: dict[str, Any] = {}
    if manifest_path.exists():
        try:
            manifest = json.loads(manifest_path.read_text())
        except (json.JSONDecodeError, OSError):
            manifest = {}

    if "types" not in manifest:
        manifest["types"] = {}
    if "retrieval_log" not in manifest:
        manifest["retrieval_log"] = []

    now = datetime.now(timezone.utc).isoformat()

    for mtype in metadata_types:
        was_wildcard = True
        if members_requested and mtype in members_requested:
            val = members_requested[mtype]
            if isinstance(val, list):
                was_wildcard = False

        existing = manifest["types"].get(mtype, {})
        prev_members = set(existing.get("retrieved_members", []))
        prev_count = existing.get("files_retrieved", 0)

        if was_wildcard:
            manifest["types"][mtype] = {
                "status": "complete",
                "strategy": "wildcard",
                "files_retrieved": counts.get(mtype, prev_count),
                "last_retrieved": now,
            }
        else:
            new_members = members_requested[mtype]
            all_members = sorted(prev_members | set(new_members))
            manifest["types"][mtype] = {
                "status": "partial" if existing else "partial",
                "strategy": "chunked",
                "files_retrieved": counts.get(mtype, 0),
                "retrieved_members": all_members,
                "chunks_done": existing.get("chunks_done", 0) + 1,
                "last_retrieved": now,
            }

    manifest["retrieval_log"].append({
        "timestamp": now,
        "types": metadata_types,
        "files_per_type": counts,
        "strategy": "wildcard" if not members_requested else "explicit",
    })
    manifest["last_updated"] = now

    manifest_path.write_text(json.dumps(manifest, indent=2))
    return manifest


async def run(arguments: dict[str, Any]) -> str:
    """Execute the retrieve_metadata tool."""
    state.set_working_dir(arguments.get("working_directory"))
    working_dir = state.get_working_dir()

    account_name = arguments.get("account_name", "").strip()
    metadata_types = arguments.get("metadata_types", [])
    members = arguments.get("members")

    if not account_name:
        return "STATUS: MISSING_CUSTOMER\n\naccount_name is required."
    if not metadata_types:
        return "STATUS: MISSING_TYPES\n\nmetadata_types array is required."

    slug = _make_slug(account_name)
    conn = _load_org_connection(working_dir, slug)

    if not conn:
        return (
            "STATUS: NOT_CONNECTED\n\n"
            f"No org-connection.json found for `{account_name}`.\n"
            "Call `connect_customer_org` first to authorize a sandbox."
        )

    org_alias = conn.get("org_alias", "")
    if not org_alias:
        return (
            "STATUS: NOT_CONNECTED\n\n"
            "org-connection.json exists but has no org_alias.\n"
            "Call `connect_customer_org` again."
        )

    metadata_dir = working_dir / "customers" / slug / "metadata"
    metadata_dir.mkdir(parents=True, exist_ok=True)

    sfdx_project = metadata_dir / "sfdx-project.json"
    if not sfdx_project.exists():
        sfdx_project.write_text(json.dumps({
            "packageDirectories": [{"path": "force-app", "default": True}],
            "namespace": "",
            "sfdcLoginUrl": "https://test.salesforce.com",
            "sourceApiVersion": API_VERSION,
        }, indent=2))
    (metadata_dir / "force-app" / "main" / "default").mkdir(
        parents=True, exist_ok=True
    )

    package_xml = _build_package_xml(metadata_types, members)
    package_xml_path = metadata_dir / "package.xml"
    package_xml_path.write_text(package_xml)

    total_members_in_call = 0
    if members:
        for mtype in metadata_types:
            val = members.get(mtype)
            if isinstance(val, list):
                total_members_in_call += len(val)

    try:
        sf_cli.run_sf(
            ["project", "retrieve", "start",
             "--manifest", str(package_xml_path),
             "--target-org", org_alias],
            cwd=metadata_dir,
            timeout=180,
        )
    except sf_cli.SfCliError as e:
        error_str = str(e)

        if "LIMIT_EXCEEDED" in error_str or "10000" in error_str:
            return (
                "STATUS: LIMIT_EXCEEDED\n\n"
                f"Salesforce 10K file limit hit for types: {metadata_types}\n"
                f"Members requested: {total_members_in_call or 'wildcard'}\n\n"
                "**Action:** Chunk into smaller batches. Use member lists "
                "from `list_metadata` discovery files to create batches of "
                "~2,000 members."
            )

        if "expired" in error_str.lower() or "invalid_session" in error_str.lower():
            return (
                "STATUS: SESSION_EXPIRED\n\n"
                "The org session has expired.\n"
                "Call `connect_customer_org` again to re-authorize, "
                "then retry this retrieval."
            )

        return (
            f"STATUS: RETRIEVE_FAILED\n\n"
            f"Retrieval failed: {e}\n\n"
            "If the error mentions a timeout, try retrieving fewer types "
            "or smaller chunks per call."
        )

    counts = _count_retrieved_files(metadata_dir, metadata_types)
    manifest_path = metadata_dir / "metadata-manifest.json"
    manifest = _update_manifest(manifest_path, metadata_types, members, counts)

    report = [f"STATUS: READY\n"]
    report.append("## Retrieved Metadata\n")

    total_files = 0
    for mtype in metadata_types:
        c = counts.get(mtype, 0)
        total_files += c
        strategy = "wildcard"
        if members and mtype in members and isinstance(members[mtype], list):
            strategy = f"{len(members[mtype])} specific members"
        report.append(f"- **{mtype}**: {c:,} files ({strategy})")

    report.append(f"\n**Total files:** {total_files:,}")
    report.append(f"**Location:** `customers/{slug}/metadata/force-app/`")
    report.append(f"**Manifest:** `customers/{slug}/metadata/metadata-manifest.json`")

    types_status = manifest.get("types", {})
    complete_types = [t for t, v in types_status.items() if v.get("status") == "complete"]
    partial_types = [t for t, v in types_status.items() if v.get("status") == "partial"]

    if complete_types:
        report.append(f"\nComplete types: {', '.join(complete_types)}")
    if partial_types:
        report.append(f"Partial types (more chunks needed): {', '.join(partial_types)}")

    return "\n".join(report)
