# Copyright (c) 2025 HaoLine Contributors
# SPDX-License-Identifier: MIT

"""
JSON Schema definition and validation for HaoLine reports.

Provides:
- INSPECTION_REPORT_SCHEMA: JSON Schema for InspectionReport (auto-generated from Pydantic)
- validate_report(): Validate a report dict against the schema using Pydantic
- ValidationError: Exception raised on validation failure

Validation uses the Pydantic models from report.py (InspectionReport, etc.)
for type-safe validation with clear error messages.
"""

from __future__ import annotations

from typing import Any

from pydantic import ValidationError as PydanticValidationError

from .report import InspectionReport


class ValidationError(Exception):
    """Raised when JSON schema validation fails."""

    def __init__(self, message: str, errors: list[str] | None = None):
        super().__init__(message)
        self.errors = errors or []


# JSON Schema for InspectionReport (Draft 7)
INSPECTION_REPORT_SCHEMA: dict[str, Any] = {
    "$schema": "http://json-schema.org/draft-07/schema#",
    "$id": "https://github.com/mdayku/HaoLine/schema/inspection-report.schema.json",
    "title": "HaoLine Inspection Report",
    "description": "Schema for ONNX model inspection reports generated by HaoLine",
    "type": "object",
    "required": ["metadata", "generated_at", "autodoc_version"],
    "properties": {
        "metadata": {
            "type": "object",
            "description": "Model metadata extracted from ONNX proto",
            "required": ["path", "ir_version", "producer_name", "opsets"],
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Path to the ONNX model file",
                },
                "ir_version": {
                    "type": "integer",
                    "minimum": 1,
                    "description": "ONNX IR version",
                },
                "producer_name": {
                    "type": "string",
                    "description": "Name of the tool that produced the model",
                },
                "producer_version": {
                    "type": "string",
                    "description": "Version of the producer tool",
                },
                "domain": {"type": "string", "description": "Model domain"},
                "model_version": {
                    "type": "integer",
                    "description": "Model version number",
                },
                "doc_string": {
                    "type": "string",
                    "description": "Model documentation string",
                },
                "opsets": {
                    "type": "object",
                    "description": "Opset versions by domain",
                    "additionalProperties": {"type": "integer", "minimum": 1},
                },
            },
        },
        "generated_at": {
            "type": "string",
            "format": "date-time",
            "description": "ISO 8601 timestamp when report was generated",
        },
        "autodoc_version": {
            "type": "string",
            "pattern": "^[0-9]+\\.[0-9]+\\.[0-9]+",
            "description": "Version of HaoLine that generated the report",
        },
        "graph_summary": {
            "type": ["object", "null"],
            "description": "Summary statistics about the ONNX graph",
            "properties": {
                "num_nodes": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Total number of nodes in graph",
                },
                "num_inputs": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Number of graph inputs",
                },
                "num_outputs": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Number of graph outputs",
                },
                "num_initializers": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Number of initializers (weights)",
                },
                "input_shapes": {
                    "type": "object",
                    "description": "Input tensor shapes by name",
                    "additionalProperties": {
                        "type": "array",
                        "items": {"type": ["integer", "string"]},
                    },
                },
                "output_shapes": {
                    "type": "object",
                    "description": "Output tensor shapes by name",
                    "additionalProperties": {
                        "type": "array",
                        "items": {"type": ["integer", "string"]},
                    },
                },
                "op_type_counts": {
                    "type": "object",
                    "description": "Count of each operator type",
                    "additionalProperties": {"type": "integer", "minimum": 0},
                },
            },
        },
        "param_counts": {
            "type": ["object", "null"],
            "description": "Parameter count statistics",
            "properties": {
                "total": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Total parameter count",
                },
                "trainable": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Trainable parameter count",
                },
                "non_trainable": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Non-trainable parameter count",
                },
                "by_op_type": {
                    "type": "object",
                    "description": "Parameters by operator type (fractional for shared weights)",
                    "additionalProperties": {"type": "number", "minimum": 0},
                },
                "shared_weights": {
                    "type": "object",
                    "description": "Information about shared weights",
                    "properties": {
                        "count": {
                            "type": "integer",
                            "minimum": 0,
                            "description": "Number of weights shared across 2+ nodes",
                        },
                        "details": {
                            "type": "object",
                            "description": "Shared weight name to list of node names using it",
                            "additionalProperties": {
                                "type": "array",
                                "items": {"type": "string"},
                            },
                        },
                    },
                },
                "precision_breakdown": {
                    "type": "object",
                    "description": "Parameter count by data type (fp32, fp16, int8, etc.)",
                    "additionalProperties": {"type": "integer", "minimum": 0},
                },
                "is_quantized": {
                    "type": "boolean",
                    "description": "Whether model uses quantized weights or ops",
                },
                "quantized_ops": {
                    "type": "array",
                    "description": "List of quantized operation types detected",
                    "items": {"type": "string"},
                },
            },
        },
        "flop_counts": {
            "type": ["object", "null"],
            "description": "FLOP estimation statistics",
            "properties": {
                "total": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Total estimated FLOPs",
                },
                "by_node_type": {
                    "type": "object",
                    "description": "FLOPs by operator type",
                    "additionalProperties": {"type": "integer", "minimum": 0},
                },
                "hotspots": {
                    "type": "array",
                    "description": "Top compute-intensive nodes",
                    "items": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string"},
                            "op_type": {"type": "string"},
                            "flops": {"type": "integer", "minimum": 0},
                        },
                    },
                },
            },
        },
        "memory_estimates": {
            "type": ["object", "null"],
            "description": "Memory usage estimates",
            "properties": {
                "model_size_bytes": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Model size in bytes",
                },
                "peak_activation_bytes": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "Peak activation memory in bytes",
                },
                "kv_cache_bytes_per_token": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "KV cache memory per token (transformers)",
                },
                "kv_cache_bytes_full_context": {
                    "type": "integer",
                    "minimum": 0,
                    "description": "KV cache for full context length",
                },
                "kv_cache_config": {
                    "type": "object",
                    "description": "KV cache configuration",
                    "properties": {
                        "num_layers": {"type": "integer", "minimum": 0},
                        "hidden_dim": {"type": "integer", "minimum": 0},
                        "seq_len": {"type": "integer", "minimum": 0},
                        "bytes_per_element": {"type": "integer", "minimum": 1},
                    },
                },
                "breakdown": {
                    "type": ["object", "null"],
                    "description": "Memory breakdown by component",
                    "properties": {
                        "weights_by_op_type": {
                            "type": "object",
                            "additionalProperties": {"type": "integer", "minimum": 0},
                        },
                        "activations_by_op_type": {
                            "type": "object",
                            "additionalProperties": {"type": "integer", "minimum": 0},
                        },
                    },
                },
            },
        },
        "detected_blocks": {
            "type": "array",
            "description": "Detected architectural blocks",
            "items": {
                "type": "object",
                "required": ["block_type", "name"],
                "properties": {
                    "block_type": {
                        "type": "string",
                        "description": "Type of block (e.g., ResidualAdd, Attention)",
                    },
                    "name": {"type": "string", "description": "Block identifier"},
                    "nodes": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Node names in this block",
                    },
                    "start_node": {"type": "string"},
                    "end_node": {"type": "string"},
                    "attributes": {
                        "type": "object",
                        "description": "Block-specific attributes",
                    },
                },
            },
        },
        "architecture_type": {
            "type": "string",
            "enum": ["transformer", "cnn", "mlp", "hybrid", "unknown"],
            "description": "Detected architecture type",
        },
        "risk_signals": {
            "type": "array",
            "description": "Detected risk signals",
            "items": {
                "type": "object",
                "required": ["id", "severity", "description"],
                "properties": {
                    "id": {
                        "type": "string",
                        "description": "Risk signal identifier",
                    },
                    "severity": {
                        "type": "string",
                        "enum": ["info", "warning", "high"],
                        "description": "Severity level",
                    },
                    "description": {
                        "type": "string",
                        "description": "Human-readable description",
                    },
                    "nodes": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Affected node names",
                    },
                    "recommendation": {
                        "type": "string",
                        "description": "Recommended action",
                    },
                },
            },
        },
        "hardware_profile": {
            "type": ["object", "null"],
            "description": "Target hardware profile",
            "properties": {
                "name": {"type": "string"},
                "vram_bytes": {"type": "integer", "minimum": 0},
                "peak_fp32_tflops": {"type": "number", "minimum": 0},
                "peak_fp16_tflops": {"type": "number", "minimum": 0},
                "memory_bandwidth_gbps": {"type": "number", "minimum": 0},
                "tdp_watts": {"type": ["integer", "null"]},
            },
        },
        "hardware_estimates": {
            "type": ["object", "null"],
            "description": "Hardware-specific estimates",
            "properties": {
                "device": {"type": "string"},
                "precision": {"type": "string"},
                "batch_size": {"type": "integer", "minimum": 1},
                "vram_required_bytes": {"type": "integer", "minimum": 0},
                "fits_in_vram": {"type": "boolean"},
                "theoretical_latency_ms": {"type": "number", "minimum": 0},
                "bottleneck": {"type": "string"},
                "compute_utilization_estimate": {"type": "number", "minimum": 0},
                "gpu_saturation": {"type": "number", "minimum": 0},
            },
        },
        "llm_summary": {
            "type": ["object", "null"],
            "description": "LLM-generated summary",
            "properties": {
                "success": {"type": "boolean"},
                "short_summary": {"type": "string"},
                "detailed_summary": {"type": "string"},
                "model": {"type": "string"},
                "error": {"type": "string"},
            },
        },
        "dataset_info": {
            "type": ["object", "null"],
            "description": "Dataset and class information",
            "properties": {
                "task": {"type": ["string", "null"]},
                "num_classes": {"type": ["integer", "null"], "minimum": 0},
                "class_names": {
                    "type": "array",
                    "items": {"type": "string"},
                },
                "source": {"type": ["string", "null"]},
            },
        },
    },
}


def validate_report(report_dict: dict[str, Any]) -> tuple[bool, list[str]]:
    """
    Validate a report dictionary against the schema.

    Uses Pydantic InspectionReport model for validation.

    Args:
        report_dict: The report as a dictionary (from to_dict()).

    Returns:
        Tuple of (is_valid, list of error messages).
    """
    try:
        InspectionReport.model_validate(report_dict)
        return True, []
    except PydanticValidationError as e:
        error_messages = []
        for error in e.errors():
            loc = " -> ".join(str(x) for x in error["loc"])
            error_messages.append(f"{loc}: {error['msg']}")
        return False, error_messages


def validate_report_strict(report_dict: dict[str, Any]) -> None:
    """
    Validate a report dictionary, raising ValidationError on failure.

    Args:
        report_dict: The report as a dictionary (from to_dict()).

    Raises:
        ValidationError: If validation fails.
    """
    is_valid, errors = validate_report(report_dict)
    if not is_valid:
        raise ValidationError(
            f"Report validation failed with {len(errors)} error(s)", errors=errors
        )


def get_schema() -> dict[str, Any]:
    """
    Return the JSON schema for InspectionReport.

    Returns the auto-generated schema from the Pydantic InspectionReport model.
    Falls back to manually-defined schema if model_json_schema() fails.
    """
    try:
        schema: dict[str, Any] = InspectionReport.model_json_schema()
        return schema
    except Exception:
        return INSPECTION_REPORT_SCHEMA.copy()


def validate_with_pydantic(report_dict: dict[str, Any]) -> InspectionReport:
    """
    Validate and parse a report dict into a Pydantic InspectionReport model.

    Args:
        report_dict: The report as a dictionary.

    Returns:
        InspectionReport instance if valid.

    Raises:
        ValidationError: If validation fails.
    """
    try:
        result: InspectionReport = InspectionReport.model_validate(report_dict)
        return result
    except PydanticValidationError as e:
        error_messages = [
            f"{' -> '.join(str(x) for x in err['loc'])}: {err['msg']}" for err in e.errors()
        ]
        raise ValidationError(
            f"Report validation failed with {len(error_messages)} error(s)",
            errors=error_messages,
        ) from e
