import argparse
import json
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, Iterable, List, Optional, Tuple

from tqdm.auto import tqdm

try:
    from openai import OpenAI  # type: ignore
except ImportError:
    OpenAI = None  # type: ignore


BatchItem = Dict[str, Any]
ModelResult = Dict[str, Any]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Infer original source artists and song titles for mashups using "
            "the OpenAI API in batched requests."
        )
    )
    parser.add_argument(
        "--input",
        required=False,
        help="Path to input JSONL file containing mashup metadata.",
    )
    parser.add_argument(
        "--output",
        required=False,
        help="Path to output JSONL file with enriched metadata.",
    )
    parser.add_argument(
        "--model",
        default="gpt-4.1-mini",
        help="OpenAI model name to use (default: gpt-4.1).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=50,
        help="Number of rows to send per OpenAI request (default: 10).",
    )
    parser.add_argument(
        "--max-rows",
        type=int,
        default=None,
        help="Optional maximum number of rows to process (for testing).",
    )
    parser.add_argument(
        "--rpm",
        type=float,
        default=None,
        help="Optional max requests per minute (simple rate limiting).",
    )
    parser.add_argument(
        "--max-retries",
        type=int,
        default=3,
        help="Maximum number of retries for a failed OpenAI request.",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=1,
        help="Number of concurrent OpenAI batch requests to run (default: 1).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help=(
            "Run a small in-memory example batch (including the provided "
            "Ariana Grande x Arca mashup) and print results instead of "
            "reading/writing files."
        ),
    )
    return parser.parse_args()


def get_openai_api_key() -> str:
    """
    Fetch the OpenAI API key.

    By default, this reads from the OPENAI_API_KEY environment variable.
    You can also hard-code a fallback here if you prefer, e.g.:

        api_key = os.environ.get("OPENAI_API_KEY") or "sk-..."

    For security, it's recommended to keep using the environment variable.
    """
    api_key = os.environ.get("OPENAI_API_KEY", "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")
    if not api_key:
        raise RuntimeError(
            "OPENAI_API_KEY environment variable is not set. "
            "Please export it before running this script."
        )
    return api_key


def ensure_openai_initialized(api_key: str) -> None:
    if OpenAI is None:
        raise RuntimeError(
            "The 'openai' Python package is not installed. "
            "Install it with `pip install openai`."
        )


def iter_jsonl(
    path: str, max_rows: Optional[int] = None
) -> Iterable[Tuple[int, Dict[str, Any]]]:
    with open(path, "r", encoding="utf-8") as f:
        for idx, line in enumerate(f):
            if max_rows is not None and idx >= max_rows:
                break
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError as e:
                print(
                    f"Skipping malformed JSON on line {idx + 1}: {e}",
                    file=sys.stderr,
                )
                continue
            yield idx, obj


def write_jsonl(path: str, rows: Iterable[Dict[str, Any]]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def build_batch_items(
    rows: List[Tuple[int, Dict[str, Any]]]
) -> List[BatchItem]:
    batch: List[BatchItem] = []
    for global_idx, row in rows:
        item: BatchItem = {"id": global_idx}

        # If there is no YouTube match, send a minimal but useful set of textual
        # metadata fields to help infer the component songs.
        if row.get("yt_match_score") is None:
            item.update(
                {
                    "song_name": row.get("song_name"),
                    "artists": row.get("artists"),
                    "album_name": row.get("album_name"),
                    "label": row.get("label"),
                    "genre": row.get("genre"),
                    "delimiter_found": row.get("delimiter_found"),
                    "mashup_keyword": row.get("mashup_keyword"),
                    "data_source": row.get("data_source"),
                    "split_parts": row.get("split_parts"),
                    "metadata": row.get("metadata"),
                }
            )
        # If we do have YouTube info, focus on title/artist plus key textual hints.
        else:
            item.update(
                {
                    "song_name": row.get("song_name"),
                    "artists": row.get("artists"),
                    "album_name": row.get("album_name"),
                    "label": row.get("label"),
                    "genre": row.get("genre"),
                    "delimiter_found": row.get("delimiter_found"),
                    "mashup_keyword": row.get("mashup_keyword"),
                    "data_source": row.get("data_source"),
                    "split_parts": row.get("split_parts"),
                    "yt_video_title": row.get("yt_video_title"),
                    "yt_channel_title": row.get("yt_channel_title"),
                    "yt_match_score": row.get("yt_match_score"),
                }
            )

        batch.append(item)
    return batch


def build_system_prompt() -> str:
    return (
        "You are an expert at parsing noisy YouTube mashup metadata and at "
        "recognizing songs and artists from incomplete or corrupted information. "
        "Given a batch of mashup video metadata items, you must infer the "
        "original songs and artists being mashed together.\n\n"
        "Rules:\n"
        "- For each item, identify the original source songs and their artists. "
        "Assume there are at least two distinct source songs in every mashup.\n"
        "- Most items are mashups of exactly 2 songs; prefer 2 source songs when "
        "that is most plausible, but allow more if clearly indicated.\n"
        "- Some items include YouTube metadata (yt_* fields); others do not and "
        "only have fields like song_name, split_parts, album_name, etc. Use "
        "whatever fields are present for each item.\n"
        "- Use your knowledge of real-world music (including underground, indie, "
        "and DJ/producer scenes) to infer likely source songs and artists even when "
        "titles or names are missing, truncated, translated, or misspelled. "
        "Correct obvious typos and normalize aliases and stylizations to canonical "
        "artist and track names when possible.\n"
        "- Prefer artist–song combinations that you know (or strongly expect) "
        "actually exist over unlikely or inconsistent ones, and avoid inventing "
        "obviously fake artists or titles. When multiple songs could match, choose "
        "the most musically plausible option given genre, era, and context, and "
        "briefly note your reasoning in parsing_notes.\n"
        "- Do not automatically assume that the mashup artist is one of the source "
        "artists; mashups are often created by third-party producers or DJs. "
        "However, in the case where the the mashup artist (indicated by the artists field)"
        "performs one of the source songs, be sure to label them as a source artist.\n"
        "- Always return source_artists and source_titles as arrays of the same "
        "length. When you know a title but not the artist (or vice versa), fill "
        "the unknown entry with the string \"unknown\". Do not copy the exact same "
        "string into both source_artists and source_titles; if you cannot tell "
        "whether a token is an artist or a title, choose the most likely role, "
        "use \"unknown\" for the other, and reflect the uncertainty in "
        "parsing_confidence and parsing_notes. If you cannot confidently identify "
        "two full song–artist pairs, still provide your best guesses for at least "
        "two sources (using \"unknown\" where needed) and set parsing_confidence "
        "to a low value.\n"
        "- Respond ONLY with valid JSON. Do not include natural-language text, "
        "backticks, comments, or any content outside the JSON.\n"
        "- The entire response MUST be a single JSON array with one object per "
        "input item, matching the specified schema exactly.\n"
        "- Each object must have keys: id, source_artists, source_titles, "
        "parsing_confidence, parsing_notes.\n"
        "- id must exactly match the id field of the corresponding input item.\n"
        "- parsing_confidence is a float between 0 and 1 representing your overall "
        "confidence in the correctness of the inferred sources.\n"
        "- source_artists and source_titles should usually each have length 2.\n"
    )


def build_user_prompt(batch_items: List[BatchItem]) -> str:
    return json.dumps(
        {
            "description": (
                "Input is a list of mashup metadata items. For each item, infer "
                "the original source songs and artists being mashed up. "
                "Return a JSON array of the same length as the input list, with "
                "one object per item following the specified schema."
            ),
            "items": batch_items,
            "expected_output_schema": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "id": {"type": "integer"},
                        "source_artists": {
                            "type": "array",
                            "items": {"type": "string"},
                        },
                        "source_titles": {
                            "type": "array",
                            "items": {"type": "string"},
                        },
                        "parsing_confidence": {
                            "type": "number",
                            "description": (
                                "Float in [0, 1] representing overall confidence "
                                "in the correctness of the inferred sources."
                            ),
                        },
                        "parsing_notes": {
                            "type": "string",
                            "description": (
                                "Free-text explanation of how the sources were "
                                "inferred, including any uncertainties or "
                                "assumptions."
                            ),
                        },
                    },
                    "required": [
                        "id",
                        "source_artists",
                        "source_titles",
                        "parsing_confidence",
                        "parsing_notes",
                    ],
                },
            },
        },
        ensure_ascii=False,
    )


def call_openai_for_batch(
    batch_items: List[BatchItem],
    model: str,
    client: Any,
    max_retries: int = 3,
    rpm: Optional[float] = None,
    last_request_ts: Optional[float] = None,
) -> Tuple[List[ModelResult], float]:
    if not batch_items:
        return [], last_request_ts or time.time()

    system_prompt = build_system_prompt()
    user_prompt = build_user_prompt(batch_items)

    attempt = 0
    while True:
        attempt += 1

        # Simple rate limiting based on requests-per-minute.
        if rpm is not None and last_request_ts is not None:
            min_interval = 60.0 / max(rpm, 1e-6)
            elapsed = time.time() - last_request_ts
            if elapsed < min_interval:
                time.sleep(min_interval - elapsed)

        try:
            completion = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=0,
            )
            content = completion.choices[0].message.content or ""
            try:
                parsed = json.loads(content)
            except json.JSONDecodeError:
                # Sometimes the model may include extra text around the JSON.
                # Try to recover by extracting the first JSON array substring.
                stripped = content.strip()
                start = stripped.find("[")
                end = stripped.rfind("]")
                if start != -1 and end != -1 and end > start:
                    candidate = stripped[start : end + 1]
                    parsed = json.loads(candidate)
                else:
                    raise
            if not isinstance(parsed, list):
                raise ValueError("Model response is not a JSON array.")
            results: List[ModelResult] = []
            for item in parsed:
                if not isinstance(item, dict):
                    continue
                results.append(item)
            return results, time.time()
        except Exception as e:  # noqa: BLE001
            print(
                f"OpenAI call failed on attempt {attempt}/{max_retries}: {e}",
                file=sys.stderr,
            )
            if attempt >= max_retries:
                # Fallback: generate default low-confidence results for each item.
                fallback_results: List[ModelResult] = []
                for bi in batch_items:
                    fallback_results.append(
                        {
                            "id": bi.get("id"),
                            "source_artists": [],
                            "source_titles": [],
                            "parsing_confidence": 0.0,
                            "parsing_notes": (
                                "Failed to parse via OpenAI API after retries; "
                                "no structured data extracted."
                            ),
                        }
                    )
                return fallback_results, time.time()
            # Exponential backoff before retrying.
            sleep_secs = min(2**attempt, 30)
            time.sleep(sleep_secs)


def merge_results_into_rows(
    batch_rows: List[Tuple[int, Dict[str, Any]]],
    results: List[ModelResult],
) -> List[Dict[str, Any]]:
    result_by_id: Dict[Any, ModelResult] = {r.get("id"): r for r in results}
    enriched_rows: List[Dict[str, Any]] = []
    for global_idx, row in batch_rows:
        r = result_by_id.get(global_idx)
        if r is None:
            # Default when model did not return anything for this id.
            row["source_artists"] = []
            row["source_titles"] = []
            row["parsing_confidence"] = 0.0
            row["parsing_notes"] = "No result returned for this row id."
        else:
            row["source_artists"] = r.get("source_artists") or []
            row["source_titles"] = r.get("source_titles") or []
            # Ensure the confidence is a float in [0, 1] if possible.
            conf = r.get("parsing_confidence")
            try:
                conf_f = float(conf)
            except (TypeError, ValueError):
                conf_f = 0.0
            if conf_f < 0.0:
                conf_f = 0.0
            if conf_f > 1.0:
                conf_f = 1.0

            notes = (r.get("parsing_notes") or "").strip()

            # If the model has clearly confused artists with titles by making the
            # arrays effectively identical, treat this as very low confidence.
            artists_norm = [str(a).strip().lower() for a in row["source_artists"]]
            titles_norm = [str(t).strip().lower() for t in row["source_titles"]]
            if artists_norm and artists_norm == titles_norm:
                if conf_f > 0.1:
                    conf_f = 0.1
                extra_note = (
                    "Model returned nearly identical values for source_artists and "
                    "source_titles; this likely reflects confusion between artist "
                    "and title roles, so confidence has been downgraded."
                )
                notes = f"{notes} {extra_note}".strip() if notes else extra_note

            row["parsing_confidence"] = conf_f
            row["parsing_notes"] = notes
        enriched_rows.append(row)
    return enriched_rows


def process_file(
    input_path: str,
    output_path: str,
    model: str,
    batch_size: int,
    max_rows: Optional[int],
    rpm: Optional[float],
    max_retries: int,
    num_workers: int,
) -> None:
    api_key = get_openai_api_key()
    ensure_openai_initialized(api_key)
    client = OpenAI(api_key=api_key)

    # If only one worker is requested, use the existing sequential batching logic.
    if num_workers <= 1:
        last_request_ts: Optional[float] = None
        batch_rows: List[Tuple[int, Dict[str, Any]]] = []
        enriched_rows: List[Dict[str, Any]] = []
        total_rows = 0
        openai_rows = 0
        num_batches = 0

        for global_idx, row in tqdm(
            iter_jsonl(input_path, max_rows=max_rows),
            desc="Processing rows",
            unit="row",
        ):
            total_rows += 1
            batch_rows.append((global_idx, row))
            openai_rows += 1
            if len(batch_rows) >= batch_size:
                num_batches += 1
                batch_items = build_batch_items(batch_rows)
                results, last_request_ts = call_openai_for_batch(
                    batch_items=batch_items,
                    model=model,
                    client=client,
                    max_retries=max_retries,
                    rpm=rpm,
                    last_request_ts=last_request_ts,
                )
                enriched_rows.extend(merge_results_into_rows(batch_rows, results))
                batch_rows = []

        # Process any remaining rows.
        if batch_rows:
            num_batches += 1
            batch_items = build_batch_items(batch_rows)
            results, last_request_ts = call_openai_for_batch(
                batch_items=batch_items,
                model=model,
                client=client,
                max_retries=max_retries,
                rpm=rpm,
                last_request_ts=last_request_ts,
            )
            enriched_rows.extend(merge_results_into_rows(batch_rows, results))

        write_jsonl(output_path, enriched_rows)
        print(
            "Finished processing.\n"
            f"  Total rows read: {total_rows}\n"
            f"  Rows sent to OpenAI: {openai_rows}\n"
            f"  OpenAI batches: {num_batches}\n"
            f"  Output written to: {output_path}",
            file=sys.stderr,
        )
        return

    # Parallel mode: use a thread pool to run multiple batch requests concurrently.
    last_request_ts = None  # rpm is not used in parallel mode
    batch_rows = []
    enriched_by_idx: Dict[int, Dict[str, Any]] = {}
    total_rows = 0
    openai_rows = 0
    num_batches = 0
    futures: List[Any] = []

    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        # First pass: read rows and submit OpenAI batch requests.
        for global_idx, row in iter_jsonl(input_path, max_rows=max_rows):
            total_rows += 1
            batch_rows.append((global_idx, row))
            openai_rows += 1
            if len(batch_rows) >= batch_size:
                num_batches += 1
                local_batch = batch_rows
                batch_rows = []
                batch_items = build_batch_items(local_batch)
                future = executor.submit(
                    call_openai_for_batch,
                    batch_items=batch_items,
                    model=model,
                    client=client,
                    max_retries=max_retries,
                    rpm=None,  # rpm limiting disabled in parallel mode
                    last_request_ts=None,
                )
                futures.append((future, local_batch))

        # Process any remaining rows.
        if batch_rows:
            num_batches += 1
            local_batch = batch_rows
            batch_rows = []
            batch_items = build_batch_items(local_batch)
            future = executor.submit(
                call_openai_for_batch,
                batch_items=batch_items,
                model=model,
                client=client,
                max_retries=max_retries,
                rpm=None,
                last_request_ts=None,
            )
            futures.append((future, local_batch))

        # Collect results from all futures, tracking progress as batches complete.
        with tqdm(
            total=total_rows,
            desc=f"Processing OpenAI batches (workers={num_workers})",
            unit="row",
        ) as pbar:
            for future, local_batch in futures:
                results, _ = future.result()
                merged_rows = merge_results_into_rows(local_batch, results)
                for (idx, _), enriched in zip(local_batch, merged_rows):
                    enriched_by_idx[idx] = enriched
                pbar.update(len(local_batch))

    # Write rows ordered by original index.
    ordered_indices = sorted(enriched_by_idx.keys())
    enriched_rows_out = [enriched_by_idx[i] for i in ordered_indices]
    write_jsonl(output_path, enriched_rows_out)
    print(
        "Finished processing.\n"
        f"  Total rows read: {total_rows}\n"
        f"  Rows sent to OpenAI: {openai_rows}\n"
        f"  OpenAI batches: {num_batches}\n"
        f"  OpenAI workers: {num_workers}\n"
        f"  Output written to: {output_path}",
        file=sys.stderr,
    )


def run_dry_run(model: str, batch_size: int, rpm: Optional[float], max_retries: int) -> None:
    """
    Run a small synthetic example batch including the Ariana Grande x Arca mashup.
    """
    api_key = get_openai_api_key()
    ensure_openai_initialized(api_key)
    client = OpenAI(api_key=api_key)

    example_rows: List[Tuple[int, Dict[str, Any]]] = []

    # Example based on the user's description.
    example_rows.append(
        (
            0,
            {
                "song_name": "Ariana Grande x Arca - One Last Time / Urchin (Mashup)",
                "artists": ["Jane Remover"],
                "album_name": None,
                "genre": (
                    "deconstructed club, repetitive, pop rap, anthemic, playful, "
                    "medley, boastful, lo-fi, mashup, dariacore, slacker rock, male "
                    "vocalist, quirky, sampling, rhythmic, female vocalist, party, "
                    "electro house, energetic, lgbt, mashup, androgynous vocals, "
                    "eclectic, electronic dance music, jersey club, dance-popbubblegum "
                    "bass, slowcore, complex, futuristic, happy, summer"
                ),
                "key": None,
                "bpm": None,
                "label": None,
                "release_date": None,
                "delimiter_found": None,
                "mashup_keyword": "mashup",
                "data_source": "dry_run_example",
                "split_parts": None,
                "metadata": None,
                "mashup_strength": None,
                "yt_video_title": (
                    "jane remover - ariana grande x arca - one last time / urchin (mashup)"
                ),
                "yt_channel_title": "peri music channel :0",
                "yt_match_score": 100.0,
                "yt_duration": None,
                "yt_url": "https://www.youtube.com/watch?v=wlTUDoNQtN4",
                "yt_views": None,
            },
        )
    )

    batch_items = build_batch_items(example_rows)
    results, _ = call_openai_for_batch(
        batch_items=batch_items,
        model=model,
        client=client,
        max_retries=max_retries,
        rpm=rpm,
        last_request_ts=None,
    )
    enriched = merge_results_into_rows(example_rows, results)
    print(json.dumps(enriched, ensure_ascii=False, indent=2))


def main() -> None:
    args = parse_args()

    if args.dry_run:
        if args.input or args.output:
            print(
                "Warning: --dry-run ignores --input/--output and uses in-memory examples only.",
                file=sys.stderr,
            )
        run_dry_run(
            model=args.model,
            batch_size=args.batch_size,
            rpm=args.rpm,
            max_retries=args.max_retries,
        )
        return

    if not args.input or not args.output:
        print(
            "Error: --input and --output are required unless --dry-run is specified.",
            file=sys.stderr,
        )
        sys.exit(1)

    process_file(
        input_path=args.input,
        output_path=args.output,
        model=args.model,
        batch_size=args.batch_size,
        max_rows=args.max_rows,
        rpm=args.rpm,
        max_retries=args.max_retries,
        num_workers=args.num_workers,
    )


if __name__ == "__main__":
    main()