from __future__ import annotations

import hashlib, multiprocessing, random
from copy import deepcopy
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from .compare import compare_steps
from .dataset import Sample, scan_dataset
from .featurescript_parser import parse_featurescript
from .lowering import lower_model
from .operation_registry import build_operation_registry
from .rebuild import rebuild_candidate
from .reports import generate_reports, read_json, write_json, write_manifest

PIPELINE_VERSION = "cadfs_to_cdsl.v2"


def _fingerprint(sample: Sample) -> str:
    value = PIPELINE_VERSION + "|" + "|".join(f"{key}:{value}" for key, value in sorted(sample.hashes.items()))
    return hashlib.sha256(value.encode()).hexdigest()


def _sample_dir(output: Path, sample_id: str) -> Path: return output / "samples" / sample_id


def _semantic_valid_prefix(
    cdsl: dict[str, Any],
    validate_semantic_cdsl: Any,
) -> tuple[dict[str, Any], dict[str, Any]] | None:
    """Return the longest contiguous semantically valid feature prefix.

    This is an artifact-recovery boundary, not permission to execute an
    invalid candidate.  The selected prefix is independently validated and
    saved under an explicit prefix filename; the complete rejected CDSL stays
    diagnostic-only.
    """
    features = cdsl.get("features") or []
    for count in range(len(features) - 1, 0, -1):
        prefix = deepcopy(cdsl)
        prefix["features"] = prefix["features"][:count]
        try:
            semantic = validate_semantic_cdsl(prefix)
        except Exception:
            continue
        return prefix, semantic
    return None


def scan(input_root: Path, output: Path) -> list[dict[str, Any]]:
    samples = scan_dataset(input_root)
    records = [sample.as_dict() for sample in samples]
    write_json(output / "dataset_index.json", {"schema": "cadfs_to_cdsl.dataset_index.v1", "input": str(input_root), "sample_count": len(records), "records": records})
    write_json(output / "operation_registry.json", build_operation_registry(samples))
    initial = [{"sample_id": item["sample_id"], "status": "scanned", "diagnostics": item["diagnostics"]} for item in records]
    write_manifest(output / "manifest.jsonl", initial)
    return records


def load_samples(input_root: Path, output: Path) -> list[Sample]:
    index = output / "dataset_index.json"
    records = read_json(index)["records"] if index.exists() else scan(input_root, output)
    return [Sample(**{key: value for key, value in item.items() if key in {"sample_id", "files", "hashes", "metadata", "diagnostics"}}) for item in records]


def select_samples(samples: list[Sample], *, sample_ids: list[str] | None = None, offset: int = 0, limit: int | None = None, seed: int | None = None) -> list[Sample]:
    if sample_ids:
        wanted = set(sample_ids); chosen = [sample for sample in samples if sample.sample_id in wanted]
        missing = wanted - {sample.sample_id for sample in chosen}
        if missing: raise ValueError("unknown sample ids: " + ", ".join(sorted(missing)))
        return chosen
    chosen = list(samples)
    if seed is not None and limit is not None:
        chosen = random.Random(seed).sample(chosen, min(limit, len(chosen))); chosen.sort(key=lambda item: item.sample_id); return chosen
    return chosen[offset:None if limit is None else offset + limit]


def convert_one(sample: Sample, output: Path, *, force: bool = False) -> dict[str, Any]:
    directory = _sample_dir(output, sample.sample_id); directory.mkdir(parents=True, exist_ok=True)
    fingerprint = _fingerprint(sample); status_path = directory / "status.json"
    if not force and status_path.exists():
        cached = read_json(status_path)
        if cached.get("input_fingerprint") == fingerprint and cached.get("conversion_status"):
            return cached
    if force:
        for name in (
            "candidate.cdsl.json", "candidate.invalid.cdsl.json", "candidate.prefix.cdsl.json",
            "bound.cdsl.json", "prefix.bound.cdsl.json", "rebuild.step", "prefix.rebuild.step",
            "rebuild.json", "prefix.rebuild.json", "rebuild.worker.json", "comparison.json",
            "comparison.worker.json",
        ):
            (directory / name).unlink(missing_ok=True)
    diagnostics = list(sample.diagnostics)
    try:
        feature_path = Path(sample.files["featurescript"])
        model = parse_featurescript(feature_path.read_text(encoding="utf-8"), sample.sample_id)
    except Exception as exc:
        missing = isinstance(exc, FileNotFoundError)
        diagnostics.append({"code": "source_missing" if missing else "parse_failed", "message": str(exc), "type": type(exc).__name__})
        write_json(directory / "diagnostics.json", diagnostics)
        final_status = "source_missing" if missing else "parse_failed"
        status = {"schema": "cadfs_to_cdsl.status.v1", "sample_id": sample.sample_id, "status": final_status, "conversion_status": final_status, "input_fingerprint": fingerprint, "diagnostic_count": len(diagnostics)}
        write_json(status_path, status); return status
    provenance = {
        "source_featurescript": str(feature_path),
        **{f"source_{key}_sha256": value for key, value in sample.hashes.items()},
        "jsonl": sample.metadata,
    }
    try:
        result = lower_model(model, provenance)
    except Exception as exc:
        diagnostics.append({"code": "lowering_failed", "message": str(exc), "type": type(exc).__name__})
        write_json(directory / "diagnostics.json", diagnostics)
        status = {
            "schema": "cadfs_to_cdsl.status.v1", "sample_id": sample.sample_id,
            "status": "lowering_failed", "conversion_status": "lowering_failed",
            "input_fingerprint": fingerprint, "diagnostic_count": len(diagnostics),
        }
        write_json(status_path, status); return status
    diagnostics.extend(result.diagnostics)
    write_json(directory / "history.json", result.history)
    if result.cdsl is None:
        write_json(directory / "diagnostics.json", diagnostics)
        status = {
            "schema": "cadfs_to_cdsl.status.v1", "sample_id": sample.sample_id,
            "status": result.status, "conversion_status": result.status,
            "input_fingerprint": fingerprint, "semantic_validation": None, "diagnostic_count": len(diagnostics),
        }
        write_json(status_path, status); return status
    try:
        from engine.cdsl_engine.semantic_validation import validate_semantic_cdsl
        semantic = validate_semantic_cdsl(result.cdsl)
    except Exception as exc:
        diagnostics.append({"code": "semantic_validation_failed", "message": str(exc), "type": type(exc).__name__})
        write_json(directory / "candidate.invalid.cdsl.json", result.cdsl)
        prefix_result = _semantic_valid_prefix(result.cdsl, validate_semantic_cdsl)
        prefix_status: dict[str, Any] | None = None
        if prefix_result is not None:
            prefix, prefix_semantic = prefix_result
            prefix_features = prefix.get("features") or []
            write_json(directory / "candidate.prefix.cdsl.json", prefix)
            prefix_status = {
                "feature_count": len(prefix_features),
                "last_feature_id": prefix_features[-1].get("id"),
                "semantic_validation": prefix_semantic,
            }
            diagnostics.append({
                "code": "semantic_valid_prefix_preserved",
                "feature_count": len(prefix_features),
                "last_feature_id": prefix_features[-1].get("id"),
                "message": "longest contiguous semantic-valid CDSL prefix was preserved separately",
            })
        write_json(directory / "diagnostics.json", diagnostics)
        status = {
            "schema": "cadfs_to_cdsl.status.v1", "sample_id": sample.sample_id,
            "status": "semantic_validation_failed", "conversion_status": "semantic_validation_failed",
            "input_fingerprint": fingerprint,
            "semantic_validation": {"valid": False, "error": {"type": type(exc).__name__, "message": str(exc)}},
            "diagnostic_count": len(diagnostics),
        }
        if prefix_status is not None:
            status["semantic_valid_prefix"] = prefix_status
        write_json(status_path, status); return status
    write_json(directory / "candidate.cdsl.json", result.cdsl)
    write_json(directory / "diagnostics.json", diagnostics)
    status = {
        "schema": "cadfs_to_cdsl.status.v1", "sample_id": sample.sample_id,
        "status": result.status, "conversion_status": result.status,
        "input_fingerprint": fingerprint, "semantic_validation": semantic, "diagnostic_count": len(diagnostics),
    }
    write_json(status_path, status); return status


def _rebuild_worker(candidate: str, step: str, result: str) -> None:
    from .rebuild import rebuild_candidate
    write_json(Path(result), rebuild_candidate(read_json(Path(candidate)), Path(step)))


def _compare_worker(gold: str, rebuilt: str, result: str) -> None:
    write_json(Path(result), compare_steps(Path(gold), Path(rebuilt)))


def _isolated(target: Any, args: tuple[str, ...], result_path: Path, timeout_seconds: float) -> str:
    result_path.unlink(missing_ok=True)
    context = multiprocessing.get_context("spawn")
    process = context.Process(target=target, args=args)
    process.start(); process.join(timeout_seconds)
    if process.is_alive():
        process.terminate(); process.join(5)
        if process.is_alive(): process.kill(); process.join()
        return "timeout"
    return "completed" if process.exitcode == 0 and result_path.exists() else "failed"


def rebuild_one(sample: Sample, output: Path, *, force: bool = False, timeout_seconds: float = 30.0) -> dict[str, Any]:
    directory = _sample_dir(output, sample.sample_id); status_path = directory / "status.json"
    status = convert_one(sample, output, force=True) if force else read_json(status_path) if status_path.exists() else convert_one(sample, output)
    # A partial conversion may still retain a self-contained, semantically
    # executable CDSL prefix.  That prefix is valuable engine evidence and
    # must be rebuilt instead of being hidden behind the conversion label.
    # ``rebuild_candidate`` keeps the two outcomes distinct by reporting a
    # runtime-ineligible candidate when no executable body can be produced.
    candidate_path = directory / "candidate.cdsl.json"
    is_prefix = False
    if not candidate_path.exists():
        candidate_path = directory / "candidate.prefix.cdsl.json"
        is_prefix = candidate_path.exists()
    if not candidate_path.exists(): return status
    rebuild_path = directory / ("prefix.rebuild.json" if is_prefix else "rebuild.json")
    step_path = directory / ("prefix.rebuild.step" if is_prefix else "rebuild.step")
    status_key = "prefix_rebuild_status" if is_prefix else "rebuild_status"
    if not force and rebuild_path.exists() and status.get(status_key):
        if status.get(status_key) != "rebuilt" or step_path.exists(): return status
    worker_result = directory / "rebuild.worker.json"
    outcome = _isolated(_rebuild_worker, (str(candidate_path), str(step_path), str(worker_result)), worker_result, timeout_seconds)
    if outcome == "completed":
        result = read_json(worker_result); worker_result.unlink(missing_ok=True)
        bound_cdsl = result.pop("bound_cdsl", None)
        prefix = result.get("last_executable_prefix")
        if isinstance(prefix, dict):
            prefix_bound_cdsl = prefix.pop("bound_cdsl", None)
            if bound_cdsl is None and prefix_bound_cdsl is not None:
                bound_cdsl = prefix_bound_cdsl
                prefix["bound_feature_count"] = prefix.get("feature_count")
        if bound_cdsl is not None:
            write_json(directory / ("prefix.bound.cdsl.json" if is_prefix else "bound.cdsl.json"), bound_cdsl)
    else:
        step_path.unlink(missing_ok=True)
        error_type = "TimeoutError" if outcome == "timeout" else "WorkerProcessError"
        message = f"rebuild exceeded {timeout_seconds:g} seconds" if outcome == "timeout" else "rebuild worker exited without a result"
        result = {"status": "rebuild_timeout" if outcome == "timeout" else "rebuild_failed", "error": {"type": error_type, "message": message}}
    write_json(rebuild_path, result)
    status[status_key] = result["status"]
    if not is_prefix:
        status["status"] = result["status"]
    write_json(status_path, status); return status


def compare_one(sample: Sample, output: Path, *, force: bool = False, compare_mode: str = "rp", timeout_seconds: float = 60.0) -> dict[str, Any]:
    directory = _sample_dir(output, sample.sample_id); status_path = directory / "status.json"
    status = read_json(status_path) if status_path.exists() else rebuild_one(sample, output, force=force)
    if status.get("rebuild_status") != "rebuilt": return status
    comparison_path = directory / "comparison.json"
    if not force and status.get("status") in {"comparison_failed", "comparison_timeout"}: return status
    if force or not comparison_path.exists():
        worker_result = directory / "comparison.worker.json"
        outcome = _isolated(_compare_worker, (sample.files["step"], str(directory / "rebuild.step"), str(worker_result)), worker_result, timeout_seconds)
        if outcome == "completed": comparison = read_json(worker_result); worker_result.unlink(missing_ok=True); write_json(comparison_path, comparison)
        else:
            status["status"] = "comparison_timeout" if outcome == "timeout" else "comparison_failed"
            error_type = "TimeoutError" if outcome == "timeout" else "WorkerProcessError"
            message = f"comparison exceeded {timeout_seconds:g} seconds" if outcome == "timeout" else "comparison worker exited without a result"
            status["comparison_error"] = {"type": error_type, "message": message}; write_json(status_path, status); return status
    else: comparison = read_json(comparison_path)
    # A forced retry can recover from a transient worker timeout. The final
    # status must describe the successful comparison rather than retaining a
    # stale infrastructure failure beside an accepted result.
    status.pop("comparison_error", None)
    status["comparison_decision"] = comparison["decision"]
    accepted = comparison[compare_mode]["passed"]
    status["status"] = "rebuilt_strict" if comparison["strict"]["passed"] else "rebuilt_approximate" if accepted else "rebuilt_rejected"
    write_json(status_path, status); return status


def run_stage(stage: str, samples: list[Sample], output: Path, *, force: bool = False, compare_mode: str = "rp", timeout_seconds: float = 30.0, workers: int = 1) -> list[dict[str, Any]]:
    def process(sample: Sample) -> dict[str, Any]:
        if stage == "convert": record = convert_one(sample, output, force=force)
        elif stage == "rebuild": record = rebuild_one(sample, output, force=force, timeout_seconds=timeout_seconds)
        elif stage == "compare": record = compare_one(sample, output, force=force, compare_mode=compare_mode, timeout_seconds=max(60.0, timeout_seconds))
        elif stage == "pipeline":
            convert_one(sample, output, force=force); rebuild_one(sample, output, force=force, timeout_seconds=timeout_seconds); record = compare_one(sample, output, force=force, compare_mode=compare_mode, timeout_seconds=max(60.0, timeout_seconds))
        else: raise ValueError(f"unknown stage {stage!r}")
        return record
    if workers == 1: records = [process(sample) for sample in samples]
    else:
        with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="cadfs") as executor:
            records = list(executor.map(process, samples))
    manifest_path = output / "manifest.jsonl"
    existing = {}
    if manifest_path.exists():
        import json
        existing = {item["sample_id"]: item for line in manifest_path.read_text(encoding="utf-8").splitlines() if line.strip() for item in [json.loads(line)]}
    existing.update({item["sample_id"]: item for item in records})
    merged = [existing[key] for key in sorted(existing)]
    write_manifest(manifest_path, merged); generate_reports(output, merged)
    return records