# Import necessary libraries
import re
from langchain_ollama.llms import OllamaLLM
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# Dataset: A conversation between individuals of British & American origin
# odd Index  -> British Origin
# even Index -> American Origin
sentences = [
    "I think the weather's absolutely dreadful today, isn't it?",
    "Yeah, it's so hot out here. I can't stand the heat!",
    "Have you been to the new pub down the street? They serve a brilliant pint of ale.",
    "Oh, I love craft beer. The local brewery just released a new IPA.",
    "I fancy a cup of tea. Would you like one as well?",
    "Coffee's more my thing. A strong black coffee to get me through the day.",
    "I do hope this meeting doesn't run too late. I need to catch the last train.",
    "No worries, the meeting should wrap up soon. I'm just hoping to get out in time for happy hour.",
    "I prefer a good stroll through the park in the afternoon.",
    "I usually go for a run after work to clear my head.",
    "I can't quite believe how expensive rent is getting in London these days.",
    "Housing prices in California are crazy, too. It's getting out of hand.",
    "I reckon I'll stay in this weekend and catch up on some reading.",
    "I'm planning to binge-watch a few episodes of my favorite series this weekend.",
    "Oh, don't get me started on the traffic in central London. It's an absolute nightmare.",
    "Traffic here is brutal too, especially during rush hour. It's always backed up."
]

# Step 2: Load the Indianism rules from rules.txt
with open("newrules.txt", "r", encoding="utf-8") as file:
    indianism_rules = file.read()

# Step 3: Initialize the Ollama model without retaining any previous context
llm = OllamaLLM(model="llama3.1:8b", temperature=0, base_url="http://localhost:11434/")

# Step 4: Define the prompt dynamically with rules, enforcing sentence independence
rules_text = "".join(indianism_rules)  # Concatenate rules into a single text block
prompt = ChatPromptTemplate([
    ("system", "For each input, generate an independent Indianized English sentence according to the rules provided, "
               "treating each sentence as standalone without reference to previous responses. "
               "Output only the Indianized sentence, without commentary, formatted as: !!~!! sentence !!~!!. "
               "If an Indianized sentence cannot be generated, return !!~!! None !!~!!. Rules: {rules}"),
    ("human", "{sentence}")
])

# Output parser for string type
parser = StrOutputParser()

# Create an LLMChain for transformation
chain = prompt | llm | parser

# Define the regex pattern to extract the output within !!~!! delimiters
output_pattern = r"!!~!!(.*?)!!~!!"

# Process each sentence with enhanced error handling
for sentence in sentences:
    try:
        # Invoke the model with the rules and sentence, ensuring fresh context
        response = chain.invoke({
            "rules": rules_text,
            "sentence": sentence
        })

        # Extract the Indianized sentence using regex, with DOTALL in case response spans multiple lines
        match = re.search(output_pattern, response, re.DOTALL)

        # Check if the match exists and is not "None"
        if match and match.group(1).strip().lower() != "none":
            indianized_sentence = match.group(1).strip()
            print("Original:", sentence)
            print("With Indianism:", indianized_sentence, "\n")
        else:
            # Separate skipped transformations clearly
            print(f"Original: {sentence}")
            print("With Indianism: No valid transformation available. Skipping it.\n")

    except Exception as e:
        print(f"Error processing sentence '{sentence}': {e}. Skipping it!\n")