import re
import os
import json
import json_repair
import sys
import argparse
from preprocessor import preprocess_remove_line_break

# Local Qwen is incapable of handling this, only use claude

project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)

from chatbot_api import llm_factory

class ChapterOverlapError(Exception):
    pass
class QuotesNotFoundError(Exception):
    pass
class ChapterPushOnError(Exception):
    pass


chapter_divide_prompt = """
Divide the following text into chapters.
Output must be in JSON format: a object/dictionary.

Requirements:

Every key should be a non-repetitive, abstract key that can be used to indicate the chapter, it should be a single word composed by capital letters and numbers. It should contain the chapter number and a abbrieviate chapter title if applicable.

Each value is a dictionary with the following keys:
"start": A short, unique, consecutive substring marks the start of the chapter, where the first character of the substring is also the first character of the chapter's main body, excluding the chapter title. Example: "Our journey begins with"
"end": A short, unique, consecutive substring marks the end of the chapter, where the last character of the substring is also the last character of the chapter. Example: "and we all went home."
These substring needs to contain the exact same characters as in the original text it quotes from, including spaces, linebreaks and punctuations etc. So it can be used to locate the chapter in the original text. Non of them should be longer than a sentence. The regions covered between start and end of different chapters should never overlap nor repeat.

Any extra content that is not a part of the story should be ignored, like title, table of contents, appendix, etc.

Place all the chapters in the original order.

Since context is limited, the provided text might only be a chunk of the whole novel.

If text of a provided chapter is explicitly completed, by conditions like the existence of the chapter title, or the end of the entire novel, you may fill in the "end" value with the substring that marks the end of the chapter. If there is no explicit ending, even if it looks like an end of chapter, you should leave the "end" value as an empty string, so it can be continued later.

A list of previously extracted chapters is provided to help you continue the extraction. If the lastest provided chapter is not complete, but is now completed in the new chunk, you should include this entry at the beginning of your response and finish it before starting a new chapter, by the following chapter content you obtained from the new chunk that belongs to the same chapter.

You should only provide the chapter provided in the text chunk. Do not provide the previously completed chapters.

The provided text may or may not include a table of contents. If it does, you should use it to help you divide the text into chapters. If it does not, you should use the text itself to divide it into chapters.

{table_of_content}

Previous Chapters:
{existing_chapters}

New Content:
{new_content}
"""

parser = argparse.ArgumentParser()
parser.add_argument("--source-file-path", "-s", type=str, default=r"novels/海底两万里/Twenty Thousand Leagues Under the Seas.txt")
args = parser.parse_args()

source_file_path = args.source_file_path
if os.path.isabs(source_file_path):
    source_file_path = os.path.abspath(source_file_path)
else:
    source_file_path = os.path.join(project_root, source_file_path)
source_dir, source_file = os.path.split(source_file_path)
source_filename, _ = os.path.splitext(source_file)
output_filename = source_filename.replace(" ", "_").lower()
processed_file_path = os.path.join(source_dir, f"{output_filename}_processed.txt")
chapters_file_path = os.path.join(source_dir, f"{output_filename}_chapters.json")

table_of_content_manual = """
FIRST PART

1. A Runaway Reef
2. The Pros and Cons
3. As Master Wishes
4. Ned Land
5. At Random!
6. At Full Steam
7. A Whale of Unknown Species
8. “Mobilis in Mobili”
9. The Tantrums of Ned Land
10. The Man of the Waters
11. The Nautilus
12. Everything through Electricity
13. Some Figures
14. The Black Current
15. An Invitation in Writing
16. Strolling the Plains
17. An Underwater Forest
18. Four Thousand Leagues Under the Pacific
19. Vanikoro
20. The Torres Strait
21. Some Days Ashore
22. The Lightning Bolts of Captain Nemo
23. “Aegri Somnia”
24. The Coral Realm

SECOND PART

1. The Indian Ocean
2. A New Proposition from Captain Nemo
3. A Pearl Worth Ten Million
4. The Red Sea
5. Arabian Tunnel
6. The Greek Islands
7. The Mediterranean in Forty-Eight Hours
8. The Bay of Vigo
9. A Lost Continent
10. The Underwater Coalfields
11. The Sargasso Sea
12. Sperm Whales and Baleen Whales
13. The Ice Bank
14. The South Pole
15. Accident or Incident?
16. Shortage of Air
17. From Cape Horn to the Amazon
18. The Devilfish
19. The Gulf Stream
20. In Latitude 47° 24’ and Longitude 17° 28’
21. A Mass Execution
22. The Last Words of Captain Nemo
23. Conclusion
"""
table_of_content_str = ("Table of content provided by the novel:\n" + table_of_content_manual) if table_of_content_manual else ""

with open(source_file_path, 'r', encoding='utf-8') as f:
    full_novel_text = preprocess_remove_line_break(f.read(), " ")

with open(processed_file_path, "w", encoding='utf-8') as f:
    f.write(full_novel_text)

chunk_size = 50000
model_type = "claude"
model = "claude-opus-4-1-20250805"

head = 0
tail = min(head + chunk_size, len(full_novel_text))

llm = llm_factory(model_type=model_type, model_name=model, temperature=0.0)
existing_chapters = {}
existing_chapters_full_text = {}
all_finished = False
while head < len(full_novel_text) and not all_finished:
    chunk = full_novel_text[head:tail]
    # Buildinig the 
    if not existing_chapters:
        provided_chapters_str = "Empty"
    else:
        provided_chapters = {}
        for chapter_key, chapter_content in existing_chapters.items():
            provided_chapters[chapter_key] = chapter_content
        provided_chapters_str = json.dumps(existing_chapters, indent=4)
        
    json_prefilled = "{\n   \""
    messages = [
        {
            "role": "user",
            "content": chapter_divide_prompt.format(table_of_content=table_of_content_str, existing_chapters=provided_chapters_str, new_content=chunk)
        }
    ]
    prefilled_message = [{
        "role": "assistant",
        "content": json_prefilled
    }]
    
    finished = False
    backup_existing_chapters = existing_chapters.copy()
    backup_response = ""
    while not finished:
        response = llm.generate(
            raw_messages=messages + prefilled_message,
        )
        response_json_str = json_prefilled + response
        if backup_response == response_json_str:
            print("Loop detected, stopping.")
            all_finished = True
            break
        backup_response = response_json_str
        response_json = json_repair.loads(response_json_str)
        if isinstance(response_json, list):
            response_json = response_json[0]
        print(response_json)
        existing_chapters.update(response_json)
        if existing_chapters == backup_existing_chapters:
            print("No new chapters found, stopping.")
            all_finished = True
            break
        existing_chapters_full_text = {}
        last_start_point = 0
        try:
            for chapter_key, chapter_content in existing_chapters.items():
                start_index = full_novel_text.find(chapter_content["start"], last_start_point)
                if start_index == -1:
                    if full_novel_text.find(chapter_content["start"]) == -1:
                        raise QuotesNotFoundError(f"{chapter_key} start not found")
                    else:
                        raise ChapterOverlapError(f"{chapter_key} start overlapped")
                ending_start_point = start_index
                if chapter_content["end"]:
                    end_index = full_novel_text.find(chapter_content["end"], start_index)
                    if end_index == -1:
                        if full_novel_text.find(chapter_content["end"]) == -1:
                            raise QuotesNotFoundError(f"{chapter_key} end not found")
                        else:
                            raise ChapterOverlapError(f"{chapter_key} end overlapped")
                    end_index += len(chapter_content["end"])
                    ending_end_point = end_index
                    print(chapter_key, start_index, end_index)
                else:
                    end_index = len(full_novel_text)
                    print(chapter_key, start_index, "...")
                existing_chapters_full_text[chapter_key] = full_novel_text[start_index:end_index]
                last_start_point = end_index
            finished = True
        except ChapterPushOnError as e:
            print(e)
            ending_start_point = ending_end_point
            finished = True
        except QuotesNotFoundError as e:
            print(e)
            finished = False
            messages.append({
                "role": "assistant",
                "content": response_json_str
            })
            messages.append({
                "role": "user",
                "content": f"{e}Cannot match the quotes. Please fix the syntax and redo the extraction."
            })
        except ChapterOverlapError as e:
            print(e)
            finished = False
            messages.append({
                "role": "assistant",
                "content": response_json_str
            })
            messages.append({
                "role": "user",
                "content": f"{e}\nOverlap detected. Please rearrange the chapter boundaries and redo the extraction."
            })
    head = ending_start_point
    print("Current head:", head)
    tail = min(head + chunk_size, len(full_novel_text))
    
    with open(chapters_file_path, 'w', encoding='utf-8') as f:
        json.dump(existing_chapters_full_text, f, indent=4)