from __future__ import annotations

import hashlib, multiprocessing, random
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 .rebuild import rebuild_candidate
from .reports import generate_reports, read_json, write_json, write_manifest

PIPELINE_VERSION = "cadfs_to_cdsl.v1"


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 scan(input_root: Path, output: Path) -> list[dict[str, Any]]:
    records = [sample.as_dict() for sample in scan_dataset(input_root)]
    write_json(output / "dataset_index.json", {"schema": "cadfs_to_cdsl.dataset_index.v1", "input": str(input_root), "sample_count": len(records), "records": records})
    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
    diagnostics = list(sample.diagnostics)
    try:
        feature_path = Path(sample.files["featurescript"])
        model = parse_featurescript(feature_path.read_text(encoding="utf-8"), sample.sample_id)
        provenance = {"source_featurescript": str(feature_path), **{f"source_{key}_sha256": value for key, value in sample.hashes.items()}, "jsonl": sample.metadata}
        result = lower_model(model, provenance); diagnostics.extend(result.diagnostics)
        write_json(directory / "history.json", result.history); write_json(directory / "diagnostics.json", diagnostics)
        if result.cdsl is not None:
            from engine.cdsl_engine.semantic_validation import validate_semantic_cdsl
            semantic = validate_semantic_cdsl(result.cdsl); write_json(directory / "candidate.cdsl.json", result.cdsl)
        else: semantic = None
        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)}
    except Exception as exc:
        missing = isinstance(exc, FileNotFoundError)
        diagnostics.append({"code": "source_missing" if missing else "parse_or_lowering_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


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) -> bool:
    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 False
    return process.exitcode == 0 and result_path.exists()


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 = read_json(status_path) if status_path.exists() else convert_one(sample, output, force=force)
    if status.get("conversion_status") != "converted_complete": return status
    rebuild_path = directory / "rebuild.json"; step_path = directory / "rebuild.step"
    if not force and rebuild_path.exists() and status.get("rebuild_status"):
        if status.get("rebuild_status") != "rebuilt" or step_path.exists(): return status
    worker_result = directory / "rebuild.worker.json"
    completed = _isolated(_rebuild_worker, (str(directory / "candidate.cdsl.json"), str(step_path), str(worker_result)), worker_result, timeout_seconds)
    if completed: result = read_json(worker_result); worker_result.unlink(missing_ok=True)
    else:
        step_path.unlink(missing_ok=True)
        result = {"status": "rebuild_timeout", "error": {"type": "TimeoutError", "message": f"rebuild exceeded {timeout_seconds:g} seconds"}}
    write_json(rebuild_path, result)
    status["rebuild_status"] = result["status"]; 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"
        completed = _isolated(_compare_worker, (sample.files["step"], str(directory / "rebuild.step"), str(worker_result)), worker_result, timeout_seconds)
        if completed: comparison = read_json(worker_result); worker_result.unlink(missing_ok=True); write_json(comparison_path, comparison)
        else:
            status["status"] = "comparison_timeout"; status["comparison_error"] = {"type": "TimeoutError", "message": f"comparison exceeded {timeout_seconds:g} seconds"}; write_json(status_path, status); return status
    else: comparison = read_json(comparison_path)
    status["comparison_decision"] = comparison["decision"]
    accepted = comparison[compare_mode]["passed"]
    status["status"] = "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