from __future__ import annotations

import re
from typing import Any

from epochler.contract.models import (
    ContractPhase,
    DatasetAsset,
    EvalSpec,
    FrozenContract,
    LockState,
    ModelSpec,
    PhaseExecutionStatus,
    PhaseName,
    PreprocessingPlan,
    TaskType,
    ensure_builtin_phases,
)

_PROPOSAL_RE = re.compile(
    r"\*\*\[PROPOSAL:\s*([^\]]+)\]\*\*\s*(.*?)(?=\n\s*\*\*\[(?:PROPOSAL|TOOL):|\Z)",
    re.IGNORECASE | re.DOTALL,
)

_LABEL_TO_COMPONENT = {
    "dataset": "dataset",
    "preprocessing": "preprocessing",
    "evaluation": "eval_spec",
    "eval spec": "eval_spec",
    "eval": "eval_spec",
    "model": "model_spec",
    "model spec": "model_spec",
    "target threshold": "target",
    "threshold": "target",
    "target": "target",
}

_PHASE_LABELS = {
    "discovery": PhaseName.DISCOVERY.value,
    "preprocessing phase": PhaseName.PREPROCESSING.value,
    "featurization": PhaseName.FEATURIZATION.value,
    "materialization": PhaseName.FEATURIZATION.value,
    "featurization/materialization": PhaseName.FEATURIZATION.value,
    "training phase": PhaseName.TRAINING.value,
    "evaluation phase": PhaseName.EVALUATION.value,
    "analysis": PhaseName.ANALYSIS.value,
    "analysis phase": PhaseName.ANALYSIS.value,
}


def sync_contract_from_text(
    contract: FrozenContract,
    text: str,
    *,
    only_components: set[str] | None = None,
    reopen_frozen: bool = False,
) -> set[str]:
    ensure_builtin_phases(contract)
    updated: set[str] = set()
    for raw_label, body in _PROPOSAL_RE.findall(text):
        component = _resolve_component(raw_label)
        phase_name = _resolve_phase(raw_label)
        if not component:
            if phase_name:
                if _update_phase_from_text(contract, phase_name, body, reopen_frozen=reopen_frozen):
                    updated.add(phase_name)
            continue
        if only_components and component not in only_components:
            continue
        if component == "target":
            before = (
                contract.target_threshold,
                contract.threshold_direction,
                contract.budget_minutes,
                contract.max_experiments,
            )
            if _parse_target(contract, body):
                after = (
                    contract.target_threshold,
                    contract.threshold_direction,
                    contract.budget_minutes,
                    contract.max_experiments,
                )
                if reopen_frozen and contract.frozen and before != after:
                    contract.frozen = False
                updated.add(component)
            continue

        current = getattr(contract, component, None)
        if current is not None and current.lock_state == LockState.FROZEN:
            if not reopen_frozen:
                continue

        parsed = _parse_component(component, body)
        if parsed is None:
            continue
        if current is not None and _component_payload(current) == _component_payload(parsed):
            continue
        setattr(contract, component, parsed)
        if reopen_frozen and contract.frozen:
            contract.frozen = False
        if phase_name:
            _update_phase_from_text(contract, phase_name, body, reopen_frozen=reopen_frozen)
        updated.add(component)
    return updated


def _resolve_component(label: str) -> str | None:
    lowered = label.strip().lower()
    if lowered in _LABEL_TO_COMPONENT:
        return _LABEL_TO_COMPONENT[lowered]
    for key, component in _LABEL_TO_COMPONENT.items():
        if key in lowered:
            return component
    return None


def _resolve_phase(label: str) -> str | None:
    lowered = label.strip().lower()
    if lowered in _PHASE_LABELS:
        return _PHASE_LABELS[lowered]
    for key, phase_name in _PHASE_LABELS.items():
        if key in lowered:
            return phase_name
    return None


def _parse_component(component: str, body: str) -> Any | None:
    if component == "dataset":
        return _parse_dataset(body)
    if component == "preprocessing":
        return _parse_preprocessing(body)
    if component == "eval_spec":
        return _parse_eval_spec(body)
    if component == "model_spec":
        return _parse_model_spec(body)
    return None


def _parse_dataset(body: str) -> DatasetAsset:
    fields = _parse_bullets(body)
    extras = _collect_extras(fields, {"path", "format", "target", "task type", "rows", "features"})

    path = _field_text(fields, "path")
    fmt = _field_text(fields, "format")
    target = _field_text(fields, "target")
    task_text = _field_text(fields, "task type") or target or body
    features_text = _field_text(fields, "features")
    rows_text = _field_text(fields, "rows") or fmt

    if features_text:
        extras["features"] = features_text
    sampling_rate = _field_text(fields, "sampling rate")
    if sampling_rate:
        extras["sampling_rate"] = sampling_rate

    return DatasetAsset(
        path=path,
        format=fmt,
        target_column=target,
        task_type=_infer_task_type(task_text),
        num_rows=_extract_first_int(rows_text),
        num_features=_extract_first_int(features_text),
        schema_info=extras,
        lock_state=LockState.PROPOSED,
    )


def _parse_preprocessing(body: str) -> PreprocessingPlan:
    fields = _parse_bullets(body)
    transforms: list[dict[str, Any]] = []

    for key, payload in fields.items():
        items = payload["items"] or ([payload["value"]] if payload["value"] else [])
        if not items:
            continue
        for item in items:
            transforms.append({"stage": key, "description": item})

    if not transforms:
        for line in _plain_lines(body):
            transforms.append({"description": line})

    text = body.lower()
    reproducibility = "adapted" if "different from" in text or "possible" in text else "exact"
    source = "adapted" if reproducibility == "adapted" else "copied"

    return PreprocessingPlan(
        source=source,
        transforms=transforms,
        reproducibility_label=reproducibility,
        lock_state=LockState.PROPOSED,
    )


def _parse_eval_spec(body: str) -> EvalSpec:
    fields = _parse_bullets(body)
    constraints = _collect_extras(fields, {"primary metric", "secondary metrics", "split strategy"})

    primary = _field_text(fields, "primary metric")
    split_strategy = _normalize_split_strategy(_field_text(fields, "split strategy"))
    secondary = _field_items(fields, "secondary metrics")
    if not secondary:
        inline_secondary = _field_text(fields, "secondary metrics")
        if inline_secondary:
            secondary = [part.strip() for part in inline_secondary.split(",") if part.strip()]

    test_set = _field_text(fields, "test set")
    if test_set:
        constraints["test_set"] = test_set

    return EvalSpec(
        task_type=_infer_task_type(body),
        primary_metric=primary,
        secondary_metrics=secondary,
        split_strategy=split_strategy,
        constraints=constraints,
        lock_state=LockState.PROPOSED,
    )


def _parse_model_spec(body: str) -> ModelSpec:
    allowed_families: list[str] = []
    search_space: dict[str, Any] = {}
    hardware_constraints: dict[str, Any] = {}
    current_family: str | None = None

    for raw_line in body.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        family_match = re.match(r"^\d+\.\s+\*\*([^*]+)\*\*:?\s*$", line)
        if family_match:
            current_family = family_match.group(1).strip()
            allowed_families.append(current_family)
            continue

        search_match = re.match(r"^(?:[-*]\s+)?search space:\s*(.+)$", line, re.IGNORECASE)
        if search_match and current_family:
            search_space[current_family] = _parse_search_space(search_match.group(1))
            continue

        if "cpu-only" in line.lower():
            hardware_constraints["execution"] = "cpu-only"

    if not allowed_families:
        fields = _parse_bullets(body)
        families = _field_items(fields, "allowed families")
        if families:
            allowed_families = families

    if "continuous" in body.lower():
        hardware_constraints.setdefault("signal_type", "continuous")

    return ModelSpec(
        allowed_families=allowed_families,
        search_space=search_space,
        hardware_constraints=hardware_constraints,
        lock_state=LockState.PROPOSED,
    )


def _parse_target(contract: FrozenContract, body: str) -> bool:
    updated = False
    threshold = None

    for raw_line in body.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if "target" in line.lower():
            threshold = _extract_float(line)
            if threshold is not None:
                contract.target_threshold = threshold
                if "<" in line:
                    contract.threshold_direction = "lower_is_better"
                elif ">" in line:
                    contract.threshold_direction = "higher_is_better"
                updated = True
        elif line.lower().startswith("- direction:"):
            value = line.split(":", 1)[1].strip()
            if value:
                contract.threshold_direction = value
                updated = True
        elif line.lower().startswith("- budget:"):
            minutes = _parse_duration_minutes(line.split(":", 1)[1].strip())
            if minutes is not None:
                contract.budget_minutes = minutes
                updated = True
        elif line.lower().startswith("- max experiments:"):
            max_experiments = _extract_first_int(line)
            if max_experiments is not None:
                contract.max_experiments = max_experiments
                updated = True

    return updated


def _parse_bullets(body: str) -> dict[str, dict[str, Any]]:
    fields: dict[str, dict[str, Any]] = {}
    current_key: str | None = None

    for raw_line in body.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        field_match = re.match(r"^[-*]\s+([^:]+):\s*(.*)$", line)
        if field_match:
            current_key = _clean_text(field_match.group(1)).lower()
            fields[current_key] = {"value": _clean_text(field_match.group(2)), "items": []}
            continue

        item_match = re.match(r"^(?:[-*]|\d+\.)\s+(.+)$", line)
        if item_match:
            item = _clean_text(item_match.group(1))
            if current_key:
                fields[current_key]["items"].append(item)
            continue

        if current_key:
            payload = fields[current_key]
            if payload["items"]:
                payload["items"][-1] = _clean_text(f"{payload['items'][-1]} {line}")
            elif payload["value"]:
                payload["value"] = _clean_text(f"{payload['value']} {line}")
            else:
                payload["value"] = _clean_text(line)

    return fields


def _plain_lines(body: str) -> list[str]:
    lines: list[str] = []
    for raw_line in body.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        cleaned = re.sub(r"^(?:[-*]|\d+\.)\s+", "", line).strip()
        if cleaned:
            lines.append(_clean_text(cleaned))
    return lines


def _field_text(fields: dict[str, dict[str, Any]], key: str) -> str:
    payload = fields.get(key.lower())
    if not payload:
        return ""
    return payload["value"]


def _field_items(fields: dict[str, dict[str, Any]], key: str) -> list[str]:
    payload = fields.get(key.lower())
    if not payload:
        return []
    return payload["items"]


def _collect_extras(
    fields: dict[str, dict[str, Any]],
    ignored: set[str],
) -> dict[str, Any]:
    extras: dict[str, Any] = {}
    for key, payload in fields.items():
        if key in ignored:
            continue
        if payload["items"]:
            extras[key.replace(" ", "_")] = payload["items"]
        elif payload["value"]:
            extras[key.replace(" ", "_")] = payload["value"]
    return extras


def _infer_task_type(text: str) -> TaskType | None:
    lower = text.lower()
    if "time series" in lower or "forecast" in lower:
        return TaskType.TIME_SERIES
    if "classification" in lower or "classify" in lower:
        return TaskType.CLASSIFICATION
    if "regression" in lower:
        return TaskType.REGRESSION
    if "nlp" in lower or "language" in lower or "text" in lower:
        return TaskType.NLP
    if "vision" in lower or "image" in lower:
        return TaskType.VISION
    if "ranking" in lower:
        return TaskType.RANKING
    return None


def _normalize_split_strategy(value: str) -> str:
    lower = value.lower()
    if "time" in lower:
        return "time_series"
    if "strat" in lower:
        return "stratified"
    if "group" in lower:
        return "grouped"
    if "holdout" in lower:
        return "holdout"
    if "fold" in lower:
        return "kfold"
    return value


def _parse_search_space(text: str) -> dict[str, Any]:
    parsed: dict[str, Any] = {}
    for match in re.finditer(r"([a-zA-Z_][a-zA-Z0-9_]*)=\[([^\]]+)\]", text):
        key = match.group(1)
        raw_values = [part.strip() for part in match.group(2).split(",") if part.strip()]
        values: list[Any] = []
        for raw_value in raw_values:
            if re.fullmatch(r"-?\d+", raw_value):
                values.append(int(raw_value))
            elif re.fullmatch(r"-?\d+\.\d+", raw_value):
                values.append(float(raw_value))
            else:
                values.append(raw_value)
        parsed[key] = values
    if not parsed and text.strip():
        parsed["raw"] = text.strip()
    return parsed


def _extract_first_int(text: str) -> int | None:
    match = re.search(r"(\d[\d,]*)", text)
    if not match:
        return None
    return int(match.group(1).replace(",", ""))


def _extract_float(text: str) -> float | None:
    match = re.search(r"(-?\d+(?:\.\d+)?)", text)
    if not match:
        return None
    return float(match.group(1))


def _parse_duration_minutes(text: str) -> int | None:
    lower = text.lower()
    value = _extract_float(lower)
    if value is None:
        return None
    if "hour" in lower:
        return int(value * 60)
    if "day" in lower:
        return int(value * 24 * 60)
    if "min" in lower:
        return int(value)
    return None


def _clean_text(text: str) -> str:
    return text.strip().strip("`").strip()


def _component_payload(component: Any) -> dict[str, Any]:
    if hasattr(component, "model_dump"):
        return component.model_dump(exclude={"lock_state"}, mode="json")
    return component


def _update_phase_from_text(
    contract: FrozenContract,
    phase_name: str,
    body: str,
    *,
    reopen_frozen: bool = False,
) -> bool:
    ensure_builtin_phases(contract)
    phase: ContractPhase = contract.phases[phase_name]
    new_description = body.strip()
    if not new_description:
        return False
    changed = phase.description.strip() != new_description
    phase.description = new_description
    phase.metadata["proposal_body"] = new_description
    if changed:
        phase.status = PhaseExecutionStatus.READY
        if phase.lock_state == LockState.FROZEN:
            phase.lock_state = LockState.PROPOSED
        if reopen_frozen and contract.frozen:
            contract.frozen = False
    return changed
