"""
Infinite Anthologies: Early Science Fiction Corpus Scanner
----------------------------------------------------------
A local Python script for screening plain-text fiction files with the OpenAI API.

IMPORTANT:
- Use your own OpenAI API key.
- Keep your key private and do not publish a real key in this file.
- The default limit is 100 files. Change MAX_FILES only after testing costs and rate limits.
"""

# Import standard Python libraries.
import csv
import json
import os
import time
from pathlib import Path

# Import the OpenAI client and error classes.
from openai import OpenAI, RateLimitError, APIConnectionError, APIStatusError

# Import tiktoken so the script can take the first N model tokens rather than N words.
import tiktoken


# ---------------------------------------------------------------------------
# USER SETTINGS
# ---------------------------------------------------------------------------

# Option A: for a simple local test, replace the placeholder below with your own key.
# Do NOT publish or share the file after inserting a real key.
OPENAI_API_KEY = "your API key"

# Option B (recommended): set an OPENAI_API_KEY environment variable instead.
# If the environment variable exists, it will be used before the placeholder above.
API_KEY = os.getenv("OPENAI_API_KEY") or OPENAI_API_KEY

# Change these paths to match your computer.
INPUT_FOLDER = Path(r"C:\path\to\your\text_files")
OUTPUT_CSV = Path(r"C:\path\to\your\analysis_results.csv")

# GPT-4o is retained here to reproduce the methodology developed for this project.
MODEL = "gpt-4o"

# Analyze only the first 1,000 GPT tokens from each text.
MAX_TEXT_TOKENS_PER_FILE = 1000

# A public/research-safe default. Set to None only after testing your budget and limits.
MAX_FILES = 100

# Wait between successful requests. Increase this if your project has a low TPM limit.
BASE_REQUEST_DELAY_SECONDS = 4.0

# Retry transient failures rather than silently losing files.
MAX_RETRIES = 8
INITIAL_RETRY_DELAY_SECONDS = 2.0
MAX_RETRY_DELAY_SECONDS = 90.0


# ---------------------------------------------------------------------------
# ANALYTICAL DEFINITIONS
# ---------------------------------------------------------------------------

# These instructions operationalise the project's concept of early science fiction.
SYSTEM_INSTRUCTIONS = """
You are classifying nineteenth- and early-twentieth-century popular fiction,
including proto-science-fiction and early science fiction.

Science fiction must have a meaningful scientific, technological, rationalised,
or extrapolative basis. Count inventions, scientific experiments, scientific
discussion, future technologies, extrapolated science, scientifically framed
social futures, planetary travel, and other phenomena explicitly grounded in
science or technology.

Do not force a science-fiction interpretation onto a text. Supernatural events,
ghosts, spiritualism, magic, dreams, prophecy, visions, psychological disturbance,
coincidence, Gothic atmosphere, or fantasy do not count as science fiction unless
the text itself gives them a scientific or technological explanation.

Use the language and genre conventions of nineteenth- and early-twentieth-century
popular fiction rather than assuming contemporary genre boundaries.

Science-fiction rating guide:
0 = no detected science-fiction elements.
1-2 = scientific material is incidental, realistic, or too slight to make the story SF.
3-4 = minor but genuine speculative-scientific content.
5-6 = substantial scientific speculation mixed with other dominant genres.
7-8 = strong science-fiction content central to the narrative.
9-10 = science, technology, or scientific extrapolation dominates the story.

SF level:
none = ratings 0-2
mild = ratings 3-4
moderate = ratings 5-6
high = ratings 7-10

OCR quality:
A = clearly readable with negligible OCR corruption.
B = good/readable with minor OCR errors.
C = reasonable OCR; recurring errors but the story remains understandable.
D = poor OCR; frequent corruption materially obstructs reading.
E = effectively unreadable.

For popular-fiction genre classification, begin with these historically useful
categories:
Gothic, Sensation, Social Problems, Didactic, Adventure, Utopian, Dystopian,
Colonial, Imperialist, Domestic, Spiritualist, War, New Woman, Allegorical,
Historical, Horror, Supernatural, Crime, Detective, Invention, Planetary,
Lost Race, Romance, Travel, Prediction.

Prefer one obvious genre. Supply a second or third only when the sampled text
genuinely indicates a hybrid. A different period-appropriate popular-fiction
genre may be used when the supplied list is clearly inadequate. Do not invent
genre evidence that is absent from the supplied text.
""".strip()


# Define the exact machine-readable output expected from the model.
OUTPUT_SCHEMA = {
    "type": "json_schema",
    "name": "early_fiction_analysis",
    "strict": True,
    "schema": {
        "type": "object",
        "properties": {
            "sf_rating": {
                "type": "integer",
                "minimum": 0,
                "maximum": 10
            },
            "sf_level": {
                "type": "string",
                "enum": ["none", "mild", "moderate", "high"]
            },
            "plot_description": {
                "type": "string"
            },
            "ocr_quality": {
                "type": "string",
                "enum": ["A", "B", "C", "D", "E"]
            },
            "genre_1": {"type": "string"},
            "genre_2": {"type": "string"},
            "genre_3": {"type": "string"}
        },
        "required": [
            "sf_rating",
            "sf_level",
            "plot_description",
            "ocr_quality",
            "genre_1",
            "genre_2",
            "genre_3"
        ],
        "additionalProperties": False
    }
}


# ---------------------------------------------------------------------------
# HELPER FUNCTIONS
# ---------------------------------------------------------------------------

# Create a tokenizer suitable for GPT-4o.
try:
    ENCODING = tiktoken.encoding_for_model(MODEL)
except KeyError:
    ENCODING = tiktoken.get_encoding("o200k_base")


# Read text while tolerating common historical/OCR encodings.
def read_text_file(file_path):
    """Read a text file without allowing one encoding problem to stop the batch."""
    raw = file_path.read_bytes()

    # Try common encodings in a sensible order.
    for encoding_name in ("utf-8-sig", "utf-8", "cp1252", "latin-1"):
        try:
            return raw.decode(encoding_name)
        except UnicodeDecodeError:
            continue

    # Final fallback replaces undecodable bytes rather than crashing.
    return raw.decode("utf-8", errors="replace")


# Keep only the first N GPT tokens from a story.
def first_n_tokens(text, max_tokens):
    """Return only the beginning of the text up to the configured token limit."""
    token_ids = ENCODING.encode(text)
    shortened_ids = token_ids[:max_tokens]
    return ENCODING.decode(shortened_ids), len(shortened_ids)


# Build the story-specific instruction.
def build_user_prompt(file_name, text_sample):
    """Construct the per-file prompt while keeping the analytical method constant."""
    return f"""
File name: {file_name}

Analyze only the supplied sample below. Do not assume later plot developments that
are not evidenced in this sample.

Return:
- sf_rating using the 0-10 guide.
- sf_level using none/mild/moderate/high.
- plot_description as a factual description of the general plot in no more than
  30 words. Do not begin with phrases such as "The text contains" or "The story features".
- ocr_quality as A-E.
- genre_1 as the strongest genre identification.
- genre_2 and genre_3 only if a genuine hybrid is evidenced; otherwise return an empty string.

TEXT SAMPLE:
{text_sample}
""".strip()


# Call OpenAI with retry logic for rate limits, connection faults, and 5xx server errors.
def analyze_with_retries(client, file_name, text_sample):
    """Analyze one file and retry transient API failures instead of silently skipping it."""
    retry_delay = INITIAL_RETRY_DELAY_SECONDS

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = client.responses.create(
                model=MODEL,
                instructions=SYSTEM_INSTRUCTIONS,
                input=build_user_prompt(file_name, text_sample),
                max_output_tokens=220,
                temperature=0.1,
                text={"format": OUTPUT_SCHEMA}
            )

            # Structured Outputs should make this valid JSON.
            result = json.loads(response.output_text)

            # Return both the result and API-reported usage for auditing.
            input_tokens = getattr(response.usage, "input_tokens", "")
            output_tokens = getattr(response.usage, "output_tokens", "")
            return result, input_tokens, output_tokens, ""

        except RateLimitError as error:
            print(
                f"Rate limit for {file_name}. "
                f"Attempt {attempt}/{MAX_RETRIES}; retrying in {retry_delay:.1f}s."
            )

        except APIConnectionError as error:
            print(
                f"Connection error for {file_name}. "
                f"Attempt {attempt}/{MAX_RETRIES}; retrying in {retry_delay:.1f}s."
            )

        except APIStatusError as error:
            # Retry transient server-side failures such as 500, 502, 503, and 504.
            if error.status_code in (500, 502, 503, 504):
                print(
                    f"OpenAI server error {error.status_code} for {file_name}. "
                    f"Attempt {attempt}/{MAX_RETRIES}; retrying in {retry_delay:.1f}s."
                )
            else:
                return None, "", "", f"OpenAI API status {error.status_code}: {error}"

        except (json.JSONDecodeError, KeyError, TypeError, ValueError) as error:
            # A malformed response is worth retrying because it may be transient.
            print(
                f"Response parsing problem for {file_name}. "
                f"Attempt {attempt}/{MAX_RETRIES}; retrying in {retry_delay:.1f}s."
            )

        except Exception as error:
            # Record unexpected errors rather than terminating the complete corpus.
            return None, "", "", f"Unexpected error: {error}"

        # Wait before another attempt and then increase the wait.
        time.sleep(retry_delay)
        retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY_SECONDS)

    return None, "", "", f"Failed after {MAX_RETRIES} attempts"


# Write one CSV row and flush it immediately.
def write_row(csv_file, writer, row):
    """Persist each result immediately so an interrupted run retains completed work."""
    writer.writerow(row)
    csv_file.flush()


# ---------------------------------------------------------------------------
# MAIN PROGRAM
# ---------------------------------------------------------------------------

def main():
    """Scan the configured folder, analyze text files, and save one CSV row per file."""

    # Refuse to run with the untouched placeholder.
    if not API_KEY or API_KEY == "your API key":
        raise ValueError(
            "Add your OpenAI API key locally or set the OPENAI_API_KEY environment variable."
        )

    # Confirm the corpus folder exists before spending API credits.
    if not INPUT_FOLDER.exists():
        raise FileNotFoundError(f"Input folder does not exist: {INPUT_FOLDER}")

    # Create an OpenAI client. Manual retry logic is used below.
    client = OpenAI(api_key=API_KEY, max_retries=0)

    # Sort files for reproducible processing order.
    text_files = sorted(
        file_path
        for file_path in INPUT_FOLDER.iterdir()
        if file_path.is_file() and file_path.suffix.lower() == ".txt"
    )

    # Apply the configured public/default file cap.
    if MAX_FILES is not None:
        text_files = text_files[:MAX_FILES]

    # Create the output folder if necessary.
    OUTPUT_CSV.parent.mkdir(parents=True, exist_ok=True)

    # Open the CSV once and flush each row after it is written.
    with OUTPUT_CSV.open("w", newline="", encoding="utf-8-sig") as csv_file:
        writer = csv.writer(csv_file)
        writer.writerow([
            "File Name",
            "Science Fiction Rating",
            "SF Level",
            "Plot Description",
            "OCR Quality",
            "Genre 1",
            "Genre 2",
            "Genre 3",
            "Sample Tokens",
            "API Input Tokens",
            "API Output Tokens",
            "Status",
            "Error"
        ])
        csv_file.flush()

        # Analyze files sequentially to make rate limiting easier to manage.
        for file_path in text_files:
            file_name = file_path.name

            try:
                content = read_text_file(file_path)

                # Empty or near-empty files still receive a CSV row.
                if not content.strip():
                    write_row(csv_file, writer, [
                        file_name, "", "", "", "", "", "", "",
                        0, "", "", "skipped_empty", ""
                    ])
                    print(f"Skipped empty file: {file_name}")
                    continue

                text_sample, sample_tokens = first_n_tokens(
                    content,
                    MAX_TEXT_TOKENS_PER_FILE
                )

                result, api_input_tokens, api_output_tokens, error_message = (
                    analyze_with_retries(client, file_name, text_sample)
                )

                # Successful API result.
                if result is not None:
                    write_row(csv_file, writer, [
                        file_name,
                        result["sf_rating"],
                        result["sf_level"],
                        result["plot_description"],
                        result["ocr_quality"],
                        result["genre_1"],
                        result["genre_2"],
                        result["genre_3"],
                        sample_tokens,
                        api_input_tokens,
                        api_output_tokens,
                        "complete",
                        ""
                    ])

                # Failed file after retries still receives a row for later review.
                else:
                    write_row(csv_file, writer, [
                        file_name, "", "", "", "", "", "", "",
                        sample_tokens, api_input_tokens, api_output_tokens,
                        "error", error_message
                    ])
                    print(f"Recorded error for: {file_name}")

                # A steady delay reduces the likelihood of token-per-minute rate limits.
                time.sleep(BASE_REQUEST_DELAY_SECONDS)

            except Exception as error:
                # No problematic filename or file content should terminate the corpus.
                write_row(csv_file, writer, [
                    file_name, "", "", "", "", "", "", "",
                    "", "", "", "file_error", str(error)
                ])
                print(f"File error recorded for {file_name}: {error}")

    print(f"Finished. Results saved to: {OUTPUT_CSV}")


# Run the program only when this file is executed directly.
if __name__ == "__main__":
    main()
