import os, json, glob, shutil
from typing import Optional, Any
from urllib.parse import quote
from lib_py.agents.base import BaseAgentAdapter, DiscoveryContext
from lib_py.verify_session import workspace_key


class GrokAgentAdapter(BaseAgentAdapter):
    @property
    def name(self) -> str:
        return 'grok'

    @property
    def own_key(self) -> str:
        return 'grok_session_id_own'

    @property
    def ready_tokens(self) -> str:
        return 'Grok|xAI|Assistant|❯|>>>'

    @property
    def exit_key(self) -> str:
        return '/exit'

    @property
    def delegate_agent_key(self) -> str:
        return 'grok-build'

    @property
    def identity_cache_fields(self) -> tuple:
        return ('session_id', 'session_jsonl')

    @property
    def input_prompt(self) -> str:
        return '❯'

    @property
    def input_placeholder(self) -> str:
        return ''

    @property
    def input_rule_pattern(self) -> str:
        return '─{10,}'

    def _ws_dir(self, ctx: DiscoveryContext) -> str:
        return quote(os.path.realpath(ctx.cwd), safe='')

    def artifact_path(self, uuid: str, ctx: DiscoveryContext) -> str:
        return f"{ctx.home_dir}/.grok/sessions/{self._ws_dir(ctx)}/{uuid}/chat_history.jsonl"

    def verify_artifact(self, uuid: str, ctx: DiscoveryContext) -> bool:
        path = self.artifact_path(uuid, ctx)
        if not os.path.exists(path):
            return False
        if ctx.epoch and os.path.getmtime(path) < ctx.epoch:
            return False
        try:
            valid_session = False
            found_cwd = None
            with open(path) as f:
                for _ in range(50):
                    line = f.readline()
                    if not line:
                        break
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        payload = json.loads(line)
                        if payload.get("session_id") == uuid or payload.get("sessionId") == uuid:
                            valid_session = True
                        if payload.get("cwd"):
                            found_cwd = payload.get("cwd")
                            break
                    except Exception:
                        pass
            if not valid_session and found_cwd is None:
                # Directory path itself encodes the workspace and session UUID
                valid_session = os.path.isdir(os.path.dirname(path))
            if not valid_session:
                return False
            if found_cwd and workspace_key(found_cwd) != workspace_key(ctx.cwd):
                return False
        except Exception:
            return False
        return True

    def purge_artifacts(self, uuid: str, ctx: DiscoveryContext) -> list:
        purged = []
        sess_dir = os.path.dirname(self.artifact_path(uuid, ctx))
        if os.path.isdir(sess_dir):
            shutil.rmtree(sess_dir, ignore_errors=True)
            purged.append(sess_dir)
        elif os.path.exists(self.artifact_path(uuid, ctx)):
            os.remove(self.artifact_path(uuid, ctx))
            purged.append(self.artifact_path(uuid, ctx))
        return purged

    def spawn_spec(self, binary: str, session_uuid: str = "", use_wrapper: bool = False) -> str:
        if session_uuid:
            return f"{binary} --session-id {session_uuid} --permission-mode bypassPermissions"
        return f"{binary} --permission-mode bypassPermissions"

    def resume_spec(self, binary: str, session_uuid: str, materialized: bool = False) -> str:
        if materialized and session_uuid:
            return f"{binary} --resume {session_uuid} --permission-mode bypassPermissions"
        if session_uuid:
            return f"{binary} --session-id {session_uuid} --permission-mode bypassPermissions"
        return f"{binary} --permission-mode bypassPermissions"

    def auth_ok(self, run_cmd: Optional[Any] = None) -> bool:
        if os.environ.get("XAI_API_KEY"):
            return True
        home = os.environ.get("HOME_DIR") or os.environ.get("HOME") or os.path.expanduser("~")
        return os.path.exists(f"{home}/.grok/auth.json")

    def discover(self, ctx: DiscoveryContext) -> list:
        root = f"{ctx.home_dir}/.grok/sessions/{self._ws_dir(ctx)}"
        if not os.path.isdir(root):
            return []
        entries = [d for d in glob.glob(f"{root}/*") if os.path.isdir(d)]
        entries.sort(key=os.path.getmtime, reverse=True)
        candidates = []
        for d in entries:
            cand = os.path.basename(d)
            if cand and self.verify_artifact(cand, ctx):
                candidates.append(cand)
        return candidates