from __future__ import annotations

import hashlib
import json
import shutil
from pathlib import Path

import numpy as np
import pytest
import soundfile as sf

from music3lab import reference_guided_append as rga
from music3lab.__main__ import build_parser


def _sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _write_float(path: Path, audio: np.ndarray) -> None:
    sf.write(path, audio.T, rga.SAMPLE_RATE, subtype="FLOAT", format="WAV")


def _fixture_bundle(tmp_path: Path) -> tuple[Path, Path]:
    source = tmp_path / "source.wav"
    t = np.arange(5000, dtype=np.float32) / rga.SAMPLE_RATE
    source_audio = np.stack((0.1 * np.sin(2*np.pi*110*t), 0.09*np.sin(2*np.pi*113*t)))
    sf.write(source, source_audio.T, rga.SAMPLE_RATE, subtype="PCM_16", format="WAV")
    bundle = tmp_path / "candidates.reference-style"
    bundle.mkdir()
    candidates = []
    profiles = []
    for index in range(8):
        phase = np.arange(rga.CANDIDATE_FRAMES, dtype=np.float32) / rga.SAMPLE_RATE
        audio = np.stack((
            0.04 * np.sin(2*np.pi*(220+index)*phase),
            0.035 * np.sin(2*np.pi*(223+index)*phase),
        )).astype(np.float32)
        path = bundle / f"candidate-{index:02d}-seed-{101+index}.wav"
        _write_float(path, audio)
        decoded = sf.read(path, dtype="float32", always_2d=True)[0].T.copy()
        candidates.append({
            "index": index, "seed": 101 + index, "wav_sha256": _sha(path),
            "decoded_pcm_sha256": hashlib.sha256(decoded.astype("<f4").tobytes()).hexdigest(),
            "eligible": index < 3, "rejections": [] if index < 3 else ["tempo drift"],
            "final_score": [0.3, 0.2, 0.2][index] if index < 3 else 0.5 + index / 100,
            "source_correlation": 0.1,
            "metrics": {"rms": 0.03, "clipped_fraction": 0.0, "tempo_bpm": 120.0,
                        "key": "C", "mode": "major", "stereo_width": 0.5},
        })
        profiles.append({"profile_sha256": f"{index:064x}", "finite": True})
    source_decoded = sf.read(source, dtype="float32", always_2d=True)[0].T.copy()
    manifest = {
        "schema_version": "music3lab.reference-style-render.v1",
        "capability": "reference_style_direct_latent_scorer_internal_text_bridge",
        "reference": {
            "source_audio_sha256": _sha(source),
            "decoded_pcm_sha256": hashlib.sha256(source_decoded.astype("<f4").tobytes()).hexdigest(),
            "tempo_bpm": 120.0, "key": "C", "mode": "major", "stereo_width": 0.5,
            "profile": {"profile_sha256": "f" * 64},
        },
        "candidates": candidates,
        "candidate_profiles": profiles,
        "semantic_digest": "a" * 64,
    }
    (bundle / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
    return source, bundle


def test_seam_formula_is_bounded_deterministic_and_combines_both_terms() -> None:
    x = np.zeros((2, rga.CROSSFADE_FRAMES), dtype=np.float32)
    y = np.linspace(-0.2, 0.2, rga.CROSSFADE_FRAMES, dtype=np.float32)[None].repeat(2, 0)
    first = rga.seam_cost(x, y)
    second = rga.seam_cost(x, y)
    assert first == second
    assert first["seam_cost"] == pytest.approx(
        0.5 * (first["waveform_mismatch"] + first["first_difference_mismatch"])
    )
    assert all(0.0 <= value <= 1.0 for value in first.values())


def test_equal_power_append_has_exact_geometry_and_exact_outside_guard() -> None:
    rng = np.random.default_rng(7)
    source = rng.normal(0, 0.05, (2, 6000)).astype(np.float32)
    candidate = rng.normal(0, 0.04, (2, rga.CANDIDATE_FRAMES)).astype(np.float32)
    output = rga.equal_power_append(source, candidate)
    assert output.shape == (2, 6000 + rga.CANDIDATE_FRAMES - rga.CROSSFADE_FRAMES)
    assert np.array_equal(output[:, : 6000-rga.CROSSFADE_FRAMES], source[:, :-rga.CROSSFADE_FRAMES])
    assert np.array_equal(output[:, 6000:], candidate[:, rga.CROSSFADE_FRAMES:])
    assert np.array_equal(output[:, 6000-rga.CROSSFADE_FRAMES], source[:, -rga.CROSSFADE_FRAMES])
    assert np.allclose(output[:, 5999], candidate[:, rga.CROSSFADE_FRAMES-1], atol=2e-7)


def test_ranking_uses_frozen_composite_and_candidate_index_tie() -> None:
    records = [
        {"index": 2, "eligible": True, "composite_score": 0.3},
        {"index": 1, "eligible": True, "composite_score": 0.3},
        {"index": 0, "eligible": False, "composite_score": 0.1},
    ]
    assert rga.select_evaluated(records)["index"] == 1


def test_manifest_source_and_candidate_swap_are_rejected(tmp_path: Path) -> None:
    source, bundle = _fixture_bundle(tmp_path)
    wrong_source = tmp_path / "wrong.wav"
    sf.write(wrong_source, np.zeros((6000, 2), dtype=np.float32), rga.SAMPLE_RATE, subtype="PCM_16")
    with pytest.raises(ValueError, match="provenance"):
        rga.evaluate_bundle(source=wrong_source, bundle=bundle, negative="")
    first = bundle / "candidate-00-seed-101.wav"
    second = bundle / "candidate-01-seed-102.wav"
    saved = first.read_bytes()
    first.write_bytes(second.read_bytes())
    with pytest.raises(ValueError, match="file hash"):
        rga.evaluate_bundle(source=source, bundle=bundle, negative="")
    first.write_bytes(saved)


def test_negative_truth_and_cli_surface(tmp_path: Path) -> None:
    source, bundle = _fixture_bundle(tmp_path)
    _, _, records, constraints, _ = rga.evaluate_bundle(
        source=source, bundle=bundle, negative="no vocals, no clipping, no tempo drift"
    )
    status = {item["name"]: item["status"] for item in constraints.to_dict()["items"]}
    assert status == {"vocals": "NOT_ENFORCEABLE", "clipping": "ENFORCEABLE", "tempo drift": "ENFORCEABLE"}
    assert constraints.to_dict()["native_negative_prompt_used"] is False
    assert [record["index"] for record in records if record["eligible"]] == [0, 1, 2]
    args = build_parser().parse_args([
        "reference-guided-append", "--audio", "a.wav", "--bundle", "b", "--output", "o.wav"
    ])
    assert args.config == "configs/reference-guided-append-v1.yaml" and not hasattr(args, "prompt")


def test_synthetic_end_to_end_roundtrip_and_retained_originals(tmp_path: Path) -> None:
    source, bundle = _fixture_bundle(tmp_path)
    result = rga.run_reference_guided_append(
        source_audio=source, candidate_bundle=bundle, output=tmp_path / "joined.wav",
        negative="no vocals, no clipping, no source copy",
        config_path="configs/reference-guided-append-v1.yaml",
    )
    payload = json.loads(result.sidecar.read_text())
    assert payload["status"] == "PASS" and len(payload["candidates"]) == 8
    assert result.selected_index == 1
    assert payload["geometry"]["net_appended_duration_seconds"] == pytest.approx(7.952834467120181)
    assert payload["validation"]["float_wav_roundtrip_sample_exact"]
    report_lines = result.report.read_text().splitlines()
    assert all(any(line.startswith(f"| {index} |") for line in report_lines) for index in range(8))
    assert _sha(result.evidence_bundle / "source-original.wav") == _sha(source)
    assert _sha(result.evidence_bundle / "candidate-01-seed-102.wav") == _sha(bundle / "candidate-01-seed-102.wav")


def test_retained_real_bundle_integrates_when_present() -> None:
    source = Path("/home/ubuntu/minimax-user-audio/incoming/loveonme_x_osh.wav")
    bundle = Path("/home/ubuntu/minimax-reference-style-evidence/user-loveonme-f040ab1-seed101.reference-style")
    if not source.exists() or not bundle.exists():
        pytest.skip("retained measured bundle unavailable")
    _, _, records, constraints, provenance = rga.evaluate_bundle(
        source=source, bundle=bundle,
        negative="no clipping, no source copy, no tempo drift, no vocals",
    )
    assert provenance["source"]["wav_sha256"] == "ff777c257a9f8044bbe08eb2933d6f4073a22770cf04ae527b8fe768982e6fea"
    assert provenance["manifest_sha256"] == "2ada3a88bc57936810e7c8b8cf920d722225d438ca697745d9869ff4eeaa7485"
    assert [record["index"] for record in records if record["eligible"]] == [0, 1, 2]
    assert rga.select_evaluated(records)["index"] == 0
    assert {item.name: item.status for item in constraints.items}["vocals"] == "NOT_ENFORCEABLE"