# Libraries
import json
import os
import re
import util.dazedwrap as dazedwrap
import threading
import time
import traceback
import tiktoken
import csv
from pathlib import Path
from colorama import Fore
from dotenv import load_dotenv
from retry import retry
from tqdm import tqdm
from util.translation import TranslationConfig, translateAI as sharedtranslateAI, getPricingConfig, calculateCost
from util.speakers import strip_speaker_prefix
import tempfile

# Globals
MODEL = os.getenv("model")
TIMEOUT = int(os.getenv("timeout"))
LANGUAGE = os.getenv("language").capitalize()
from util.paths import PROMPT_PATH, VOCAB_PATH

PROMPT = PROMPT_PATH.read_text(encoding="utf-8")
VOCAB = VOCAB_PATH.read_text(encoding="utf-8")
LOCK = threading.Lock()
WIDTH = int(os.getenv("width"))
LISTWIDTH = int(os.getenv("listWidth"))
NOTEWIDTH = int(os.getenv("noteWidth"))
MAXHISTORY = 10
ESTIMATE = ""
TOKENS = [0, 0]
NAMESLIST = []
NAMES = False  # Output a list of all the character names found
BRFLAG = False  # If the game uses <br> instead
FIXTEXTWRAP = True  # Overwrites textwrap
IGNORETLTEXT = True  # Ignores all translated text.
MISMATCH = []  # Lists files that thdata a mismatch error (Length of GPT list response is wrong)
FILENAME = None
BRACKETNAMES = False

# CSV Configuration Settings (configurable via GUI)
CSV_DELIMITER = ","  # CSV delimiter character (comma, semicolon, tab)
SOURCE_COLUMN = 0  # Which column has the source text to translate (0-indexed)
TARGET_COLUMN = 1  # Which column to write translations to
SPEAKER_COLUMN = 1  # Which column has speaker names (-1 = none)
SKIP_HEADER_ROW = True  # Skip the first row (header)
USE_TARGET_IF_NOT_EMPTY = False  # Use target column text if not empty (T++ style)
SKIP_IF_TARGET_TRANSLATED = True  # Skip batches whose targets are already translated
WRITE_TO_NEXT_COLUMN = False  # Write to column after target instead of overwriting
PARSE_NAME_TAGS = False  # Parse :name[] tags in text
PARSE_M_MARKERS = False  # Parse \M markers in text
REMOVE_FURIGANA = True  # Remove furigana annotations ＜＝＞
SKIP_COMMENT_ROWS = False  # Skip rows starting with 'comment'

# Regex - Need to change this if you want to translate from/to other languages. Default is Japanese Regex
LANGREGEX = r"[一-龠ぁ-ゔァ-ヴーａ-ｚＡ-Ｚ０-９\uFF61-\uFF9F]+"

# Get pricing configuration based on the model
PRICING_CONFIG = getPricingConfig(MODEL)
INPUTAPICOST = PRICING_CONFIG["inputAPICost"]
OUTPUTAPICOST = PRICING_CONFIG["outputAPICost"]
BATCHSIZE = PRICING_CONFIG["batchSize"]
FREQUENCY_PENALTY = PRICING_CONFIG["frequencyPenalty"]

# tqdm Globals
BAR_FORMAT = "{l_bar}{bar:10}{r_bar}{bar:-10b}"
POSITION = 0
LEAVE = False
PBAR = None
ENCODING = "utf8"

# Initialize Translation Config
TRANSLATION_CONFIG = TranslationConfig(
    model=MODEL,
    language=LANGUAGE,
    prompt=PROMPT,
    vocab=VOCAB,
    langRegex=LANGREGEX,
    batchSize=BATCHSIZE,
    maxHistory=MAXHISTORY,
    estimateMode=False  # Will be set dynamically based on ESTIMATE
)
LEAVE = False

def handleCSV(filename, estimate):
    global ESTIMATE, TOKENS, FILENAME
    ESTIMATE = estimate
    FILENAME = filename

    if not ESTIMATE:
        with open("translated/" + filename, "w+t", newline="", encoding=ENCODING, errors="xmlcharrefreplace") as writeFile:
            # Translate
            start = time.time()
            translatedData = openFiles(filename, writeFile)

            # Print Result
            end = time.time()
            tqdm.write(getResultString(translatedData, end - start, filename))
            with LOCK:
                TOKENS[0] += translatedData[1][0]
                TOKENS[1] += translatedData[1][1]
    else:
        # Translate
        start = time.time()
        translatedData = openFilesEstimate(filename)

        # Print Result
        end = time.time()
        tqdm.write(getResultString(translatedData, end - start, filename))
        with LOCK:
            TOKENS[0] += translatedData[1][0]
            TOKENS[1] += translatedData[1][1]

    # Print Total
    totalString = getResultString(["", TOKENS, None], end - start, "TOTAL")

    # Print any errors on maps
    if len(MISMATCH) > 0:
        return totalString + Fore.RED + f"\nMismatch Errors: {MISMATCH}" + Fore.RESET
    else:
        return totalString


def openFiles(filename, writeFile):
    with open("files/" + filename, "r", encoding=ENCODING) as readFile, writeFile:
        translatedData = parseCSV(readFile, writeFile, filename)

    return translatedData


def openFilesEstimate(filename):
    with open("files/" + filename, "r", encoding="utf8") as readFile:
        translatedData = parseCSV(readFile, "", filename)

    return translatedData


def getResultString(translatedData, translationTime, filename):
    # File Print String
    cost = calculateCost(translatedData[1][0], translatedData[1][1], MODEL)
    totalTokenstring = (
        Fore.YELLOW + "[Input: " + str(translatedData[1][0]) + "]"
        "[Output: "
        + str(translatedData[1][1])
        + "]" "[Cost: ${:,.4f}".format(cost)
        + "]"
    )
    timeString = Fore.BLUE + "[" + str(round(translationTime, 1)) + "s]"

    if translatedData[2] is None:
        # Success
        return filename + ": " + totalTokenstring + timeString + Fore.GREEN + " \u2713 " + Fore.RESET
    else:
        # Fail
        try:
            raise translatedData[2]
        except Exception as e:
            traceback.print_exc()
            errorString = str(e) + Fore.RED
            return filename + ": " + totalTokenstring + timeString + Fore.RED + " \u2717 " + errorString + Fore.RESET


def parseCSV(readFile, writeFile, filename):
    totalTokens = [0, 0]
    totalLines = 0
    global LOCK

    # Get total for progress bar
    totalLines = len(readFile.readlines())
    readFile.seek(0)

    reader = csv.reader(readFile, delimiter=CSV_DELIMITER)
    if not ESTIMATE:
        writer = csv.writer(
            writeFile,
            delimiter=CSV_DELIMITER,
        )
    else:
        writer = ""

    with tqdm(bar_format=BAR_FORMAT, position=POSITION, total=totalLines, leave=LEAVE) as pbar:
        pbar.desc = filename
        pbar.total = totalLines

        # Grab All Rows
        data = []
        for row in reader:
            data.append(row)

        try:
            response = translateCSV(data, pbar, writeFile, writer, filename, None)
            totalTokens[0] = response[0]
            totalTokens[1] = response[1]
        except Exception:
            traceback.print_exc()
    return [data, totalTokens, None]


def flush_progress_csv(writeFile, writer, rows):
    """Flush current CSV progress to the already-open output file (Windows-safe)."""
    try:
        if ESTIMATE or writeFile is None:
            return
        with LOCK:
            writeFile.seek(0)
            # Recreate writer at current position to avoid state issues
            tmp_writer = csv.writer(writeFile, delimiter=CSV_DELIMITER)
            tmp_writer.writerows(rows)
            writeFile.truncate()
            writeFile.flush()
            os.fsync(writeFile.fileno())
    except Exception:
        traceback.print_exc()


def _actual_target_column():
    """Column index translations are written to."""
    return TARGET_COLUMN + 1 if WRITE_TO_NEXT_COLUMN else TARGET_COLUMN


def _target_is_translated(row):
    """True when the write-target cell is non-empty and has no Japanese left."""
    actual_target = _actual_target_column()
    if len(row) <= actual_target:
        return False
    target = row[actual_target]
    if not isinstance(target, str) or not target.strip():
        return False
    return not re.search(LANGREGEX, target)


def _row_source_text(row):
    """Source text for a row, honoring USE_TARGET_IF_NOT_EMPTY."""
    if USE_TARGET_IF_NOT_EMPTY and len(row) > TARGET_COLUMN and row[TARGET_COLUMN]:
        return row[TARGET_COLUMN]
    if len(row) > SOURCE_COLUMN and row[SOURCE_COLUMN]:
        return row[SOURCE_COLUMN]
    return ""


def _is_candidate_row(data, i):
    """True when this row would normally be collected for translation."""
    if SKIP_HEADER_ROW and i == 0:
        return False
    if SKIP_COMMENT_ROWS and len(data[i]) > 0 and "comment" in str(data[i][0]).lower():
        return False
    if len(data[i]) <= SOURCE_COLUMN:
        return False
    return bool(_row_source_text(data[i]))


def _collect_process_indices(data):
    """Row indices to translate.

    When SKIP_IF_TARGET_TRANSLATED is on, candidates are checked in groups of
    BATCHSIZE (same as the translation batch size):
    - If every target in the group is already translated, skip the group.
    - If any row still needs work, keep the group but only queue the
      untranslated rows so finished targets are left alone.
    """
    candidates = [i for i in range(len(data)) if _is_candidate_row(data, i)]
    if not SKIP_IF_TARGET_TRANSLATED:
        return candidates

    process = []
    batch_size = max(1, int(BATCHSIZE))
    for start in range(0, len(candidates), batch_size):
        batch = candidates[start:start + batch_size]
        if all(_target_is_translated(data[i]) for i in batch):
            continue
        process.extend(i for i in batch if not _target_is_translated(data[i]))
    return process


def translateCSV(data, pbar, writeFile, writer, filename, translatedList):
    """
    Unified CSV translation function using configurable settings.
    
    Uses global settings:
    - SOURCE_COLUMN: column index for source text
    - TARGET_COLUMN: column index to write translations
    - SPEAKER_COLUMN: column index for speaker names (-1 = none)
    - SKIP_HEADER_ROW: whether to skip first row
    - USE_TARGET_IF_NOT_EMPTY: use existing target text if present (T++ style)
    - SKIP_IF_TARGET_TRANSLATED: skip fully-translated candidate batches
    - WRITE_TO_NEXT_COLUMN: write to column after target
    - PARSE_NAME_TAGS: parse :name[] tags
    - PARSE_M_MARKERS: parse \\M markers
    - REMOVE_FURIGANA: remove furigana
    - SKIP_COMMENT_ROWS: skip rows with 'comment' in first column
    """
    global LOCK, ESTIMATE, PBAR
    PBAR = pbar
    translatedText = ""
    totalTokens = [0, 0]
    stringList = []
    process_indices = _collect_process_indices(data)

    try:
        # Translate
        for i in process_indices:
            # Get source text
            jaString = _row_source_text(data[i])
            speaker = ""
            
            # Handle speaker column if configured
            if SPEAKER_COLUMN >= 0 and len(data[i]) > SPEAKER_COLUMN and data[i][SPEAKER_COLUMN]:
                speakerResponse = getSpeaker(data[i][SPEAKER_COLUMN])
                totalTokens[0] += speakerResponse[1][0]
                totalTokens[1] += speakerResponse[1][1]
                speaker = speakerResponse[0]
                data[i][SPEAKER_COLUMN] = speaker
            
            # Parse :name[] tags if configured
            if PARSE_NAME_TAGS and ':name' in jaString:
                match = re.search(r":name\[([^\]]+?)\]\n([\w\W]*)", jaString)
                if match:
                    # Translate speaker name
                    response = getSpeaker(match.group(1))
                    speaker = response[0]
                    totalTokens[0] += response[1][0]
                    totalTokens[1] += response[1][1]
                    data[i][SOURCE_COLUMN] = data[i][SOURCE_COLUMN].replace(match.group(1), speaker)
                    
                    # Extract text portion
                    jaString = match.group(2)
                    # Remove voice markers
                    voMatch = re.search(r"\\[vfF]+\[.+]", jaString)
                    if voMatch:
                        jaString = jaString.replace(voMatch.group(0), "")
            
            # Parse \M markers if configured
            if PARSE_M_MARKERS and '\\M' in jaString:
                match = re.search(r"\\M.+\n([\w\W]*)", jaString)
                if match:
                    jaString = match.group(1)
                    voMatch = re.search(r"\\[vfF]+\[.+]", jaString)
                    if voMatch:
                        jaString = jaString.replace(voMatch.group(0), "")
            
            # Remove furigana if configured
            if REMOVE_FURIGANA:
                jaString = re.sub(r"＜(.*)＝.*＞", r"\1", jaString)
            
            # Store original for replacement
            ojaString = jaString
            
            # Remove textwrap
            jaString = jaString.replace("\\n", " ")
            jaString = jaString.replace("\n", " ")
            
            # Pass 1: Collect strings
            if not translatedList:
                if speaker:
                    stringList.append(f"[{speaker}]: {jaString}")
                else:
                    stringList.append(jaString)
            
            # Pass 2: Apply translations
            else:
                # Grab and pop translation
                translatedText = translatedList[0]
                translatedList.pop(0)
                
                # Remove speaker prefix from translation if present
                translatedText = strip_speaker_prefix(translatedText)
                
                # Add wordwrap
                translatedText = dazedwrap.wrapText(translatedText, WIDTH)
                translatedText = translatedText.replace("\n", "\\n")
                
                # Determine target column
                actual_target = _actual_target_column()
                
                # Ensure row has enough columns
                while len(data[i]) <= actual_target:
                    data[i].append("")
                
                # Set data
                if PARSE_NAME_TAGS or PARSE_M_MARKERS:
                    # Replace original text portion
                    data[i][actual_target] = data[i][actual_target].replace(ojaString, translatedText) if data[i][actual_target] else translatedText
                else:
                    data[i][actual_target] = translatedText
                
                flush_progress_csv(writeFile, writer, data)

        # EOF - Process collected strings
        if len(stringList) > 0:
            # Set Progress
            pbar.total = len(stringList)
            pbar.refresh()

            # Translate
            response = translateAI(stringList, "")
            totalTokens[0] += response[1][0]
            totalTokens[1] += response[1][1]
            translatedList = response[0]

            # Set Strings
            if len(stringList) == len(translatedList):
                translateCSV(data, pbar, writeFile, writer, filename, translatedList)
            # Mismatch
            else:
                with LOCK:
                    if filename not in MISMATCH:
                        MISMATCH.append(filename)

            # Write all Data
            with LOCK:
                if not ESTIMATE:
                    for row in data:
                        writer.writerow(row)
            flush_progress_csv(writeFile, writer, data)

    except Exception:
        traceback.print_exc()

        # Write all Data
        with LOCK:
            if not ESTIMATE:
                for row in data:
                    writer.writerow(row)
        flush_progress_csv(writeFile, writer, data)
        return totalTokens

    return totalTokens

# Save some money and enter the character before translation
def getSpeaker(speaker):
    match speaker:
        case "ファイン":
            return ["Fine", [0, 0]]
        case "":
            return ["", [0, 0]]
        case _:
            # Find Speaker
            for i in range(len(NAMESLIST)):
                if speaker == NAMESLIST[i][0]:
                    return [NAMESLIST[i][1], [0, 0]]
                
            if speaker == "？？？":
                return ["???", [0, 0]]

            # Translate and Store Speaker
            response = translateAI(
                f"{speaker}",
                "Reply with the " + LANGUAGE + " translation of the NPC name.",
                False,
            )
            response[0] = response[0].title()
            response[0] = response[0].replace("'S", "'s")
            response[0] = response[0].replace("Speaker: ", "")

            # Retry if name doesn't translate for some reason
            if re.search(r"([a-zA-Z？?])", response[0]) == None:
                response = translateAI(
                    f"{speaker}",
                    "Reply with the " + LANGUAGE + " translation of the NPC name.",
                    False,
                )
                response[0] = response[0].title()
                response[0] = response[0].replace("'S", "'s")

            speakerList = [speaker, response[0]]
            NAMESLIST.append(speakerList)
            return response
    return [speaker, [0, 0]]

def translateAI(text, history, history_ctx=None):
    """
    Legacy wrapper function for the new shared translation utility.
    This maintains compatibility with existing code while using the new shared implementation.
    """
    global PBAR, MISMATCH, FILENAME
    
    # Update config estimate mode based on global ESTIMATE
    TRANSLATION_CONFIG.estimateMode = bool(ESTIMATE)
    
    # Call the new shared translation function
    return sharedtranslateAI(
        text=text,
        history=history,
        config=TRANSLATION_CONFIG,
        filename=FILENAME,
        pbar=PBAR,
        lock=LOCK,
        mismatchList=MISMATCH
    )