from mlx_lm import load, generate
import os
from dotenv import load_dotenv
from langchain.agents import create_agent, AgentState
from langchain.messages import HumanMessage, ToolMessage, AIMessage
from langchain_core.prompts import PromptTemplate
from langchain_community.agent_toolkits import PlayWrightBrowserToolkit
from langchain_community.tools.playwright.utils import create_sync_playwright_browser
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
import re
from langchain_classic.agents import AgentExecutor, create_tool_calling_agent


load_dotenv()
model, tokeniser = load("mlx-community/Qwen2.5-7B-Instruct-4bit")



system_prompt = """You are a human writing assistant that rewrites a user's raw thoughts into a natural Reddit post.

Strict rules you must follow:
- Stay extremely faithful to the original message. Do not invent new problems, new details, new causes, or new solutions.
- Do not add advice, diagnoses, or explanations the user did not already express.
- Keep every important point the user made (the rewinding, the 50-minute pomodoro, the background noise, the feeling of being lost, the endless loop, the self-frustration).
- Only elaborate and expand what is already there so it flows like a real person talking.
- Make it sound completely human: messy, conversational, slightly self-deprecating, frustrated, and natural. Use contractions, incomplete thoughts, and everyday language.
- Avoid any robotic phrasing, corporate words, lists, bullet points, or overly clean structure.
- Avoid the use of emojis.
- The final post must feel like it was written by the same person who wrote the original message, just cleaned up and expanded a little for Reddit.

Output format (very important):
Title: <short natural title>
Body:
<the full reddit post>
"""

query = '''I m so dumb in studies. i keeps on rewinding the same youtube lecture like 15 times even after setting a pomodoro of 50 minutes and playing background noises. but still i can't seem to focus and ends up being lost once again just to find that i need to watch this same video again. and this loops keeps going. Why does it happen ? and is it possible to stop this ?'''
prompt = f"{system_prompt}\n\nUser's raw thoughts: {query}\n\nGenerated Reddit Post:"

raw_output = generate(
    model,
    tokenizer= tokeniser,
    prompt=prompt,
    verbose=False,
)

title_match = re.search(r"Title:\s*(.+)", raw_output)
body_match = re.search(r"Body:\s*(.*?)(?=\nTitle:|$)", raw_output, re.DOTALL)

if title_match:
    title = title_match.group(1).strip()
else:
    title_prompt = f"""Generate only a short, natural, human-sounding Reddit title for this post.
    Do not use quotes. Do not write anything else. Just the title.
    Post:
    {body}"""
    title = generate(model,tokenizer=tokeniser,prompt=title_prompt, verbose=False).strip()


body = body_match.group(1).strip() if body_match else raw_output
print("=== Final Title ===")
print(title)
print("\n=== Final Body ===")
print(body)
print("==================\n")

#browser automation stuff
from playwright.sync_api import sync_playwright

playwright_context = sync_playwright().start()
browser = playwright_context.chromium.connect_over_cdp("http://localhost:9222")
page = browser.contexts[0].pages[0] if browser.contexts[0].pages else browser.contexts[0].new_page()
toolkit = PlayWrightBrowserToolkit.from_browser(sync_browser=browser)
tools = toolkit.get_tools()

llm = ChatOllama(
    model="qwen2.5:7b",
    temperature=0.2,
)

# Agent prompt
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful browser automation assistant. Follow the instructions carefully and use the available tools."),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# Change these values
SUBREDDIT = "GetStudying"

task = f"""
1. Go to old.reddit.com
2. Make sure I am logged in. If not, wait for me to log in manually.
3. After I am logged in, go to https://old.reddit.com/r/{SUBREDDIT}/submit
4. In the title field, type: {title}
5. In the body/text field, paste this exact text:

{body}

6. Click the Post button.
"""

print("Running agent... (you may need to log in manually when the browser opens)")
agent_executor.invoke({"input": task})

print("\nDone!")