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

import argparse
import json
import math
import os
from pathlib import Path
from typing import Any

from psycopg import sql

from _job_common import connect, emit_payload, fetch_job_context, resolve_scope_summary, update_job_status

VECTOR_TABLE_BY_DIM = {
    192: 'audio_embedding_vector_192',
    768: 'audio_embedding_vector_768',
}
MODEL_RUNTIME_REQUIREMENTS = {
    'mert': ('numpy', 'torch', 'torchaudio', 'transformers'),
    'muq': ('numpy', 'torch', 'torchaudio', 'transformers'),
    'ecapa': ('numpy', 'torch', 'torchaudio', 'speechbrain'),
}
ALLOWED_VECTOR_TABLES = set(VECTOR_TABLE_BY_DIM.values())


def fetch_scope_windows(conn, target_scope: str) -> list[dict[str, object]]:
    if not target_scope.startswith('reference_set:'):
        raise SystemExit(f'unsupported target_scope for embedding worker: {target_scope}')
    set_name = target_scope.split(':', 1)[1]
    rows = conn.execute(
        """
        SELECT
            aw.window_id,
            aw.asset_id,
            aw.window_index,
            aw.start_sec,
            aw.end_sec,
            aw.duration_sec,
            aw.recording_id,
            aw.work_id,
            aw.canonical_song_id,
            ra.storage_uri,
            ra.ingest_status,
            aw.active_for_index
        FROM reference_set_registry rs
        JOIN reference_set_member rsm ON rsm.reference_set_id = rs.reference_set_id
        JOIN audio_window aw ON aw.recording_id = rsm.recording_id
        JOIN recording_asset ra ON ra.asset_id = aw.asset_id
        WHERE rs.set_name = %s
          AND aw.active_for_index = TRUE
          AND ra.ingest_status = 'ready'
        ORDER BY aw.window_id;
        """,
        (set_name,),
    ).fetchall()
    return [
        {
            'window_id': int(row[0]),
            'asset_id': int(row[1]),
            'window_index': int(row[2]),
            'start_sec': float(row[3]),
            'end_sec': float(row[4]),
            'duration_sec': float(row[5]),
            'recording_id': int(row[6]),
            'work_id': int(row[7]),
            'canonical_song_id': int(row[8]),
            'storage_uri': row[9],
            'ingest_status': row[10],
            'active_for_index': bool(row[11]),
        }
        for row in rows
    ]


def detect_runtime(model_name: str) -> dict[str, Any]:
    checks: dict[str, Any] = {'model_name': model_name, 'requirements': list(MODEL_RUNTIME_REQUIREMENTS.get(model_name, ('numpy',)))}
    availability: dict[str, bool] = {}
    missing: list[str] = []
    for package_name in checks['requirements']:
        try:
            __import__(package_name)
            availability[package_name] = True
        except Exception:  # noqa: BLE001
            availability[package_name] = False
            missing.append(package_name)
    checks['availability'] = availability
    checks['missing_dependencies'] = missing
    checks['ready'] = not missing
    return checks


def validate_vector_table(conn, vector_table: str | None, expected_dim: int | None) -> dict[str, Any]:
    payload = {
        'requested_vector_table': vector_table,
        'expected_dim': expected_dim,
        'allowed_vector_tables': sorted(ALLOWED_VECTOR_TABLES),
        'resolved': False,
        'table_exists': False,
        'reason': None,
    }
    if not vector_table:
        payload['reason'] = 'missing_vector_table'
        return payload
    if vector_table not in ALLOWED_VECTOR_TABLES:
        payload['reason'] = 'vector_table_not_allowlisted'
        return payload
    dim_from_table = 192 if vector_table.endswith('_192') else 768 if vector_table.endswith('_768') else None
    if expected_dim is not None and dim_from_table is not None and dim_from_table != expected_dim:
        payload['reason'] = 'vector_table_dim_mismatch'
        return payload
    row = conn.execute('SELECT to_regclass(%s);', (vector_table,)).fetchone()
    payload['table_exists'] = bool(row and row[0])
    if not payload['table_exists']:
        payload['reason'] = 'vector_table_missing_in_schema'
        return payload
    payload['resolved'] = True
    return payload


def build_artifact_path(artifact_dir: Path, *, extraction_job_id: int, window_id: int) -> Path:
    artifact_dir.mkdir(parents=True, exist_ok=True)
    return artifact_dir / f'job{extraction_job_id}_window{window_id}.json'


def vector_literal(values: list[float]) -> str:
    return '[' + ','.join(f'{value:.10f}' for value in values) + ']'


def compute_vector_norm(values: list[float]) -> float:
    return math.sqrt(sum(value * value for value in values))


def upsert_audio_embedding(
    conn,
    *,
    feature_set_id: int,
    extraction_job_id: int,
    vector_table: str,
    window: dict[str, object],
    embedding_uri: str,
    embedding: list[float],
    checksum: str | None,
    metadata_json: dict[str, object],
) -> tuple[int, str]:
    row = 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, %s, %s, %s, %s, %s,
            %s, %s, %s, %s, %s,
            TRUE, %s::jsonb
        )
        ON CONFLICT (feature_set_id, window_id) WHERE window_id IS NOT NULL
        DO UPDATE SET
            extraction_job_id = EXCLUDED.extraction_job_id,
            asset_id = EXCLUDED.asset_id,
            recording_id = EXCLUDED.recording_id,
            work_id = EXCLUDED.work_id,
            canonical_song_id = EXCLUDED.canonical_song_id,
            embedding_storage_mode = EXCLUDED.embedding_storage_mode,
            embedding_uri = EXCLUDED.embedding_uri,
            vector_norm = EXCLUDED.vector_norm,
            checksum = EXCLUDED.checksum,
            is_indexed = EXCLUDED.is_indexed,
            metadata_json = EXCLUDED.metadata_json
        RETURNING embedding_id, xmax = 0 AS inserted;
        """,
        (
            feature_set_id,
            extraction_job_id,
            window['asset_id'],
            window['window_id'],
            window['recording_id'],
            window['work_id'],
            window['canonical_song_id'],
            'pgvector_inline',
            embedding_uri,
            compute_vector_norm(embedding),
            checksum,
            json.dumps(metadata_json, ensure_ascii=False),
        ),
    ).fetchone()
    embedding_id = int(row[0])
    inserted = bool(row[1])
    conn.execute(
        sql.SQL(
            """
            INSERT INTO {vector_table} (embedding_id, embedding)
            VALUES (%s, %s::vector)
            ON CONFLICT (embedding_id)
            DO UPDATE SET embedding = EXCLUDED.embedding;
            """
        ).format(vector_table=sql.Identifier(vector_table)),
        (embedding_id, vector_literal(embedding)),
    )
    return embedding_id, 'inserted' if inserted else 'updated'


def fail_job(
    conn,
    *,
    job,
    blockers: list[str],
    output_target: str,
    resolved_vector_table: str | None,
    artifact_dir: Path,
    scope: dict[str, Any],
    scope_windows: list[dict[str, object]],
    missing_windows: list[dict[str, object]],
    runtime_report: dict[str, Any],
    vector_table_report: dict[str, Any],
) -> dict[str, Any]:
    return update_job_status(
        conn,
        job.extraction_job_id,
        status='failed',
        expected_status='running',
        output_count=0,
        metadata_patch={
            'worker': 'run_embedding_job',
            'output_target': output_target,
            'vector_table': resolved_vector_table,
            'dry_run': False,
            'write_target_table': output_target,
            'artifact_dir': str(artifact_dir),
            'execution_mode': 'preflight_failure',
            'failure_reason': 'preflight_failed',
            'preflight_blockers': blockers,
            'scope_window_count': len(scope_windows),
            'missing_window_count': len(missing_windows),
            'missing_window_samples': missing_windows[:5],
            'runtime_report': runtime_report,
            'vector_table_report': vector_table_report,
            'target_scope_summary': scope,
        },
        set_finished_at=True,
    )


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--dsn', default=os.environ.get('PG_DSN'))
    ap.add_argument('--schema', default=os.environ.get('PG_SCHEMA', 'acr_test'))
    ap.add_argument('--job-id', type=int, default=int(os.environ.get('EXTRACTION_JOB_ID', '0')))
    ap.add_argument('--model-name', default=os.environ.get('MODEL_NAME'))
    ap.add_argument('--model-version', default=os.environ.get('MODEL_VERSION'))
    ap.add_argument('--vector-table', default=os.environ.get('VECTOR_TABLE'))
    ap.add_argument('--output-target', default=os.environ.get('OUTPUT_TARGET', 'audio_embedding'))
    ap.add_argument('--complete-dry-run', action='store_true')
    ap.add_argument('--artifact-dir', default=os.environ.get('ARTIFACT_DIR', 'data/pgvector_eval/music20/phase1_embeddings'))
    ap.add_argument('--output')
    args = ap.parse_args()

    if not args.dsn:
        raise SystemExit('missing --dsn or PG_DSN')
    if not args.job_id:
        raise SystemExit('missing --job-id or EXTRACTION_JOB_ID')

    artifact_dir = Path(args.artifact_dir)

    with connect(args.dsn, args.schema) as conn:
        job = fetch_job_context(conn, args.job_id)
        if job.model_name == 'chromaprint':
            raise SystemExit(f'feature_extraction_job={args.job_id} is not an embedding job')
        if job.feature_name != 'semantic_embedding' or job.feature_level != 'window':
            raise SystemExit(
                f'feature_extraction_job={args.job_id} does not match embedding feature contract: '
                f'{job.feature_name}/{job.feature_level}'
            )
        if args.model_name and job.model_name != args.model_name:
            raise SystemExit(f'model mismatch: job={job.model_name} cli={args.model_name}')
        if args.model_version and job.model_version != args.model_version:
            raise SystemExit(f'model version mismatch: job={job.model_version} cli={args.model_version}')

        expected_dim = job.embedding_dim or job.output_embedding_dim
        resolved_vector_table = args.vector_table or VECTOR_TABLE_BY_DIM.get(expected_dim or -1)
        scope = resolve_scope_summary(conn, job.target_scope)
        scope_windows = fetch_scope_windows(conn, job.target_scope)
        runtime_report = detect_runtime(job.model_name)
        vector_table_report = validate_vector_table(conn, resolved_vector_table, expected_dim)

        running = update_job_status(
            conn,
            job.extraction_job_id,
            status='running',
            expected_status='pending',
            input_count=len(scope_windows),
            metadata_patch={
                'worker': 'run_embedding_job',
                'output_target': args.output_target,
                'vector_table': resolved_vector_table,
                'dry_run': bool(args.complete_dry_run),
                'target_scope_summary': scope,
                'execution_mode': 'dry_run' if args.complete_dry_run else 'preflight',
                'runtime_report': runtime_report,
                'vector_table_report': vector_table_report,
                'scope_window_count': len(scope_windows),
            },
            set_started_at=True,
        )
        completed = None
        failed = None
        processed_windows: list[dict[str, object]] = []

        if args.complete_dry_run:
            completed = update_job_status(
                conn,
                job.extraction_job_id,
                status='completed',
                expected_status='running',
                output_count=0,
                metadata_patch={
                    'worker': 'run_embedding_job',
                    'output_target': args.output_target,
                    'vector_table': resolved_vector_table,
                    'dry_run': True,
                    'dry_run_result': 'completed_without_feature_write',
                    'write_target_table': args.output_target,
                    'scope_window_count': len(scope_windows),
                    'runtime_report': runtime_report,
                    'vector_table_report': vector_table_report,
                },
                set_finished_at=True,
            )
        else:
            missing_windows: list[dict[str, object]] = []
            for window in scope_windows:
                asset_path = Path(str(window['storage_uri']))
                if not asset_path.exists():
                    missing_windows.append({
                        'window_id': window['window_id'],
                        'asset_id': window['asset_id'],
                        'storage_uri': str(asset_path),
                        'reason': 'missing_audio',
                    })

            blockers: list[str] = []
            if missing_windows:
                blockers.append('unreadable_audio_assets')
            if not vector_table_report['resolved']:
                blockers.append(str(vector_table_report['reason']))
            if not runtime_report['ready']:
                blockers.append('model_runtime_unavailable')

            if blockers:
                failed = fail_job(
                    conn,
                    job=job,
                    blockers=blockers,
                    output_target=args.output_target,
                    resolved_vector_table=resolved_vector_table,
                    artifact_dir=artifact_dir,
                    scope=scope,
                    scope_windows=scope_windows,
                    missing_windows=missing_windows,
                    runtime_report=runtime_report,
                    vector_table_report=vector_table_report,
                )
            else:
                failed = update_job_status(
                    conn,
                    job.extraction_job_id,
                    status='failed',
                    expected_status='running',
                    output_count=0,
                    metadata_patch={
                        'worker': 'run_embedding_job',
                        'output_target': args.output_target,
                        'vector_table': resolved_vector_table,
                        'dry_run': False,
                        'write_target_table': args.output_target,
                        'artifact_dir': str(artifact_dir),
                        'execution_mode': 'write_attempt',
                        'failure_reason': 'encoder_inference_not_implemented',
                        'scope_window_count': len(scope_windows),
                        'runtime_report': runtime_report,
                        'vector_table_report': vector_table_report,
                        'next_expected_step': 'replace the guarded failure path with real model inference while keeping the same upsert contract',
                    },
                    set_finished_at=True,
                )

    emit_payload(
        {
            'worker': 'run_embedding_job',
            'schema': args.schema,
            'job': job.__dict__,
            'target_scope_summary': scope,
            'scope_window_count': len(scope_windows),
            'status_after_start': running,
            'status_after_complete': completed,
            'status_after_failed': failed,
            'resolved_vector_table': resolved_vector_table,
            'vector_table_report': vector_table_report,
            'runtime_report': runtime_report,
            'processed_windows': processed_windows,
            'notes': [
                'this worker now validates planner -> job -> scope windows -> PostgreSQL failure semantics',
                'real model inference should replace the guarded failure path without changing the job contract or idempotent upsert keys',
            ],
        },
        args.output,
    )


if __name__ == '__main__':
    main()