"""Schema discovery for JSONL metadata."""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import yaml

from .record import parse_jsonl

MAX_ENUM_VALUES = 20


@dataclass
class SchemaEntry:
    """Discovered schema for a single metadata key."""

    type: str
    count: int
    values: Optional[List[Any]] = None
    example: Optional[Any] = None
    description: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        d = {"type": self.type, "count": self.count}
        if self.values is not None:
            d["values"] = self.values
        if self.example is not None:
            d["example"] = self.example
        if self.description is not None:
            d["description"] = self.description
        return d


@dataclass
class CollectionSchema:
    """Schema for an entire collection: record count + metadata keys."""

    record_count: int = 0
    metadata_keys: Dict[str, SchemaEntry] = field(default_factory=dict)


def _json_type(value: Any) -> str:
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, (int, float)):
        return "number"
    if isinstance(value, str):
        return "string"
    if isinstance(value, list):
        return "array"
    if isinstance(value, dict):
        return "object"
    return "string"


def discover_schema(path: Union[str, Path]) -> Dict[str, SchemaEntry]:
    """Scan a JSONL file and discover metadata key schemas."""
    key_counts: Dict[str, int] = {}
    key_types: Dict[str, str] = {}
    key_values: Dict[str, set] = {}
    key_example: Dict[str, Any] = {}

    for record in parse_jsonl(path):
        if not record.metadata:
            continue
        for key, value in record.metadata.items():
            key_counts[key] = key_counts.get(key, 0) + 1
            key_types[key] = _json_type(value)

            if key not in key_example:
                key_example[key] = value

            # Track enumerable values; set to None for non-scalar types
            if key not in key_values:
                key_values[key] = set()
            if key_values[key] is None:
                continue
            if isinstance(value, (str, int, float, bool)):
                key_values[key].add(value)
            else:
                key_values[key] = None

    result = {}
    for key in key_counts:
        values_set = key_values.get(key)
        if values_set is not None and len(values_set) <= MAX_ENUM_VALUES:
            if all(isinstance(v, str) for v in values_set):
                values = sorted(str(v) for v in values_set)
            else:
                values = list(values_set)
            entry = SchemaEntry(
                type=key_types[key],
                count=key_counts[key],
                values=values,
            )
        else:
            entry = SchemaEntry(
                type=key_types[key],
                count=key_counts[key],
                example=key_example[key],
            )
        result[key] = entry

    return result


def load_schema_yaml(path: Union[str, Path]) -> Dict[str, CollectionSchema]:
    """Load schema.yaml into collection schemas.

    Returns a dict mapping collection name to CollectionSchema.
    """
    path = Path(path)
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        return {}

    result = {}
    for coll_name, coll_data in data.items():
        if not isinstance(coll_data, dict):
            continue
        record_count = coll_data.get("record_count", 0)
        keys = {}
        for key_name, key_data in coll_data.get("metadata_keys", {}).items():
            if not isinstance(key_data, dict):
                continue
            keys[key_name] = SchemaEntry(
                type=key_data.get("type", "string"),
                count=key_data.get("count", 0),
                values=key_data.get("values"),
                example=key_data.get("example"),
                description=key_data.get("description"),
            )
        result[coll_name] = CollectionSchema(
            record_count=record_count, metadata_keys=keys
        )

    return result


def save_schema_yaml(
    schemas: Dict[str, CollectionSchema], path: Union[str, Path]
) -> None:
    """Write schema.yaml from collection schemas.

    Produces human-readable YAML with comments.
    """
    path = Path(path)

    data = {}
    for coll_name, coll_schema in schemas.items():
        coll_dict = {"record_count": coll_schema.record_count, "metadata_keys": {}}
        for key_name, entry in coll_schema.metadata_keys.items():
            key_dict = {}
            if entry.description is not None:
                key_dict["description"] = entry.description
            key_dict["type"] = entry.type
            key_dict["count"] = entry.count
            if entry.values is not None:
                key_dict["values"] = entry.values
            if entry.example is not None:
                key_dict["example"] = entry.example
            coll_dict["metadata_keys"][key_name] = key_dict
        data[coll_name] = coll_dict

    header = "# Auto-generated by arkiv. Edit freely.\n"
    body = yaml.dump(
        data,
        default_flow_style=False,
        allow_unicode=True,
        sort_keys=False,
    )
    path.write_text(header + body, encoding="utf-8")


def merge_schema(
    auto: Dict[str, SchemaEntry],
    curated: Dict[str, SchemaEntry],
) -> Dict[str, SchemaEntry]:
    """Merge auto-discovered schema with curated schema.

    Live fields (type, count) come from auto.
    Stable fields (description, values if curated) come from curated.
    Keys in curated but not in auto are preserved with count=0.
    """
    result = {}

    # Start with all auto-discovered keys
    for key, auto_entry in auto.items():
        curated_entry = curated.get(key)
        if curated_entry:
            result[key] = SchemaEntry(
                type=auto_entry.type,
                count=auto_entry.count,
                values=curated_entry.values
                if curated_entry.values is not None
                else auto_entry.values,
                description=curated_entry.description,
            )
        else:
            result[key] = auto_entry

    # Preserve curated-only keys (not found in data)
    for key, curated_entry in curated.items():
        if key not in result:
            result[key] = SchemaEntry(
                type=curated_entry.type or "string",
                count=0,
                values=curated_entry.values,
                description=curated_entry.description,
            )

    return result
