agent.py

import re
import os
from ollama_client import ask_ollama
from artifacts import ArtifactManager
from templates import render_page, slugify, VALID_ACCENTS

MAX_REGEN_ATTEMPTS = 1  # how many times to retry generation after a validation failure


class JijiAgent:
    def __init__(self, pet):
        self.pet = pet
        self.artifacts = ArtifactManager()
        self.codestyle = self._load_codestyle()

    def _load_codestyle(self) -> str:
        style_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "jiji_codestyle.md")
        try:
            with open(style_path, "r", encoding="utf-8") as f:
                return f.read()
        except FileNotFoundError:
            return ""

    # ─────────────────────────────────────────────
    # ROUTING (already deterministic — unchanged)
    # ─────────────────────────────────────────────

    def _fast_route(self, user_input: str) -> str:
        cleaned = user_input.lower().strip()
        existing_files = list(self.artifacts.metadata.keys())

        vision_triggers = [
            "screenshot", "what's on my screen", "what is on my screen",
            "look at my screen", "see this", "ocr", "read my screen",
            "what do you see", "what tabs", "what's open", "whats on my screen",
            "on my screen", "my screen", "my browser", "what am i looking at"
        ]
        code_create_triggers = [
            "landing page", "website", "webpage",
            "write html", "write css", "generate html", "generate css",
            "make a page", "make a site", "make a website", "build me a",
            "create a page", "create a site", "create a website",
            "make me a page", "make me a site", "make me a website",
            "self-contained", "single file html",
            "write me a script", "write a function", "generate a script",
        ]
        code_modify_triggers = [
            "change the", "change it", "edit the", "modify", "update the",
            "add a", "add the", "remove the", "fix the", "make it", "make the",
            "turn it", "switch the", "adjust", "move the", "rename",
            "background", "colour", "color", "font", "button", "navbar",
            "footer", "header", "section", "dark", "light", "padding", "margin"
        ]

        if any(kw in cleaned for kw in vision_triggers):
            return "VISION"
        if any(kw in cleaned for kw in code_create_triggers):
            return "CODE"
        if existing_files and any(kw in cleaned for kw in code_modify_triggers):
            return "CODE"
        return "CHAT"

    def _is_modification(self, user_input: str) -> bool:
        cleaned = user_input.lower().strip()
        create_signals = [
            "make me a", "make a", "create a", "build me", "build a",
            "generate a", "write me a", "new page", "new file", "from scratch", "brand new"
        ]
        if any(s in cleaned for s in create_signals):
            return False
        modify_signals = [
            "change", "edit", "modify", "update", "add", "remove", "fix",
            "make it", "make the", "turn it", "switch", "adjust", "rename",
            "background", "color", "colour", "font", "button", "navbar",
            "footer", "header", "section", "dark mode", "light mode"
        ]
        return any(s in cleaned for s in modify_signals)

    # ─────────────────────────────────────────────
    # MAIN ENTRY POINT
    # ─────────────────────────────────────────────

    def process_request(self, user_input: str, clipboard_data: str = "") -> dict:
        if clipboard_data:
            intent = "CHAT"
        else:
            intent = self._fast_route(user_input)

        print(f"[Jiji] Intent: {intent}")

        full_input = user_input
        if clipboard_data:
            full_input += f"\n\n[Clipboard Content]:\n{clipboard_data}"

        if intent == "CODE":
            return self._handle_code(full_input)
        elif intent == "VISION":
            return self._handle_vision(full_input)
        else:
            return self._handle_chat(full_input)

    # ─────────────────────────────────────────────
    # HANDLERS
    # ─────────────────────────────────────────────

    def _handle_chat(self, user_input: str) -> dict:
        response = ask_ollama(
            question=user_input,
            system_prompt=self.pet.get_system_prompt(),
            memory=self.pet.get_memory(),
            custom_options={"num_predict": 600, "temperature": 0.8},
            model="chat"
        )
        self.pet.add_to_memory("user", user_input)
        self.pet.add_to_memory("assistant", response)
        return {"text": response, "action": "chat", "state": "idle"}

    def _handle_code(self, user_input: str) -> dict:
        existing_files = list(self.artifacts.metadata.keys())
        is_edit = self._is_modification(user_input) and bool(existing_files)

        if is_edit:
            return self._handle_code_edit(user_input)
        return self._handle_code_create(user_input)

    # --- CREATE: template-based, model only picks accent + writes body markup ---

    def _handle_code_create(self, user_input: str) -> dict:
        system_prompt = (
            "This is not a conversation. You output page content only, nothing else.\n\n"
            "You do not write <!DOCTYPE>, <html>, <head>, <body>, <style> tags, or "
            "code fences — those are handled outside of you. Do NOT say 'Sure' or "
            "explain what you're about to do. Your response must start with the "
            "literal text 'TITLE:' as its very first characters.\n\n"
            "Output EXACTLY this format:\n\n"
            "TITLE: <short page title>\n"
            f"ACCENT: <one of: {', '.join(sorted(VALID_ACCENTS))}>\n"
            "BODY:\n"
            "<the nav/hero/section/card markup and copy that goes inside <body>. "
            "Use classes .btn .card .hero .grid where relevant — their CSS "
            "already exists, just use the class names.>\n\n"
            "Example, for \"make me a page for a coffee shop\":\n"
            "TITLE: Morning Brew Coffee\n"
            "ACCENT: amber\n"
            "BODY:\n"
            "<nav><div>Morning Brew</div><div><a href=\"#menu\">Menu</a></div></nav>\n"
            "<div class=\"hero\"><h1>Coffee That Wakes You Up</h1>"
            "<p>Small batch, locally roasted, every morning.</p>"
            "<button class=\"btn\">View Menu</button></div>\n"
        )

        result = self._generate_create_page(user_input, system_prompt)

        if result is None:
            return {
                "text": "hmm, generation failed. try a simpler prompt?",
                "action": "chat",
                "state": "idle"
            }

        title, full_html = result
        filename = slugify(title)
        save_result = self.artifacts.save_artifact(filename, full_html)

        return {
            "text": f"done! made you \"{title}\" ✨",
            "action": "code_gen",
            "file_path": save_result["path"],
            "version": save_result["version"],
            "state": "happy"
        }

    def _generate_create_page(self, user_input: str, system_prompt: str):
        """
        Tries the TITLE/ACCENT/BODY format first (template-guaranteed compliance).
        If the model ignores that and writes a full document instead — which small
        models do more often than you'd hope — falls back to extracting and
        validating that full document rather than hard-failing. Retries once with
        a reinforced instruction if neither path produces usable output.
        """
        prompt = user_input

        for attempt in range(MAX_REGEN_ATTEMPTS + 1):
            response = ask_ollama(
                question=prompt,
                system_prompt=system_prompt,
                memory=[],
                custom_options={"num_predict": 1800, "temperature": 0.2, "num_ctx": 8192},
                model="code",
                skip_clean=True,
            )

            title, accent, body = self._parse_page_response(response)

            if body:
                full_html = render_page(title=title, accent_key=accent, body_html=body)
                issues = self.artifacts.validate_html(full_html)
                if issues:
                    print(f"[Jiji] Template output failed validation (unexpected): {issues}")
                return title, full_html

            # Model ignored the format — see if it wrote a usable full document anyway.
            print(f"[Jiji] Format not followed (attempt {attempt + 1}). Response start:\n{response[:300]}")
            legacy_html = self.artifacts.extract_code(response)
            if legacy_html:
                issues = self.artifacts.validate_html(legacy_html)
                if not issues:
                    guessed_title = self._guess_title(legacy_html) or "Untitled"
                    return guessed_title, legacy_html
                print(f"[Jiji] Fallback full-document also failed validation: {issues}")

            prompt = (
                f"{user_input}\n\n"
                "Your previous response did not follow the required format — it "
                "started with an explanation or a full HTML document instead. "
                "Do NOT write <!DOCTYPE>, <html>, <head>, <body>, <style>, or code "
                "fences. Do NOT explain anything. Your response must start with "
                "the literal text 'TITLE:' as its very first characters."
            )

        return None

    def _guess_title(self, html: str) -> str:
        match = re.search(r'<title>(.*?)</title>', html, re.IGNORECASE | re.DOTALL)
        return match.group(1).strip() if match else ""

    def _parse_page_response(self, response: str) -> tuple[str, str, str]:
        title_match = re.search(r'TITLE:\s*(.+)', response)
        accent_match = re.search(r'ACCENT:\s*(\w+)', response, re.IGNORECASE)
        body_match = re.search(r'BODY:\s*\n([\s\S]*)', response)

        title = title_match.group(1).strip() if title_match else "Untitled"
        accent = accent_match.group(1).strip().lower() if accent_match else "purple"
        if accent not in VALID_ACCENTS:
            accent = "purple"
        body = body_match.group(1).strip() if body_match else ""

        # strip stray code fences if the model added them out of habit
        body = re.sub(r'^```[a-zA-Z]*\n?', '', body)
        body = re.sub(r'```$', '', body).strip()

        return title, accent, body

    # --- EDIT: model still needs full-file access, but now validated + reuses the known filename ---

    def _handle_code_edit(self, user_input: str) -> dict:
        existing_files = list(self.artifacts.metadata.keys())
        latest = max(
            existing_files,
            key=lambda f: self.artifacts.metadata[f].get("updated", "")
        )
        file_path = self.artifacts.metadata[latest]["path"]

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                current_code = f.read()
        except Exception as e:
            print(f"[Jiji] Could not read file: {e}")
            return {
                "text": "couldn't find the file to edit — try making it again?",
                "action": "chat",
                "state": "idle"
            }

        edit_prompt = (
            f"CURRENT FILE:\n```html\n{current_code}\n```\n\n"
            f"USER REQUEST: {user_input}\n\n"
            "Return the COMPLETE updated file. Self-contained. All CSS inside <style>. "
            "No external stylesheets."
        )

        system_prompt = (
            "You are a code generation engine. You output files, nothing else.\n\n"
            "OUTPUT FORMAT — follow exactly:\n"
            "```html\n"
            "The complete HTML file.\n"
            "```\n\n"
            f"STYLE GUIDE:\n{self.codestyle}\n\n"
            "ABSOLUTE RULES:\n"
            "- ZERO explanations outside the code block.\n"
            "- ALL CSS inside <style> tag. NEVER use <link> stylesheets.\n"
            "- body background MUST be dark (#0a0a0f or similar). NEVER white.\n"
            "- text color must be light (#e8e2f8 or white). NEVER black on white.\n"
        )

        clean_code, response = self._generate_and_validate_edit(edit_prompt, system_prompt)

        if not clean_code:
            return {
                "text": "hmm, edit failed. try a simpler prompt?",
                "action": "chat",
                "state": "idle"
            }

        # Always save back to the SAME filename we started with — we already
        # know it, we don't need (or want) the model telling us a new one.
        save_result = self.artifacts.save_artifact(latest, clean_code)

        display_text = re.sub(
            r'```(?:[a-zA-Z0-9+_#-]+)?\n[\s\S]*?```',
            '',
            response,
            flags=re.DOTALL
        ).strip()
        if not display_text:
            display_text = "done! ✨"

        return {
            "text": display_text,
            "action": "code_gen",
            "file_path": save_result["path"],
            "version": save_result["version"],
            "state": "happy"
        }

    def _generate_and_validate_edit(self, edit_prompt: str, system_prompt: str) -> tuple[str, str]:
        """Runs the edit generation, validates, retries once with a correction note if needed."""
        attempt_prompt = edit_prompt

        for attempt in range(MAX_REGEN_ATTEMPTS + 1):
            response = ask_ollama(
                question=attempt_prompt,
                system_prompt=system_prompt,
                memory=[],
                custom_options={"num_predict": 3000, "temperature": 0.2, "num_ctx": 8192},
                model="code",
                skip_clean=True,
            )
            clean_code = self.artifacts.extract_code(response)
            if not clean_code:
                continue

            issues = self.artifacts.validate_html(clean_code)
            if not issues:
                return clean_code, response

            print(f"[Jiji] Edit failed validation (attempt {attempt + 1}): {issues}")
            if attempt < MAX_REGEN_ATTEMPTS:
                attempt_prompt = (
                    f"{edit_prompt}\n\n"
                    f"Your previous attempt violated these rules: {', '.join(issues)}. "
                    "Fix them and return the complete corrected file."
                )

        # Ran out of attempts — return whatever we last got rather than fail outright
        return clean_code, response

    def _handle_vision(self, user_input: str) -> dict:
        try:
            from screen_reader import ScreenReader
            import time
            reader = ScreenReader()

            if not reader.is_available():
                return {
                    "text": "can't see your screen — pytesseract or pillow isn't installed. run: pip install pillow pytesseract and install tesseract-ocr",
                    "action": "chat",
                    "state": "idle"
                }

            time.sleep(0.8)
            screen_text = reader.capture_screen_text()

        except Exception as e:
            return {
                "text": f"screen reading broke: {str(e)[:80]}",
                "action": "chat",
                "state": "idle"
            }

        enriched = (
            f"{user_input}\n\n"
            f"[Screen Content — OCR extracted text from the user's screen]:\n{screen_text}\n\n"
            "Answer based on what you can see in the screen content above."
        )
        response = ask_ollama(
            question=enriched,
            system_prompt=self.pet.get_system_prompt(),
            memory=self.pet.get_memory(),
            custom_options={"num_predict": 300, "temperature": 0.7},
            model="chat"
        )
        self.pet.add_to_memory("user", user_input)
        self.pet.add_to_memory("assistant", response)
        return {"text": response, "action": "vision", "state": "idle"}