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 track (artists and title) for remixes using "
            "the OpenAI API in batched requests."
        )
    )
    parser.add_argument(
        "--input",
        required=False,
        help="Path to input JSONL file containing remix 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-mini).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=50,
        help="Number of rows to send per OpenAI request (default: 50).",
    )
    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 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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",  # noqa: E501
    )
    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]:
    """
    Build the list of batch items sent to the model for remix parsing.

    Each item contains a subset of metadata fields that help identify the
    original source track for the remix.
    """
    batch: List[BatchItem] = []
    for global_idx, row in rows:
        item: BatchItem = {"id": global_idx}

        # Core textual metadata.
        item.update(
            {
                "title": row.get("title"),
                "artists": row.get("artists"),
                "album_name": row.get("album_name"),
                "label": row.get("label"),
                "genre": row.get("genre"),
                "bpm": row.get("bpm"),
                "key": row.get("key"),
                "duration": row.get("duration"),
                "release_date": row.get("release_date"),
                "data_source": row.get("data_source"),
                "existing_source_artists": row.get("source_artists"),
                "existing_source_title": row.get("source_title"),
                "remix_confidence": row.get("remix_confidence"),
                "output_id": row.get("output_id"),
                "source_id": row.get("source_id"),
                "source_votes": row.get("source_votes"),
                "s3_filepath": row.get("s3_filepath"),
            }
        )

        batch.append(item)
    return batch


def build_system_prompt() -> str:
    return (
        "You are an expert at parsing noisy remix metadata and at recognizing"
        " songs and artists from incomplete or corrupted information. Given a"
        " batch of remix metadata items, you must infer the original source"
        " song that is being remixed.\n\n"
        "Rules:\n"
        "- For each item, identify the single original source song and its"
        " artist or artists. Most items correspond to a remix of exactly one"
        " track.\n"
        "- Use whatever metadata is present, including title, artists, album"
        " name, label, genre, BPM, key, duration, release date, and data"
        " source.\n"
        "- Use your knowledge of real-world music (including underground,"
        " indie, DJ/producer scenes, and remix culture) to infer the most"
        " likely original track 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 and historically plausible"
        " option given genre, era, and context, and briefly note your"
        " reasoning in parsing_notes.\n"
        "- Do not automatically assume that the listed remix artist is the"
        " original source artist; many remixes are created by third-party"
        " producers or DJs.\n"
        "- In many cases, the original artist appears in the `artists` field"
        " alongside the remixer (for example: 'Original Artist - Track Name"
        " (Remixer Name Remix)'). When the formatting clearly separates"
        " original vs remixer, treat the original performer as the"
        " source_artist and the remixer as a non-source contributor.\n"
        "- If the listed artist is clearly the performer of the original"
        " track (for example, a self-remix or when only one artist is"
        " mentioned), treat them as a source artist.\n"
        "- Always return source_artists as an array of one or more strings and"
        " source_title as a single string. If you know a title but not the"
        " artist (or vice versa), fill the unknown entry with the string"
        ' \"unknown\" and explain the uncertainty in parsing_notes.\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.\n"
        "- Each object must have keys: id, source_artists, source_title,"
        " 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 source track.\n"
    )


def build_user_prompt(batch_items: List[BatchItem]) -> str:
    return json.dumps(
        {
            "description": (
                "Input is a list of remix metadata items. For each item, infer "
                "the original source song (title and artist(s)) being remixed. "
                "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_title": {"type": "string"},
                        "parsing_confidence": {
                            "type": "number",
                            "description": (
                                "Float in [0, 1] representing overall "
                                "confidence in the correctness of the "
                                "inferred source track."
                            ),
                        },
                        "parsing_notes": {
                            "type": "string",
                            "description": (
                                "Free-text explanation of how the source was "
                                "inferred, including any uncertainties or "
                                "assumptions."
                            ),
                        },
                    },
                    "required": [
                        "id",
                        "source_artists",
                        "source_title",
                        "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_title": "",
                            "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["parsing_confidence"] = 0.0
            row["parsing_notes"] = "No result returned for this row id."
        else:
            source_artists = r.get("source_artists") or []
            if not isinstance(source_artists, list):
                source_artists = [str(source_artists)]
            row["source_artists"] = [str(a) for a in source_artists]

            source_title = r.get("source_title")
            if source_title is None:
                source_title = ""
            row["source_title"] = str(source_title)

            # 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()

            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 a sequential batching logic.
    if num_workers <= 1:
        last_request_ts: Optional[float] = None
        batch_rows: List[Tuple[int, Dict[str, Any]]] = []
        enriched_by_idx: Dict[int, 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
            # Use sources_parsed and data_source as gates for OpenAI calls.
            if row.get("sources_parsed"):
                enriched_by_idx[global_idx] = row
                continue

            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,
                )
                merged_rows = merge_results_into_rows(batch_rows, results)
                for (idx, _), enriched in zip(batch_rows, merged_rows):
                    enriched_by_idx[idx] = enriched
                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,
            )
            merged_rows = merge_results_into_rows(batch_rows, results)
            for (idx, _), enriched in zip(batch_rows, merged_rows):
                enriched_by_idx[idx] = enriched

        # 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"  Output written to: {output_path}",
            file=sys.stderr,
        )
        return

    # Parallel mode: use a thread pool to run multiple batch requests concurrently.
    batch_rows: List[Tuple[int, Dict[str, Any]]] = []
    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

            if row.get("sources_parsed"):
                enriched_by_idx[global_idx] = row
                continue

            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=openai_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 for remix parsing.
    """
    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 remix row.
    example_rows.append(
        (
            0,
            {
                "title": "Artist X - Song Y (Artist Z Remix)",
                "artists": ["Artist Z"],
                "album_name": None,
                "genre": "electronic, house, remix",
                "key": None,
                "bpm": 124,
                "label": "Example Label",
                "release_date": "2020-01-01",
                "duration": 360,
                "data_source": "dry_run_example_remix",
                "source_artists": None,
                "source_title": None,
                "remix_confidence": None,
                "sources_parsed": False,
                "output_id": None,
                "source_id": None,
                "source_votes": None,
                "id": "example-0",
                "s3_filepath": 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()