import argparse
import concurrent.futures as futures
import datetime as dt
import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Dict, Iterator, List, Optional


# Optional: allow running directly via `python -m` when inside repo
THIS_DIR = Path(__file__).resolve().parent


def read_jsonl(path: Path) -> Iterator[Dict]:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def write_jsonl(path: Path, rows: List[Dict], mode: str = "a") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open(mode, encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def setup_logging(log_fp: Path) -> None:
    log_fp.parent.mkdir(parents=True, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(log_fp, mode="a", encoding="utf-8"),
            logging.StreamHandler(sys.stdout),
        ],
    )


def ffprobe_duration_seconds(audio_path: Path) -> float:
    import subprocess

    result = subprocess.run(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            str(audio_path),
        ],
        capture_output=True,
        text=True,
        check=True,
    )
    return float(result.stdout.strip())


def _gemini_audio_part(local_fp: Path):
    # Reuse the modality used in our runner (opus accepted, sent as audio/mpeg part).
    from google.genai import types

    with open(local_fp, "rb") as f:
        audio_bytes = f.read()
    return types.Part.from_bytes(data=audio_bytes, mime_type="audio/mpeg")


def _load_google_creds(creds_path: str = "/home/vibert/data/credentials/google-3.json") -> None:
    """Set GOOGLE_APPLICATION_CREDENTIALS to a local JSON service account key."""
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_path


def call_gemini_enhanced_lyrics(
    local_audio: Path,
    prompt_text: str,
    project: str = "ml-gemini-455703",
    location: str = "us-central1",
    model: str = "gemini-2.5-pro",
    temperature: float = 0.3,
) -> str:
    from google import genai
    from google.genai import types

    client = genai.Client(vertexai=True, project=project, location=location)

    # Thinking mode requested: includeThoughts False, thinkingBudget 0 (no budget) on 2.5-pro.
    # Enable thinking mode without specifying budget (avoid INVALID_ARGUMENT on some models)
    config = types.GenerateContentConfig(
        candidateCount=1,
        temperature=temperature,
        max_output_tokens=65536,
        response_mime_type="text/plain",
        thinkingConfig=types.ThinkingConfig(
            includeThoughts=False,
        ),
    )

    audio_part = _gemini_audio_part(local_audio)
    # Prompt first, then audio part.
    resp = client.models.generate_content(
        model=model,
        contents=[prompt_text, audio_part],
        config=config,
    )

    # Extract plain text
    if not resp.candidates:
        raise RuntimeError("Gemini returned no candidates")
    parts = resp.candidates[0].content.parts
    if not parts:
        raise RuntimeError("Gemini returned empty content parts")
    text = parts[0].text if hasattr(parts[0], "text") else str(parts[0])
    return text.strip()


def ensure_local_file_from_s3(s3_uri: str, local_path: Path, timeout: int = 600) -> None:
    """If local file is missing, download from S3 to local_path using AWS CLI."""
    if local_path.exists():
        return

    if not s3_uri.startswith("s3://"):
        raise ValueError(f"Invalid s3 uri: {s3_uri}")

    import subprocess

    local_path.parent.mkdir(parents=True, exist_ok=True)
    logging.info(f"aws s3 cp {s3_uri} {local_path}")
    subprocess.run(["aws", "s3", "cp", s3_uri, str(local_path), "--no-progress"], check=True)


def process_one(
    row: Dict,
    prompt_text: str,
    audio_root: Path,
) -> Dict:
    """Process a single item: ensure audio, call Gemini, return output row."""
    _id = row.get("id")
    s3_filepath = row.get("s3_filepath")
    views = row.get("views")
    duration_s = row.get("duration_s")

    # Local filepath policy
    local_filepath = audio_root / f"{_id}.opus"
    if not local_filepath.exists():
        if not s3_filepath:
            raise FileNotFoundError(f"Missing local file and s3_filepath for id={_id}")
        ensure_local_file_from_s3(s3_filepath, local_filepath)

    # If duration missing, probe
    if not duration_s:
        try:
            duration_s = ffprobe_duration_seconds(local_filepath)
        except Exception:
            duration_s = None

    enhanced = call_gemini_enhanced_lyrics(local_filepath, prompt_text)

    return {
        "id": _id,
        "s3_filepath": s3_filepath,
        "local_filepath": str(local_filepath),
        "duration_s": duration_s,
        "views": views,
        "enhanced_lyrics_gemini": enhanced,
    }


def main():
    parser = argparse.ArgumentParser(description="Run Gemini captioning for Deezer dataset")
    parser.add_argument(
        "--input",
        type=Path,
        default=Path("/home/minz/llm_tagging/data/deezer_metas.jsonl"),
        help="Path to input JSONL (deezer metas)",
    )
    parser.add_argument(
        "--prompt",
        type=Path,
        default=THIS_DIR / "prompt.txt",
        help="Prompt text file path",
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("/home/vibert/data/deezer_captions") / dt.datetime.now().strftime("%Y%m%d_%H%M%S"),
        help="Output directory (default uses timestamp)",
    )
    parser.add_argument(
        "--min-views",
        type=int,
        default=100_000,
        help="Only process rows with views >= this value",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Limit number of items (e.g., 200 for pilot)",
    )
    parser.add_argument(
        "--test-first-10",
        action="store_true",
        help="Process first 10 items only (preliminary test)",
    )
    parser.add_argument(
        "--max-workers",
        type=int,
        default=50,
        help="Parallelism (recommended: 50)",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=50,
        help="Results are saved every N successes",
    )
    parser.add_argument(
        "--google-creds",
        type=Path,
        default=Path("/home/vibert/data/credentials/google-3.json"),
        help="Path to Google service account JSON for Vertex AI",
    )
    parser.add_argument(
        "--shard-index",
        type=int,
        default=None,
        help="Index of this shard (0..num_shards-1)",
    )
    parser.add_argument(
        "--num-shards",
        type=int,
        default=None,
        help="Total number of shards to partition selected items",
    )

    args = parser.parse_args()

    outdir = args.outdir
    outdir.mkdir(parents=True, exist_ok=True)
    setup_logging(outdir / "progress.log")
    logging.info("Starting Deezer Gemini captioning run")
    logging.info(f"Input: {args.input}")
    logging.info(f"Outdir: {outdir}")
    logging.info(f"Min views: {args.min_views}")

    # Save run meta for reproducibility
    meta = {
        "input": str(args.input),
        "outdir": str(outdir),
        "min_views": args.min_views,
        "limit": args.limit,
        "max_workers": args.max_workers,
        "chunk_size": args.chunk_size,
        "ts": dt.datetime.now().isoformat(),
    }
    with (outdir / "run_meta.json").open("w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    # Load credentials and prompt
    _load_google_creds(str(args.google_creds))
    prompt_text = args.prompt.read_text(encoding="utf-8").strip()

    # Read and filter input
    rows = []
    total_scanned = 0
    for rec in read_jsonl(args.input):
        total_scanned += 1
        if rec.get("views", 0) >= args.min_views:
            rows.append(rec)
        if args.test_first_10 and len(rows) >= 10:
            break
        if args.limit and len(rows) >= args.limit:
            break

    # Optional sharding
    if args.num_shards and args.shard_index is not None:
        orig_len = len(rows)
        rows = [r for i, r in enumerate(rows) if i % args.num_shards == args.shard_index]
        logging.info(
            f"Applied sharding: shard_index={args.shard_index} num_shards={args.num_shards} -> {len(rows)}/{orig_len} items"
        )

    logging.info(f"Scanned {total_scanned} lines, selected {len(rows)} items")
    write_jsonl(outdir / "selected.jsonl", rows, mode="w")

    # Audio root
    audio_root = Path("/app2/suno/data/raw_audio_opus_v0")

    results_path = outdir / "results.jsonl"
    errors_path = outdir / "errors.jsonl"
    checkpoint_path = outdir / "checkpoint.json"

    success_batch: List[Dict] = []
    error_batch: List[Dict] = []

    processed = 0
    succeeded = 0
    failed = 0

    t_start = time.time()

    def _task(rec: Dict) -> Dict:
        return process_one(rec, prompt_text, audio_root)

    with futures.ThreadPoolExecutor(max_workers=args.max_workers) as ex:
        future_to_id = {ex.submit(_task, r): r.get("id") for r in rows}

        for fut in futures.as_completed(future_to_id):
            _id = future_to_id[fut]
            processed += 1
            try:
                out = fut.result()
                success_batch.append(out)
                succeeded += 1
                logging.info(f"OK id={_id} (processed={processed}, success={succeeded}, fail={failed})")
            except Exception as e:
                failed += 1
                err = {"id": _id, "error": str(e)}
                error_batch.append(err)
                logging.warning(f"ERR id={_id}: {e}")

            # Periodic save
            if len(success_batch) >= args.chunk_size:
                write_jsonl(results_path, success_batch, mode="a")
                success_batch.clear()

            if len(error_batch) >= max(5, args.chunk_size // 5):
                write_jsonl(errors_path, error_batch, mode="a")
                error_batch.clear()

            # Lightweight checkpoint for monitoring
            with checkpoint_path.open("w", encoding="utf-8") as f:
                json.dump(
                    {
                        "processed": processed,
                        "succeeded": succeeded,
                        "failed": failed,
                        "elapsed_s": round(time.time() - t_start, 2),
                    },
                    f,
                )

    # Flush remaining
    if success_batch:
        write_jsonl(results_path, success_batch, mode="a")
    if error_batch:
        write_jsonl(errors_path, error_batch, mode="a")

    elapsed = time.time() - t_start
    logging.info(
        f"Done. processed={processed} succeeded={succeeded} failed={failed} elapsed={round(elapsed, 2)}s"
    )


if __name__ == "__main__":
    main()