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

import argparse
import json
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from statistics import median
from typing import Any

import psycopg

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SCHEMA_SQL = ROOT / 'sql' / 'acr_pg_schema_v2.sql'
DEFAULT_REFERENCE = ROOT / 'data' / 'pgvector_eval' / 'music20' / 'reference_embeddings.jsonl'
DEFAULT_QUERY = ROOT / 'data' / 'pgvector_eval' / 'music20' / 'query_embeddings.jsonl'
DEFAULT_OUTPUT = ROOT / 'data' / 'pgvector_eval' / 'music20' / 'live_pgvector_report.json'


@dataclass
class EntityIds:
    canonical_song_id: int
    work_id: int
    recording_id: int
    asset_id: int
    window_id: int
    embedding_id: int


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines() if line.strip()]


def pad_embedding(vec: list[float], target_dim: int = 192) -> list[float]:
    if len(vec) > target_dim:
        raise ValueError(f'embedding dim {len(vec)} > target {target_dim}')
    if len(vec) == target_dim:
        return vec
    return vec + [0.0] * (target_dim - len(vec))


def vec_literal(vec: list[float]) -> str:
    return '[' + ','.join(f'{x:.10f}' for x in vec) + ']'


def compute_metrics(ranks: list[int], topk: int) -> dict[str, Any]:
    if not ranks:
        return {'count': 0}
    return {
        'count': len(ranks),
        'top1': round(sum(1 for r in ranks if r == 1) / len(ranks), 6),
        'top3': round(sum(1 for r in ranks if r <= 3) / len(ranks), 6),
        f'top{topk}': round(sum(1 for r in ranks if r <= topk) / len(ranks), 6),
        'mrr': round(sum(1.0 / r for r in ranks) / len(ranks), 6),
        'mean_rank': round(sum(ranks) / len(ranks), 4),
        'median_rank': median(ranks),
    }


def aggregate_song_scores(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[row['song_id']].append(row)
    ranked = []
    for song_id, vals in grouped.items():
        vals.sort(key=lambda x: x['score'], reverse=True)
        scores = [v['score'] for v in vals]
        max_sim = scores[0]
        top3_avg = sum(scores[:3]) / min(3, len(scores))
        vote = len(scores)
        combined = 0.6 * max_sim + 0.3 * top3_avg + 0.1 * min(vote / 10.0, 1.0)
        ranked.append({
            'song_id': song_id,
            'canonical_song_id': vals[0]['canonical_song_id'],
            'evidence_window_id': vals[0]['window_id'],
            'combined_score': combined,
            'max_sim': max_sim,
            'top3_avg': top3_avg,
            'vote': vote,
        })
    ranked.sort(key=lambda x: x['combined_score'], reverse=True)
    return ranked


def reset_schema(conn: psycopg.Connection, schema: str) -> None:
    conn.execute(f'DROP SCHEMA IF EXISTS {schema} CASCADE;')
    conn.execute(f'CREATE SCHEMA {schema};')
    conn.execute(f'SET search_path TO {schema}, public;')


def apply_schema(conn: psycopg.Connection, schema_sql: Path) -> None:
    sql_text = schema_sql.read_text(encoding='utf-8')
    conn.execute(sql_text)


def seed_registry(conn: psycopg.Connection) -> tuple[int, int, int, int]:
    model_id = conn.execute(
        """
        INSERT INTO model_registry (
            model_name, model_family, model_version, model_source, model_uri,
            license_name, input_sample_rate, default_window_sec, default_hop_sec,
            output_embedding_dim, pooling_supported, metadata_json
        ) VALUES (
            'local_chroma24', 'chroma_baseline', 'v1', 'repo-local-eval',
            'acr-engine/scripts/live_pgvector_music20_eval.py', 'internal-eval',
            22050, 8.0, 8.0, 24, ARRAY['mean_std'],
            '{"storage_padding":"zero-pad to vector(192) for pgvector compatibility"}'::jsonb
        )
        ON CONFLICT (model_name, model_version) DO UPDATE
        SET updated_at = NOW()
        RETURNING model_id;
        """
    ).fetchone()[0]

    feature_set_id = conn.execute(
        """
        INSERT INTO feature_set_registry (
            model_id, feature_name, feature_level, extraction_granularity,
            window_sec, hop_sec, embedding_dim, pooling_strategy, layer_selection,
            normalize_l2, distance_metric, quantization_type, feature_schema_version,
            config_json, status
        ) VALUES (
            %s, 'chroma24_songid_eval', 'window', 'window',
            8.0, 8.0, 24, 'mean_std', 'na', TRUE, 'cosine', NULL, 'v1',
            '{"physical_storage":"audio_embedding_vector_192","padding":"zero"}'::jsonb,
            'active'
        )
        RETURNING feature_set_id;
        """,
        (model_id,),
    ).fetchone()[0]

    reference_set_id = conn.execute(
        """
        INSERT INTO reference_set_registry (set_name, description, encoder_scope, status, metadata_json)
        VALUES (
            'music20_live_reference',
            '20-song local live pgvector evaluation reference set',
            'local_chroma24',
            'active',
            '{"purpose":"live_pgvector_music20_eval"}'::jsonb
        )
        ON CONFLICT (set_name) DO UPDATE SET updated_at = NOW()
        RETURNING reference_set_id;
        """
    ).fetchone()[0]

    retrieval_index_id = conn.execute(
        """
        INSERT INTO retrieval_index_registry (
            feature_set_id, index_name, index_backend, index_type, storage_uri,
            shard_no, row_count, index_status, config_json, built_at
        ) VALUES (
            %s, 'music20_live_pgvector_hnsw', 'pgvector', 'hnsw_cosine',
            'postgres://d2@127.0.0.1/d2#acr_test.audio_embedding_vector_192',
            0, 0, 'active', '{"physical_dim":192,"logical_dim":24}'::jsonb, NOW()
        )
        RETURNING retrieval_index_id;
        """,
        (feature_set_id,),
    ).fetchone()[0]

    return model_id, feature_set_id, reference_set_id, retrieval_index_id


def ingest_references(conn: psycopg.Connection, refs: list[dict[str, Any]], feature_set_id: int, reference_set_id: int) -> dict[str, EntityIds]:
    entities: dict[str, EntityIds] = {}
    for idx, row in enumerate(refs):
        song_id = str(row['song_id'])
        canonical_song_id = conn.execute(
            """
            INSERT INTO canonical_song (biz_song_code, title, title_norm, primary_artist, primary_artist_norm, rights_status, metadata_json)
            VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb)
            RETURNING canonical_song_id;
            """,
            (song_id, f'Song {song_id}', f'song {song_id}', f'Artist {song_id}', f'artist {song_id}', 'protected', json.dumps({'source': 'music20_live_eval'})),
        ).fetchone()[0]
        work_id = conn.execute(
            """
            INSERT INTO work (canonical_song_id, work_code, work_title, work_title_norm, composer, publisher, metadata_json)
            VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb)
            RETURNING work_id;
            """,
            (canonical_song_id, f'work-{song_id}', f'Song {song_id}', f'song {song_id}', f'Composer {song_id}', 'Unknown', json.dumps({'note': '1:1 work for eval'})),
        ).fetchone()[0]
        recording_id = conn.execute(
            """
            INSERT INTO recording (
                work_id, canonical_song_id, recording_code, recording_title, artist_name,
                album_name, version_type, is_reference, reference_priority, duration_sec, metadata_json
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, TRUE, %s, %s, %s::jsonb)
            RETURNING recording_id;
            """,
            (work_id, canonical_song_id, f'rec-{song_id}', f'Song {song_id} Reference', f'Artist {song_id}', 'music20', 'master_reference', 100 + idx, 8.0, json.dumps({'source_audio_path': row['audio_path']})),
        ).fetchone()[0]
        asset_id = conn.execute(
            """
            INSERT INTO recording_asset (
                recording_id, asset_role, storage_uri, storage_scheme, file_ext, mime_type,
                sample_rate, channels, codec_name, duration_sec, normalized_storage_uri,
                ingest_status, metadata_json
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb)
            RETURNING asset_id;
            """,
            (recording_id, 'reference_audio', row['audio_path'], 'file', Path(row['audio_path']).suffix.lstrip('.'), 'audio/wav', 22050, 1, 'pcm_s16le', 8.0, row['audio_path'], 'ready', json.dumps({'type': 'reference'})),
        ).fetchone()[0]
        window_id = conn.execute(
            """
            INSERT INTO audio_window (
                asset_id, recording_id, work_id, canonical_song_id,
                window_index, start_sec, end_sec, duration_sec,
                segment_role, segment_type, quality_score, active_for_index, metadata_json
            ) VALUES (%s, %s, %s, %s, 0, 0.0, 8.0, 8.0, 'reference', 'full_clip', 1.0, TRUE, %s::jsonb)
            RETURNING window_id;
            """,
            (asset_id, recording_id, work_id, canonical_song_id, json.dumps({'source_audio_path': row['audio_path']})),
        ).fetchone()[0]
        embedding_id = conn.execute(
            """
            INSERT INTO audio_embedding (
                feature_set_id, extraction_job_id, asset_id, window_id, recording_id, work_id,
                canonical_song_id, embedding_storage_mode, embedding_uri, vector_norm, checksum,
                is_indexed, metadata_json
            ) VALUES (%s, NULL, %s, %s, %s, %s, %s, %s, NULL, %s, NULL, TRUE, %s::jsonb)
            RETURNING embedding_id;
            """,
            (feature_set_id, asset_id, window_id, recording_id, work_id, canonical_song_id, 'pgvector_inline_192_padded', 1.0, json.dumps({'logical_embedding_dim': len(row['embedding'])})),
        ).fetchone()[0]
        conn.execute(
            'INSERT INTO audio_embedding_vector_192 (embedding_id, embedding) VALUES (%s, %s::vector);',
            (embedding_id, vec_literal(pad_embedding(row['embedding']))),
        )
        conn.execute(
            'INSERT INTO reference_set_member (reference_set_id, recording_id, member_role) VALUES (%s, %s, %s);',
            (reference_set_id, recording_id, 'hot_reference'),
        )
        entities[song_id] = EntityIds(canonical_song_id, work_id, recording_id, asset_id, window_id, embedding_id)
    return entities


def _expect_insert_failure(
    conn: psycopg.Connection,
    sql: str,
    params: tuple[Any, ...],
    case_name: str,
    expected_guard: str,
) -> dict[str, Any]:
    try:
        with conn.transaction():
            conn.execute(sql, params)
        return {
            'passed': False,
            'case': case_name,
            'expected_guard': expected_guard,
            'note': 'bad lineage insert unexpectedly succeeded',
        }
    except Exception as exc:
        return {
            'passed': True,
            'case': case_name,
            'expected_guard': expected_guard,
            'error_type': type(exc).__name__,
            'message': str(exc).splitlines()[0],
        }


def run_lineage_negative_tests(conn: psycopg.Connection, entity: EntityIds, feature_set_id: int) -> dict[str, Any]:
    recording_case = _expect_insert_failure(
        conn,
        """
        INSERT INTO recording (
            work_id, canonical_song_id, recording_code, recording_title, artist_name,
            album_name, version_type, is_reference, reference_priority, duration_sec, metadata_json
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, FALSE, %s, %s, %s::jsonb);
        """,
        (
            entity.work_id,
            entity.canonical_song_id + 999999,
            f'bad-rec-{entity.recording_id}',
            'Bad Recording Lineage',
            'Bad Artist',
            'bad-album',
            'bad_lineage_probe',
            9999,
            8.0,
            json.dumps({'probe': 'recording_lineage_negative'}),
        ),
        'recording_lineage_mismatch',
        'recording.canonical_song_id must equal work.canonical_song_id',
    )

    audio_window_case = _expect_insert_failure(
        conn,
        """
        INSERT INTO audio_window (
            asset_id, recording_id, work_id, canonical_song_id, window_index,
            start_sec, end_sec, duration_sec, segment_role, segment_type, quality_score, active_for_index
        ) VALUES (%s, %s, %s, %s, 999, 0.0, 8.0, 8.0, 'reference', 'bad_lineage', 0.0, TRUE);
        """,
        (entity.asset_id, entity.recording_id + 999999, entity.work_id, entity.canonical_song_id),
        'audio_window_lineage_mismatch',
        'audio_window recording/work/song lineage must match recording_asset + recording parents',
    )

    audio_embedding_case = _expect_insert_failure(
        conn,
        """
        INSERT INTO audio_embedding (
            feature_set_id, extraction_job_id, asset_id, window_id, recording_id, work_id,
            canonical_song_id, embedding_storage_mode, embedding_uri, vector_norm, checksum,
            is_indexed, metadata_json
        ) VALUES (%s, NULL, %s, %s, %s, %s, %s, %s, NULL, %s, NULL, FALSE, %s::jsonb);
        """,
        (
            feature_set_id,
            entity.asset_id,
            entity.window_id,
            entity.recording_id,
            entity.work_id,
            entity.canonical_song_id + 999999,
            'pgvector_inline_192_padded',
            1.0,
            json.dumps({'probe': 'audio_embedding_lineage_negative'}),
        ),
        'audio_embedding_lineage_mismatch',
        'audio_embedding recording/work/song lineage must match the parent audio_window',
    )

    cases = [recording_case, audio_window_case, audio_embedding_case]
    return {
        'passed': all(item['passed'] for item in cases),
        'cases': cases,
    }


def fetch_raw_candidates(conn: psycopg.Connection, feature_set_id: int, query_vec: list[float], topn: int) -> list[dict[str, Any]]:
    rows = conn.execute(
        """
        SELECT
            cs.biz_song_code AS song_id,
            ae.canonical_song_id,
            aw.window_id,
            1 - (aev.embedding <=> %s::vector) AS score
        FROM audio_embedding_vector_192 aev
        JOIN audio_embedding ae ON ae.embedding_id = aev.embedding_id
        JOIN canonical_song cs ON cs.canonical_song_id = ae.canonical_song_id
        JOIN audio_window aw ON aw.window_id = ae.window_id
        WHERE ae.feature_set_id = %s
        ORDER BY aev.embedding <=> %s::vector
        LIMIT %s;
        """,
        (vec_literal(pad_embedding(query_vec)), feature_set_id, vec_literal(pad_embedding(query_vec)), topn),
    ).fetchall()
    return [
        {
            'song_id': r[0],
            'canonical_song_id': r[1],
            'window_id': r[2],
            'score': float(r[3]),
        }
        for r in rows
    ]


def persist_candidates(conn: psycopg.Connection, query_id: str, retrieval_index_id: int, feature_set_id: int, ranked: list[dict[str, Any]], topk: int) -> None:
    for i, item in enumerate(ranked[:topk], start=1):
        conn.execute(
            """
            INSERT INTO retrieval_candidate (
                query_id, retrieval_index_id, feature_set_id, source_lane,
                candidate_level, candidate_id, evidence_window_id, raw_score,
                normalized_score, rank_no, metadata_json
            ) VALUES (%s, %s, %s, 'semantic', 'canonical_song', %s, %s, %s, %s, %s, %s::jsonb);
            """,
            (query_id, retrieval_index_id, feature_set_id, item['canonical_song_id'], item['evidence_window_id'], item['max_sim'], item['combined_score'], i, json.dumps({'vote': item['vote'], 'song_id': item['song_id']})),
        )


def persist_decision(conn: psycopg.Connection, query_id: str, ranked: list[dict[str, Any]]) -> None:
    top = ranked[0] if ranked else None
    conn.execute(
        """
        INSERT INTO match_decision (
            query_id, canonical_song_id, work_id, recording_id,
            decision_status, decision_score, decision_reason, metadata_json
        ) VALUES (%s, %s, NULL, NULL, %s, %s, %s, %s::jsonb);
        """,
        (
            query_id,
            top['canonical_song_id'] if top else None,
            'matched' if top else 'no_match',
            top['combined_score'] if top else None,
            'top1 semantic candidate from live pgvector eval' if top else 'no candidate',
            json.dumps({'top_song_id': top['song_id']} if top else {}),
        ),
    )


def evaluate_live(conn: psycopg.Connection, feature_set_id: int, retrieval_index_id: int, queries: list[dict[str, Any]], topn: int, topk: int) -> dict[str, Any]:
    by_type: dict[str, list[int]] = defaultdict(list)
    examples: dict[str, list[dict[str, Any]]] = defaultdict(list)
    confusion_focus: dict[str, dict[str, Any]] = {}

    for idx, q in enumerate(queries):
        qtype = str(q['query_type'])
        query_id = f'music20-q{idx:04d}-t{qtype}-song{q["song_id"]}'
        raw_rows = fetch_raw_candidates(conn, feature_set_id, q['embedding'], topn)
        ranked = aggregate_song_scores(raw_rows)
        gold = str(q['song_id'])
        rank = next((i + 1 for i, item in enumerate(ranked) if item['song_id'] == gold), len(ranked) + 1)
        by_type[qtype].append(rank)
        persist_candidates(conn, query_id, retrieval_index_id, feature_set_id, ranked, topk)
        persist_decision(conn, query_id, ranked)
        if len(examples[qtype]) < 5:
            examples[qtype].append({
                'query_id': query_id,
                'song_id': gold,
                'rank': rank,
                'top3': ranked[:3],
            })

    for qtype in ('7', '8', '16'):
        ranks = by_type.get(qtype, [])
        confusion_focus[qtype] = {
            'query_type': int(qtype),
            'metrics': compute_metrics(ranks, topk),
            'interpretation': {
                '7': 'light confusion / transformed query',
                '8': 'harder confusion bucket',
                '16': 'strong confusion or far-domain bucket',
            }[qtype],
        }

    all_ranks = [r for ranks in by_type.values() for r in ranks]
    return {
        'backend': 'postgresql+pgvector-live',
        'note': 'Reference embeddings are stored in schema v2; 24-d logical embeddings are zero-padded to vector(192) for physical storage.',
        'overall': compute_metrics(all_ranks, topk),
        'by_query_type': {qtype: compute_metrics(ranks, topk) for qtype, ranks in by_type.items()},
        'confusion_focus': confusion_focus,
        'examples': examples,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--dsn', required=True)
    ap.add_argument('--schema', default='acr_test')
    ap.add_argument('--schema-sql', default=str(DEFAULT_SCHEMA_SQL))
    ap.add_argument('--reference-embeddings-jsonl', default=str(DEFAULT_REFERENCE))
    ap.add_argument('--query-embeddings-jsonl', default=str(DEFAULT_QUERY))
    ap.add_argument('--output', default=str(DEFAULT_OUTPUT))
    ap.add_argument('--topn', type=int, default=20)
    ap.add_argument('--topk', type=int, default=10)
    ap.add_argument('--reset-schema', action='store_true')
    args = ap.parse_args()

    refs = load_jsonl(Path(args.reference_embeddings_jsonl))
    queries = load_jsonl(Path(args.query_embeddings_jsonl))

    with psycopg.connect(args.dsn, autocommit=True) as conn:
        if args.reset_schema:
            reset_schema(conn, args.schema)
        else:
            conn.execute(f'CREATE SCHEMA IF NOT EXISTS {args.schema};')
            conn.execute(f'SET search_path TO {args.schema}, public;')
        apply_schema(conn, Path(args.schema_sql))
        model_id, feature_set_id, reference_set_id, retrieval_index_id = seed_registry(conn)
        entities = ingest_references(conn, refs, feature_set_id, reference_set_id)
        lineage_check = run_lineage_negative_tests(conn, next(iter(entities.values())), feature_set_id)
        report = evaluate_live(conn, feature_set_id, retrieval_index_id, queries, args.topn, args.topk)
        conn.execute('UPDATE retrieval_index_registry SET row_count = %s WHERE retrieval_index_id = %s;', (len(refs), retrieval_index_id))
        counts = {
            'canonical_song': conn.execute('SELECT count(*) FROM canonical_song;').fetchone()[0],
            'work': conn.execute('SELECT count(*) FROM work;').fetchone()[0],
            'recording': conn.execute('SELECT count(*) FROM recording;').fetchone()[0],
            'recording_asset': conn.execute('SELECT count(*) FROM recording_asset;').fetchone()[0],
            'audio_window': conn.execute('SELECT count(*) FROM audio_window;').fetchone()[0],
            'audio_embedding': conn.execute('SELECT count(*) FROM audio_embedding;').fetchone()[0],
            'retrieval_candidate': conn.execute('SELECT count(*) FROM retrieval_candidate;').fetchone()[0],
            'match_decision': conn.execute('SELECT count(*) FROM match_decision;').fetchone()[0],
        }

    payload = {
        'schema': args.schema,
        'dsn_redacted': 'postgres://d2:***@127.0.0.1:5432/d2',
        'input': {
            'reference_embeddings_jsonl': args.reference_embeddings_jsonl,
            'query_embeddings_jsonl': args.query_embeddings_jsonl,
            'reference_count': len(refs),
            'query_count': len(queries),
        },
        'registry': {
            'model_id': model_id,
            'feature_set_id': feature_set_id,
            'reference_set_id': reference_set_id,
            'retrieval_index_id': retrieval_index_id,
        },
        'table_counts': counts,
        'lineage_negative_test': lineage_check,
        'evaluation': report,
    }

    out = Path(args.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(payload, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()