#!/usr/bin/env /usr/local/miniconda3/bin/python
from __future__ import annotations

import argparse
import hashlib
import importlib
import json
import wave
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
import sys
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.engines.chromaprint_matcher import ChromaprintMatcher, load_audio_mono

MERT_MODEL_ID = 'm-a-p/MERT-v1-95M'
_MERT_RUNTIME = None


def load_jsonl(path: Path):
    for line in path.read_text().splitlines():
        line = line.strip()
        if line:
            yield json.loads(line)


def module_available(name: str) -> bool:
    try:
        importlib.import_module(name)
        return True
    except Exception:
        return False


def semantic_runtime_available() -> tuple[bool, list[str]]:
    required = ['torch', 'torchaudio', 'transformers']
    missing = [m for m in required if not module_available(m)]
    return (len(missing) == 0, missing)


def read_wav_stats(path: Path, start_ms: int, end_ms: int) -> dict:
    with wave.open(str(path), 'rb') as wf:
        rate = wf.getframerate()
        sampwidth = wf.getsampwidth()
        n_channels = wf.getnchannels()
        start_frame = int(start_ms * rate / 1000)
        end_frame = int(end_ms * rate / 1000)
        wf.setpos(min(start_frame, wf.getnframes()))
        frames = wf.readframes(max(end_frame - start_frame, 0))
    digest = hashlib.sha256(frames).hexdigest()
    if sampwidth == 1:
        energy = sum(abs(b - 128) for b in frames[: min(len(frames), 4000)])
    else:
        energy = sum(abs(int.from_bytes(frames[i:i+2], 'little', signed=True)) for i in range(0, min(len(frames), 8000), 2))
    return {'digest': digest, 'energy': energy, 'rate': rate, 'channels': n_channels, 'bytes_read': len(frames)}


def extract_matcher_fingerprint(path: Path, start_ms: int, end_ms: int) -> dict | None:
    try:
        matcher = ChromaprintMatcher(sr=16000)
        y, _ = load_audio_mono(str(path), sr=matcher.sr)
        start = int(start_ms * matcher.sr / 1000)
        end = int(end_ms * matcher.sr / 1000)
        segment = y[start:end]
        hashes = matcher.extract_hashes(segment)
        digest = hashlib.sha256(json.dumps(hashes[:128]).encode('utf-8')).hexdigest()
        return {
            'fingerprint_value': digest[:32],
            'checksum': f'chromaprint:{digest[:16]}',
            'metadata_json': {'hash_count': len(hashes), 'hash_sample': hashes[:8]},
        }
    except Exception:
        return None


def load_mert_runtime():
    global _MERT_RUNTIME
    if _MERT_RUNTIME is not None:
        return _MERT_RUNTIME
    import torch
    import torchaudio
    from transformers import Wav2Vec2FeatureExtractor, AutoModel
    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MERT_MODEL_ID, trust_remote_code=True)
    model = AutoModel.from_pretrained(MERT_MODEL_ID, trust_remote_code=True)
    model.eval()
    _MERT_RUNTIME = {
        'torch': torch,
        'torchaudio': torchaudio,
        'feature_extractor': feature_extractor,
        'model': model,
        'sample_rate': int(feature_extractor.sampling_rate),
        'hidden_size': int(getattr(model.config, 'hidden_size', 768)),
    }
    return _MERT_RUNTIME


def extract_mert_embedding(asset_path: Path, start_ms: int, end_ms: int) -> dict:
    rt = load_mert_runtime()
    torch = rt['torch']
    samples, sr = load_audio_mono(str(asset_path), sr=rt['sample_rate'])
    samples = np.asarray(samples, dtype=np.float32)
    start_frame = int(start_ms * sr / 1000)
    end_frame = int(end_ms * sr / 1000)
    segment = samples[start_frame:end_frame]
    if segment.size == 0:
        raise ValueError('empty segment for MERT extraction')
    inputs = rt['feature_extractor'](
        segment,
        sampling_rate=sr,
        return_tensors='pt',
    )
    with torch.no_grad():
        outputs = rt['model'](**inputs)
    emb = outputs.last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy().astype(np.float32)
    digest = hashlib.sha256(emb.tobytes()).hexdigest()
    return {
        'embedding_dim': int(emb.shape[0]),
        'embedding_uri': f"inline-mert://{digest[:16]}:{start_ms}:{end_ms}",
        'vector_table_name': f"audio_embedding_vector_{int(emb.shape[0])}_placeholder",
        'checksum': f"emb:{digest[:16]}",
        'metadata_json': {
            'semantic_backend': 'mert_runtime',
            'embedding_preview': [float(x) for x in emb[:8]],
            'model_id': MERT_MODEL_ID,
            'sample_rate': sr,
        },
    }


def build_semantic_feature(asset_path: Path, stats: dict, start_ms: int, end_ms: int, runtime_ok: bool, missing: list[str]) -> dict:
    if runtime_ok:
        try:
            mert = extract_mert_embedding(asset_path, start_ms, end_ms)
            return {
                'feature_type': 'embedding',
                'model_name': 'mert-v1-95m',
                'model_version': 'hf-main',
                'feature_set_name': 'mert_5s_hop2.5_v1',
                'feature_schema_ver': 'v1',
                'embedding_dim': mert['embedding_dim'],
                'embedding_uri': mert['embedding_uri'],
                'vector_table_name': mert['vector_table_name'],
                'checksum': mert['checksum'],
                'metadata_json': mert['metadata_json'],
            }
        except Exception as exc:
            return {
                'feature_type': 'embedding',
                'model_name': 'semantic_runtime_ready_placeholder',
                'model_version': 'awaiting_real_adapter',
                'feature_set_name': 'semantic_runtime_ready_5s',
                'feature_schema_ver': 'v1',
                'embedding_dim': 8,
                'embedding_uri': f"runtime-ready://{stats['digest'][:16]}:{start_ms}:{end_ms}",
                'vector_table_name': 'audio_embedding_vector_8_placeholder',
                'checksum': f"emb:{stats['digest'][:16]}",
                'metadata_json': {'semantic_backend': 'runtime_ready_placeholder', 'runtime_error': str(exc)},
            }
    return {
        'feature_type': 'embedding',
        'model_name': 'local_wavehash_embed',
        'model_version': 'v1',
        'feature_set_name': 'wavehash_embed_5s',
        'feature_schema_ver': 'v1',
        'embedding_dim': 8,
        'embedding_uri': f"inline://{stats['digest'][:16]}:{start_ms}:{end_ms}",
        'vector_table_name': 'audio_embedding_vector_8_placeholder',
        'checksum': f"emb:{stats['digest'][:16]}",
        'metadata_json': {
            'energy': stats['energy'],
            'rate': stats['rate'],
            'channels': stats['channels'],
            'semantic_backend': 'local_fallback',
            'runtime_missing': missing,
        },
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-manifest', required=True)
    parser.add_argument('--output-manifest', required=True)
    parser.add_argument('--report-output')
    args = parser.parse_args()

    in_path = Path(args.input_manifest).resolve()
    out_path = Path(args.output_manifest).resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    report_path = Path(args.report_output).resolve() if args.report_output else None
    if report_path:
        report_path.parent.mkdir(parents=True, exist_ok=True)

    runtime_ok, missing_runtime = semantic_runtime_available()

    rows = []
    feature_count = 0
    wav_windows_seen = 0
    matcher_fp_count = 0
    fallback_fp_count = 0
    semantic_runtime_ready_count = 0
    semantic_fallback_count = 0

    for row in load_jsonl(in_path):
        asset = row['asset']
        asset_path = Path(asset['storage_uri'])
        for window in row.get('windows', []):
            features = window.setdefault('features', [])
            if asset_path.suffix.lower() == '.wav' and asset_path.exists():
                wav_windows_seen += 1
                stats = read_wav_stats(asset_path, window['start_ms'], window['end_ms'])
                matcher_fp = extract_matcher_fingerprint(asset_path, window['start_ms'], window['end_ms'])
                if matcher_fp is not None:
                    fp = {
                        'feature_type': 'fingerprint',
                        'model_name': 'chromaprint_matcher',
                        'model_version': 'phase1_local',
                        'feature_set_name': 'chromaprint_matcher_5s',
                        'fingerprint_value': matcher_fp['fingerprint_value'],
                        'checksum': matcher_fp['checksum'],
                        'metadata_json': matcher_fp['metadata_json'],
                    }
                    matcher_fp_count += 1
                else:
                    fp = {
                        'feature_type': 'fingerprint',
                        'model_name': 'local_wavehash',
                        'model_version': 'v1',
                        'feature_set_name': 'wavehash_5s',
                        'fingerprint_value': stats['digest'][:32],
                        'checksum': f"fp:{stats['digest'][:16]}",
                        'metadata_json': {'energy': stats['energy'], 'bytes_read': stats['bytes_read']},
                    }
                    fallback_fp_count += 1

                emb = build_semantic_feature(asset_path, stats, window['start_ms'], window['end_ms'], runtime_ok, missing_runtime)
                if runtime_ok:
                    semantic_runtime_ready_count += 1
                else:
                    semantic_fallback_count += 1

                features.extend([fp, emb])
                feature_count += 2
        rows.append(row)

    out_path.write_text('\n'.join(json.dumps(r, ensure_ascii=False) for r in rows) + ('\n' if rows else ''))
    report = {
        'input_manifest': str(in_path),
        'output_manifest': str(out_path),
        'rows': len(rows),
        'wav_windows_seen': wav_windows_seen,
        'features_added': feature_count,
        'matcher_fingerprint_count': matcher_fp_count,
        'fallback_fingerprint_count': fallback_fp_count,
        'semantic_runtime_available': runtime_ok,
        'semantic_runtime_missing': missing_runtime,
        'semantic_runtime_ready_count': semantic_runtime_ready_count,
        'semantic_fallback_count': semantic_fallback_count,
    }
    if report_path:
        report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2))
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())