"""
Step package validation and certificate verification.

This module provides the core validation logic for step packages,
checking contract compliance, certificate validity, and code integrity.
"""

import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple

from .crypto import (
    compute_distribution_hash,
    verify_signature,
    sign_data,
    create_signable_json,
)


# Certificate format version
CERT_FORMAT_VERSION = "cf-pipeline-cert-v2"

CF_NS = "https://cogniflow.org/ns#"
CF_NS_HTTP = "http://cogniflow.org/ns#"

PLUGIN_ARTIFACT_TYPES = {
    "cf:PluginArtifact",
    f"{CF_NS}PluginArtifact",
    f"{CF_NS_HTTP}PluginArtifact",
}
STEP_IMPLEMENTATION_TYPES = {
    "cf:StepImplementation",
    f"{CF_NS}StepImplementation",
    f"{CF_NS_HTTP}StepImplementation",
}

IMPLEMENTATION_ARTIFACT_KEYS = (
    "cf:implementationArtifact",
    f"{CF_NS}implementationArtifact",
    f"{CF_NS_HTTP}implementationArtifact",
)
OPERATION_NAME_KEYS = (
    "cf:operationName",
    f"{CF_NS}operationName",
    f"{CF_NS_HTTP}operationName",
)
OPERATION_CONTRACT_KEYS = (
    "cf:operationContract",
    f"{CF_NS}operationContract",
    f"{CF_NS_HTTP}operationContract",
)
PLUGIN_NAME_KEYS = (
    "cf:pluginName",
    f"{CF_NS}pluginName",
    f"{CF_NS_HTTP}pluginName",
    "schema:name",
    "https://schema.org/name",
)
PLUGIN_VERSION_KEYS = (
    "cf:pluginVersion",
    f"{CF_NS}pluginVersion",
    f"{CF_NS_HTTP}pluginVersion",
    "schema:version",
    "https://schema.org/version",
)
PLUGIN_CERTIFICATE_KEYS = (
    "cf:pluginCertificate",
    f"{CF_NS}pluginCertificate",
    f"{CF_NS_HTTP}pluginCertificate",
)
ARTIFACT_PATH_KEYS = (
    "cf:artifactPath",
    f"{CF_NS}artifactPath",
    f"{CF_NS_HTTP}artifactPath",
)
PLUGIN_LIST_SYMBOL_KEYS = (
    "cf:pluginListSymbol",
    f"{CF_NS}pluginListSymbol",
    f"{CF_NS_HTTP}pluginListSymbol",
)
PLUGIN_RESOLVE_SYMBOL_KEYS = (
    "cf:pluginResolveSymbol",
    f"{CF_NS}pluginResolveSymbol",
    f"{CF_NS_HTTP}pluginResolveSymbol",
)


@dataclass
class OperationInfo:
    """Information about a single operation in a step package."""
    name: str
    inputs: List[Dict[str, Any]]
    outputs: List[Dict[str, Any]]
    parameters: List[Dict[str, Any]] = field(default_factory=list)
    description: str = ""


@dataclass
class ValidationResult:
    """Result of validating a step package."""
    valid: bool
    plugin_name: str = ""
    plugin_version: str = ""
    operations: List[OperationInfo] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    distribution_hash: str = ""
    
    def to_certificate(self, sign_key_b64: Optional[str] = None) -> Dict[str, Any]:
        """
        Convert validation result to a certificate.
        
        Args:
            sign_key_b64: Optional Ed25519 private key for signing.
                         If not provided, certificate will have signature=null
                         (suitable for certification requests).
            self_signed: If True, marks certificate as self-signed (local trust only).
                        Self-signed certificates include the public key for verification.
        
        Returns:
            Certificate dictionary ready for embedding in steps.jsonld (cf:pluginCertificate)
        """
        # Determine trust level BEFORE creating the cert dict
        trust_level = "ca-signed" if sign_key_b64 else "unsigned"
        
        cert = {
            "format": CERT_FORMAT_VERSION,
            "plugin_name": self.plugin_name,
            "plugin_version": self.plugin_version,
            "distribution_hash": self.distribution_hash,
            "certified_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "operations": [
                {
                    "name": op.name,
                    "inputs": op.inputs,
                    "outputs": op.outputs,
                    "parameters": op.parameters,
                }
                for op in self.operations
            ],
            "signature": None,
            "trust_level": trust_level,
            "signer_public_key": None,  # Only set for self-signed
        }
        
        if sign_key_b64:
            signable = create_signable_json(cert)
            cert["signature"] = sign_data(signable, sign_key_b64)
        
        return cert
    
    def to_self_signed_certificate(self) -> Dict[str, Any]:
        """
        Create a self-signed certificate with embedded keypair.
        
        Self-signed certificates:
        - Include the public key so anyone can verify the hash
        - Are marked as "self-signed" trust level
        - Don't require CA involvement
        - Still protect against tampering after creation
        
        Returns:
            Certificate with embedded public key for self-verification
        """
        from .crypto import generate_keypair
        
        private_key, public_key = generate_keypair()
        
        cert = {
            "format": CERT_FORMAT_VERSION,
            "plugin_name": self.plugin_name,
            "plugin_version": self.plugin_version,
            "distribution_hash": self.distribution_hash,
            "certified_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "operations": [
                {
                    "name": op.name,
                    "inputs": op.inputs,
                    "outputs": op.outputs,
                    "parameters": op.parameters,
                }
                for op in self.operations
            ],
            "signature": None,
            "trust_level": "self-signed",
            "signer_public_key": public_key,
        }
        
        signable = create_signable_json(cert)
        cert["signature"] = sign_data(signable, private_key)
        
        return cert


@dataclass
class CertificateInfo:
    """Information extracted from a certificate."""
    format_version: str
    plugin_name: str
    plugin_version: str
    distribution_hash: str
    certified_at: str
    operations: List[Dict[str, Any]]
    signature: Optional[str]
    is_signed: bool
    trust_level: str  # "unsigned", "self-signed", "ca-signed"
    signer_public_key: Optional[str]  # Only for self-signed
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CertificateInfo":
        """Parse certificate from dictionary."""
        return cls(
            format_version=data.get("format", "unknown"),
            plugin_name=data.get("plugin_name", ""),
            plugin_version=data.get("plugin_version", ""),
            distribution_hash=data.get("distribution_hash", ""),
            certified_at=data.get("certified_at", ""),
            operations=data.get("operations", []),
            signature=data.get("signature"),
            is_signed=data.get("signature") is not None,
            trust_level=data.get("trust_level", "unsigned"),
            signer_public_key=data.get("signer_public_key"),
        )


def _iter_nodes(doc: Dict[str, Any]) -> List[Dict[str, Any]]:
    graph = doc.get("@graph")
    if isinstance(graph, list):
        return [node for node in graph if isinstance(node, dict)]
    if isinstance(doc, dict):
        return [doc]
    return []


def _has_type(node: Dict[str, Any], candidates: set[str]) -> bool:
    raw = node.get("@type")
    types: List[str] = []
    if isinstance(raw, str):
        types.append(raw)
    elif isinstance(raw, list):
        for entry in raw:
            if isinstance(entry, str):
                types.append(entry)
            elif isinstance(entry, dict):
                entry_id = entry.get("@id")
                if isinstance(entry_id, str):
                    types.append(entry_id)
    elif isinstance(raw, dict):
        entry_id = raw.get("@id")
        if isinstance(entry_id, str):
            types.append(entry_id)
    return any(t in candidates for t in types)


def _extract_text(value: Any) -> Optional[str]:
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        if isinstance(value.get("@value"), str):
            return value.get("@value")
        if isinstance(value.get("@id"), str):
            return value.get("@id")
    if isinstance(value, list):
        for item in value:
            text = _extract_text(item)
            if text:
                return text
    return None


def _extract_text_list(value: Any) -> List[str]:
    values: List[str] = []
    if isinstance(value, str):
        values.append(value)
    elif isinstance(value, dict):
        entry_value = value.get("@value")
        if isinstance(entry_value, str):
            values.append(entry_value)
        entry_id = value.get("@id")
        if isinstance(entry_id, str):
            values.append(entry_id)
    elif isinstance(value, list):
        for item in value:
            values.extend(_extract_text_list(item))
    return values


def _extract_id_values(value: Any) -> List[str]:
    values: List[str] = []
    if isinstance(value, str):
        values.append(value)
    elif isinstance(value, dict):
        entry_id = value.get("@id")
        if isinstance(entry_id, str):
            values.append(entry_id)
    elif isinstance(value, list):
        for item in value:
            values.extend(_extract_id_values(item))
    return values


def _first_property(node: Dict[str, Any], keys: Tuple[str, ...]) -> Any:
    for key in keys:
        if key in node:
            return node.get(key)
    return None


def _extract_json_value(value: Any) -> Optional[Dict[str, Any]]:
    if isinstance(value, dict):
        if "@value" in value and isinstance(value.get("@value"), dict):
            return value.get("@value")
        return value
    if isinstance(value, list):
        for item in value:
            extracted = _extract_json_value(item)
            if extracted is not None:
                return extracted
    return None


def _find_steps_jsonld(plugin_path: Path) -> Tuple[Optional[Path], Optional[str]]:
    path = Path(plugin_path).resolve()
    if path.is_file():
        if path.name.endswith(".jsonld") or path.name.endswith(".json"):
            return path, None
        return None, f"Expected a steps.jsonld file or package directory, got {path}"

    candidate = path / "steps.jsonld"
    if candidate.exists():
        return candidate, None

    matches = list(path.rglob("steps.jsonld"))
    if len(matches) == 1:
        return matches[0], None
    if not matches:
        return None, f"Missing steps.jsonld in {path}"

    match_list = ", ".join(str(match) for match in matches)
    return None, f"Multiple steps.jsonld files found under {path}: {match_list}"


def load_steps_jsonld(plugin_path: Path) -> Tuple[Path, Dict[str, Any]]:
    steps_path, error = _find_steps_jsonld(plugin_path)
    if error or steps_path is None:
        raise FileNotFoundError(error or "steps.jsonld not found")
    try:
        document = json.loads(steps_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:  # pragma: no cover - should be caught by tests
        raise ValueError(f"Invalid JSON in steps.jsonld: {exc}") from exc
    return steps_path, document


def _select_plugin_artifact(
    nodes: List[Dict[str, Any]]
) -> Tuple[Optional[Dict[str, Any]], Optional[str], List[Dict[str, Any]], Optional[str]]:
    plugin_nodes: Dict[str, Dict[str, Any]] = {}
    for node in nodes:
        node_id = node.get("@id")
        if not isinstance(node_id, str):
            continue
        if _has_type(node, PLUGIN_ARTIFACT_TYPES):
            plugin_nodes[node_id] = node

    impl_nodes = [node for node in nodes if _has_type(node, STEP_IMPLEMENTATION_TYPES)]

    referenced: set[str] = set()
    for impl in impl_nodes:
        artifact_refs = _extract_id_values(_first_property(impl, IMPLEMENTATION_ARTIFACT_KEYS))
        for ref in artifact_refs:
            if ref in plugin_nodes:
                referenced.add(ref)

    plugin_ids = sorted(referenced or plugin_nodes.keys())

    if not plugin_ids:
        return None, None, impl_nodes, "Missing PluginArtifact in steps.jsonld"
    if len(plugin_ids) > 1:
        return None, None, impl_nodes, f"Multiple PluginArtifact entries found: {plugin_ids}"

    plugin_id = plugin_ids[0]
    return plugin_nodes[plugin_id], plugin_id, impl_nodes, None


def select_plugin_artifact(doc: Dict[str, Any]) -> Tuple[Dict[str, Any], str]:
    node, plugin_id, _, error = _select_plugin_artifact(_iter_nodes(doc))
    if error or node is None or plugin_id is None:
        raise ValueError(error or "PluginArtifact not found")
    return node, plugin_id


def validate_step_package(plugin_path: Path) -> ValidationResult:
    """
    Validate a step package against the cf-pipeline contract.
    
    This checks:
    - steps.jsonld exists and has required metadata
    - Plugin artifact metadata is present
    - Operations have required fields (name, inputs, outputs) or contract
    - Input/output specifications are valid
    
    Args:
        plugin_path: Path to the step package root directory
        
    Returns:
        ValidationResult with validation status and details
    """
    result = ValidationResult(valid=False)

    steps_path, error = _find_steps_jsonld(plugin_path)
    if error or steps_path is None:
        result.errors.append(error or "Missing steps.jsonld")
        return result

    try:
        steps_doc = json.loads(steps_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        result.errors.append(f"Invalid JSON in steps.jsonld: {exc}")
        return result

    nodes = _iter_nodes(steps_doc)
    plugin_node, plugin_id, impl_nodes, error = _select_plugin_artifact(nodes)
    if error or plugin_node is None or plugin_id is None:
        result.errors.append(error or "Missing PluginArtifact")
        return result

    plugin_name = _extract_text(_first_property(plugin_node, PLUGIN_NAME_KEYS))
    plugin_version = _extract_text(_first_property(plugin_node, PLUGIN_VERSION_KEYS))
    artifact_paths = _extract_text_list(_first_property(plugin_node, ARTIFACT_PATH_KEYS))
    list_symbol = _extract_text(_first_property(plugin_node, PLUGIN_LIST_SYMBOL_KEYS))
    resolve_symbol = _extract_text(_first_property(plugin_node, PLUGIN_RESOLVE_SYMBOL_KEYS))

    missing_fields: List[str] = []
    if not plugin_name:
        missing_fields.append("plugin_name")
    if not plugin_version:
        missing_fields.append("plugin_version")
    if not artifact_paths:
        missing_fields.append("artifact_path")
    if not list_symbol:
        missing_fields.append("list_symbol")
    if not resolve_symbol:
        missing_fields.append("resolve_symbol")

    if missing_fields:
        result.errors.append(f"Missing required plugin fields: {missing_fields}")
        return result

    result.plugin_name = plugin_name or ""
    result.plugin_version = plugin_version or ""

    operations: List[OperationInfo] = []
    for impl in impl_nodes:
        artifact_refs = _extract_id_values(_first_property(impl, IMPLEMENTATION_ARTIFACT_KEYS))
        if plugin_id not in artifact_refs:
            continue

        op_name = _extract_text(_first_property(impl, OPERATION_NAME_KEYS))
        if not op_name:
            continue

        contract_value = _extract_json_value(_first_property(impl, OPERATION_CONTRACT_KEYS))
        inputs = contract_value.get("inputs", []) if isinstance(contract_value, dict) else []
        outputs = contract_value.get("outputs", []) if isinstance(contract_value, dict) else []
        parameters = contract_value.get("parameters", []) if isinstance(contract_value, dict) else []
        description = _extract_text(
            impl.get("cf:description")
            or impl.get(f"{CF_NS}description")
            or impl.get(f"{CF_NS_HTTP}description")
        ) or ""

        has_io_fields = bool(inputs) or bool(outputs)
        has_contract = contract_value is not None

        if not has_io_fields and not has_contract:
            result.warnings.append(
                f"Operation '{op_name}' has minimal spec (no inputs/outputs/contract)"
            )

        for j, inp in enumerate(inputs):
            if isinstance(inp, dict) and "name" not in inp:
                result.errors.append(f"Operation '{op_name}' input {j} missing 'name'")
            if isinstance(inp, dict) and "type" not in inp:
                result.warnings.append(
                    f"Operation '{op_name}' input '{inp.get('name', j)}' missing 'type'"
                )

        for j, out in enumerate(outputs):
            if isinstance(out, dict) and "name" not in out:
                result.errors.append(f"Operation '{op_name}' output {j} missing 'name'")
            if isinstance(out, dict) and "type" not in out:
                result.warnings.append(
                    f"Operation '{op_name}' output '{out.get('name', j)}' missing 'type'"
                )

        operations.append(
            OperationInfo(
                name=op_name,
                inputs=inputs if isinstance(inputs, list) else [],
                outputs=outputs if isinstance(outputs, list) else [],
                parameters=parameters if isinstance(parameters, list) else [],
                description=description,
            )
        )

    if not operations:
        result.errors.append("No operations defined in steps.jsonld (missing cf:operationName)")
        return result

    result.operations = operations

    if not result.errors:
        result.valid = True
        result.distribution_hash = compute_distribution_hash(steps_path.parent)

    return result


def verify_certificate(
    plugin_path: Path,
    public_key_b64: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Verify the certificate embedded in a step package.
    
    This checks:
    1. Certificate exists in steps.jsonld (cf:pluginCertificate)
    2. Certificate format is recognized
    3. Signature is valid (if public_key provided)
    4. Distribution hash matches current code
    
    Args:
        plugin_path: Path to the step package root directory
        public_key_b64: Ed25519 public key for signature verification.
                       If not provided, only hash verification is performed.
    
    Returns:
        Dictionary with verification results:
        {
            "valid": bool,
            "certificate": CertificateInfo or None,
            "signature_valid": bool or None,
            "hash_matches": bool or None,
            "current_hash": str,
            "errors": List[str],
        }
    """
    result = {
        "valid": False,
        "certificate": None,
        "signature_valid": None,
        "hash_matches": None,
        "current_hash": "",
        "errors": [],
    }

    steps_path, error = _find_steps_jsonld(plugin_path)
    if error or steps_path is None:
        result["errors"].append(error or "Missing steps.jsonld")
        return result

    try:
        steps_doc = json.loads(steps_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        result["errors"].append(f"Invalid JSON in steps.jsonld: {exc}")
        return result

    nodes = _iter_nodes(steps_doc)
    plugin_node, _, _, error = _select_plugin_artifact(nodes)
    if error or plugin_node is None:
        result["errors"].append(error or "Missing PluginArtifact")
        return result

    cert_data = _extract_json_value(_first_property(plugin_node, PLUGIN_CERTIFICATE_KEYS))
    if not cert_data:
        result["errors"].append("No certificate found in steps.jsonld")
        return result
    
    cert_info = CertificateInfo.from_dict(cert_data)
    result["certificate"] = cert_info
    
    # Check format version
    if cert_info.format_version != CERT_FORMAT_VERSION:
        result["errors"].append(
            f"Unknown certificate format: {cert_info.format_version} "
            f"(expected {CERT_FORMAT_VERSION})"
        )
        return result
    
    # Compute current distribution hash
    current_hash = compute_distribution_hash(steps_path.parent)
    result["current_hash"] = current_hash
    result["trust_level"] = cert_info.trust_level
    
    # Verify hash matches
    result["hash_matches"] = (current_hash == cert_info.distribution_hash)
    if not result["hash_matches"]:
        result["errors"].append(
            f"Distribution hash mismatch! Certificate hash: {cert_info.distribution_hash}, "
            f"Current hash: {current_hash}. Code may have been tampered with."
        )
    
    # Determine which key to use for verification
    verification_key = public_key_b64
    
    # For self-signed certificates, use the embedded public key
    if cert_info.trust_level == "self-signed" and cert_info.signer_public_key:
        verification_key = cert_info.signer_public_key
        result["self_signed"] = True
    
    # Verify signature
    if verification_key and cert_info.is_signed:
        signable = create_signable_json(cert_data)
        result["signature_valid"] = verify_signature(
            signable, cert_info.signature, verification_key
        )
        if not result["signature_valid"]:
            result["errors"].append(
                "Invalid certificate signature! Certificate may have been tampered with."
            )
    elif verification_key and not cert_info.is_signed:
        result["errors"].append("Certificate is not signed")
        result["signature_valid"] = False
    elif not verification_key and cert_info.is_signed:
        # No key available, can only verify hash
        result["signature_valid"] = None
    
    # Determine overall validity
    result["valid"] = (
        len(result["errors"]) == 0 and
        result["hash_matches"] == True and
        (result["signature_valid"] == True if verification_key else True)
    )
    
    return result


def sign_certificate_request(
    request_path: Path,
    private_key_b64: str,
    output_path: Optional[Path] = None,
) -> Dict[str, Any]:
    """
    Sign a certification request file.
    
    Args:
        request_path: Path to the .cert-request.json file
        private_key_b64: Ed25519 private key for signing
        output_path: Where to write the signed certificate (default: same name with .cert.json)
    
    Returns:
        The signed certificate dictionary
    """
    with open(request_path, "r", encoding="utf-8") as f:
        cert_data = json.load(f)
    
    # Sign the certificate
    signable = create_signable_json(cert_data)
    cert_data["signature"] = sign_data(signable, private_key_b64)
    
    # Write output
    if output_path is None:
        output_path = request_path.with_suffix("").with_suffix(".cert.json")
    
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(cert_data, f, indent=2)
    
    return cert_data
