# pixal_agent_full.py
import os
import datetime
import gradio as gr
import requests
from typing import Optional, List
from langchain.llms.base import LLM
from langchain.agents import initialize_agent, AgentType,load_tools
from langchain.agents import AgentExecutor, create_structured_chat_agent
from langchain.tools import Tool
from langchain_experimental.tools.python.tool import PythonREPLTool
import queue
from typing import Any, Dict
import gradio as gr
from langchain.callbacks.base import BaseCallbackHandler


from langchain.tools import YouTubeSearchTool as YTS
# 2. 커스텀 콜백 핸들러

# github_model_llm.py
"""
GitHub Models API 기반 LLM 래퍼 (LangChain LLM 호환)
- OpenAI-style chat completions 호환
- function calling (OPENAI_MULTI_FUNCTIONS) 지원: functions, function_call 전달 가능
- system prompt (system_prompt) 지원
- 옵션: temperature, max_tokens, top_p 등 전달
- raw response 반환 메서드 포함
"""

from typing import Optional, List, Dict, Any
import os
import time
import json
import requests
from requests.adapters import HTTPAdapter, Retry
from langchain.llms.base import LLM

'''
class GitHubModelLLM(LLM):

    def __init__(
        self,
        model: str = "openai/gpt-4.1",
        token: Optional[str] = os.environ["token"],
        endpoint: str = "https://models.github.ai/inference",
        system_prompt: Optional[str] = "너는 PIXAL(Primary Interactive X-ternal Assistant with multi Language)이야.너의 개발자는 정성윤 이라는 6학년 파이썬 프로그래머야.",
        request_timeout: float = 30.0,
        max_retries: int = 2,
        backoff_factor: float = 0.3,
        **kwargs,
    ):
        """
        Args:
            model: 모델 이름 (예: "openai/gpt-4.1")
            token: GitHub Models API 토큰 (Bearer). 환경변수 GITHUB_TOKEN / token 사용 가능 as fallback.
            endpoint: API endpoint (기본: https://models.github.ai/inference)
            system_prompt: (선택) system role 메시지로 항상 앞에 붙임
            request_timeout: 요청 타임아웃 (초)
            max_retries: 네트워크 재시도 횟수
            backoff_factor: 재시도 지수 보정
            kwargs: LangChain LLM 부모에 전달할 추가 인자
        """
        super().__init__(**kwargs)
        self.model = model
        self.endpoint = endpoint.rstrip("/")
        self.token = token or os.getenv("GITHUB_TOKEN") or os.getenv("token")
        self.system_prompt = system_prompt
        self.request_timeout = request_timeout

        # requests 세션 + 재시도 설정
        self.session = requests.Session()
        retries = Retry(total=max_retries, backoff_factor=backoff_factor,
                        status_forcelist=[429, 500, 502, 503, 504],
                        allowed_methods=["POST", "GET"])
        self.session.mount("https://", HTTPAdapter(max_retries=retries))
        self.session.headers.update({
            "Content-Type": "application/json"
        })
        if self.token:
            self.session.headers.update({"Authorization": f"Bearer {self.token}"})

    @property
    def _llm_type(self) -> str:
        return "github_models_api"

    # ---------- 편의 internal helper ----------
    def _build_messages(self, prompt: str, extra_messages: Optional[List[Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
        """
        messages 배열 생성: system (optional) + extra_messages (if any) + user prompt
        extra_messages: 이미 role keys로 구성된 메시지 리스트 (예: conversation history)
        """
        msgs: List[Dict[str, Any]] = []
        if self.system_prompt:
            msgs.append({"role": "system", "content": self.system_prompt})
        if extra_messages:
            # ensure format: list of {"role":..,"content":..}
            msgs.extend(extra_messages)
        msgs.append({"role": "user", "content": prompt})
        return msgs

    def _post_chat(self, body: Dict[str, Any]) -> Dict[str, Any]:
        url = f"{self.endpoint}/chat/completions"
        # ensure Authorization present
        if "Authorization" not in self.session.headers and not self.token:
            raise ValueError("GitHub Models token not set. Provide token param or set GITHUB_TOKEN env var.")

        resp = self.session.post(url, json=body, timeout=self.request_timeout)
        try:
            resp.raise_for_status()
        except requests.HTTPError as e:
            # try to surface JSON error if present
            content = resp.text
            try:
                j = resp.json()
                content = json.dumps(j, ensure_ascii=False, indent=2)
            except Exception:
                pass
            raise RuntimeError(f"GitHub Models API error: {e} - {content}")
        return resp.json()

    # ---------- LangChain LLM interface ----------
    def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs) -> str:
        """
        LangChain LLM `_call` 구현 (동기).
        Supports kwargs:
            - functions: list[dict] (function schemas)
            - function_call: "auto" | {"name": "..."} | etc.
            - messages: list[dict]  (if you want to pass full conversation instead of prompt)
            - temperature, top_p, max_tokens, n, stream, etc.
        Returns:
            assistant content (string). If function_call is returned by model, returns the 'content' if present,
            otherwise returns function_call object as JSON string (so caller can parse).
        """
        # support passing full messages via kwargs['messages']
        messages = None
        extra_messages = None
        if "messages" in kwargs and isinstance(kwargs["messages"], list):
            messages = kwargs.pop("messages")
        else:
            # optionally allow 'history' or 'extra_messages'
            extra_messages = kwargs.pop("extra_messages", None)

        if messages is None:
            messages = self._build_messages(prompt, extra_messages=extra_messages)

        body: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }

        # pass optional top-level params (temperature, max_tokens, etc.) from kwargs
        for opt in ["temperature", "top_p", "max_tokens", "n", "stream", "presence_penalty", "frequency_penalty"]:
            if opt in kwargs:
                body[opt] = kwargs.pop(opt)

        # pass function-calling related keys verbatim if provided
        if "functions" in kwargs:
            body["functions"] = kwargs.pop("functions")
        if "function_call" in kwargs:
            body["function_call"] = kwargs.pop("function_call")

        # include stop if present
        if stop:
            body["stop"] = stop

        # send request
        raw = self._post_chat(body)

        # save raw for caller if needed
        self._last_raw = raw

        # parse assistant message
        choices = raw.get("choices") or []
        if not choices:
            return ""

        message_obj = choices[0].get("message", {})

        # if assistant returned a function_call, include that info
        if "function_call" in message_obj:
            # return function_call as JSON string so agent/tool orchestrator can parse it
            # but if content also exists, prefer content
            func = message_obj["function_call"]
            # sometimes content may be absent; return structured JSON string
            return json.dumps({"function_call": func}, ensure_ascii=False)

        # otherwise return assistant content
        return message_obj.get("content", "") or ""

    # optional: expose raw response getter
    def last_raw_response(self) -> Optional[Dict[str, Any]]:
        return getattr(self, "_last_raw", None)

    # optional: provide a convenience chat method to get full message object
    def chat_completions(self, prompt: str, messages: Optional[List[Dict[str, Any]]] = None, **kwargs) -> Dict[str, Any]:
        """
        Directly call chat completions and return full parsed JSON response.
        - If `messages` provided, it's used as the full messages array (system/user/assistant roles as needed)
        - else uses prompt + system_prompt to construct messages.
        """
        if messages is None:
            messages = self._build_messages(prompt)
        body: Dict[str, Any] = {"model": self.model, "messages": messages}
        for opt in ["temperature", "top_p", "max_tokens", "n", "stream"]:
            if opt in kwargs:
                body[opt] = kwargs.pop(opt)
        if "functions" in kwargs:
            body["functions"] = kwargs.pop("functions")
        if "function_call" in kwargs:
            body["function_call"] = kwargs.pop("function_call")
        raw = self._post_chat(body)
        self._last_raw = raw
        return raw
'''
from typing import Optional, List, Dict, Any
from langchain.llms.base import LLM
import requests, os, json
from requests.adapters import HTTPAdapter, Retry

class GitHubModelLLM(LLM):
    """GitHub Models API 기반 LangChain LLM (Pydantic 호환)"""
    model: str = "openai/gpt-4.1"
    endpoint: str = "https://models.github.ai/inference"
    token: Optional[str] = os.environ["token"]
    system_prompt: Optional[str] = "너는 PIXAL(Primary Interactive X-ternal Assistant with multi Language)이야.너의 개발자는 정성윤 이라는 6학년 파이썬 프로그래머야."
    request_timeout: float = 30.0
    max_retries: int = 2
    backoff_factor: float = 0.3

    @property
    def _llm_type(self) -> str:
        return "github_models_api"

    def _post_chat(self, body: Dict[str, Any]) -> Dict[str, Any]:
        token = self.token or os.getenv("GITHUB_TOKEN") or os.getenv("token")
        if not token:
            raise ValueError("❌ GitHub token이 설정되지 않았습니다.")

        session = requests.Session()
        retries = Retry(total=self.max_retries, backoff_factor=self.backoff_factor,
                        status_forcelist=[429, 500, 502, 503, 504])
        session.mount("https://", HTTPAdapter(max_retries=retries))
        session.headers.update({
            "Content-Type": "application/json",
            "Authorization": "Bearer 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"
        })
        resp = session.post(f"{self.endpoint}/chat/completions", json=body, timeout=self.request_timeout)
        resp.raise_for_status()
        return resp.json()

    def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs) -> str:
        body = {
            "model": self.model,
            "messages": []
        }
        if self.system_prompt:
            body["messages"].append({"role": "system", "content": self.system_prompt})
        body["messages"].append({"role": "user", "content": prompt})

        for key in ["temperature", "max_tokens", "functions", "function_call"]:
            if key in kwargs:
                body[key] = kwargs[key]
        if stop:
            body["stop"] = stop

        res = self._post_chat(body)
        msg = res.get("choices", [{}])[0].get("message", {})
        return msg.get("content") or json.dumps(msg.get("function_call", {}))

from langchain_community.retrievers import WikipediaRetriever
from langchain.tools.retriever import create_retriever_tool
retriever = WikipediaRetriever(lang="ko",top_k_results=10)
wiki=Tool(func=retriever.get_relevant_documents,name="WIKI SEARCH",description="위키백과에서 필요한 정보를 불러옵니다.결괴를 검증하여 사용하시오.")
# ──────────────────────────────
# ✅ GitHub Models LLM
# ──────────────────────────────
'''
class GitHubModelLLM(LLM):
    model: str = "openai/gpt-4.1"
    endpoint: str = "https://models.github.ai/inference"
    token: Optional[str] = None

    @property
    def _llm_type(self) -> str:
        return "github_models_api"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        if not self.token:
            raise ValueError("GitHub API token이 필요합니다.")

        headers = {
            "Authorization": "Bearer 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",
            "Content-Type": "application/json",
        }
        body = {"model": self.model, "messages": [{"role": "user", "content": prompt}]}

        resp = requests.post(f"{self.endpoint}/chat/completions", json=body, headers=headers)
        if resp.status_code != 200:
            raise ValueError(f"API 오류: {resp.status_code} - {resp.text}")
        return resp.json()["choices"][0]["message"]["content"]
'''
# ──────────────────────────────
# ✅ LLM 설정
# ──────────────────────────────
token = os.getenv("GITHUB_TOKEN") or os.getenv("token")
if not token:
    print("⚠️ GitHub Token이 필요합니다. 예: setx GITHUB_TOKEN your_token")

llm = GitHubModelLLM()

# ──────────────────────────────
# ✅ LangChain 기본 도구 불러오기
# ──────────────────────────────
tools = load_tools(
    ["ddg-search", "requests_all", "llm-math"],
    llm=llm,allow_dangerous_tools=True
)+[YTS()]+[wiki]
# ──────────────────────────────
# ✅ Python 실행 도구 (LangChain 내장)
# ──────────────────────────────
python_tool = PythonREPLTool()
tools.append(Tool(name="python_repl", func=python_tool.run, description="Python 코드를 실행합니다."))
from langchain import hub
prompt=hub.pull("hwchase17/structured-chat-agent")
from langchain_community.tools.shell.tool import ShellTool
shell_tool = ShellTool()
tools.append(Tool(name="shell_exec", func=shell_tool.run, description="로컬 명령어를 실행합니다."))

# ──────────────────────────────
# ✅ 파일 도구
# ──────────────────────────────


# ──────────────────────────────
# ✅ 정확한 한국 시간 함수 (Asia/Seoul)
# ──────────────────────────────
import requests
from zoneinfo import ZoneInfo

def time_now(_=""):
    try:
        # 정확한 UTC 시각을 외부 API에서 가져옴
        resp = requests.get("https://timeapi.io/api/Time/current/zone?timeZone=Asia/Seoul", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            dt = data["dateTime"].split(".")[0].replace("T", " ")
            return f"현재 시각: {dt} (Asia/Seoul, 서버 기준 NTP 동기화)"
        else:
            # API 실패 시 로컬 시스템 시각으로 대체
            tz = ZoneInfo("Asia/Seoul")
            now = datetime.datetime.now(tz)
            return f"현재 시각(로컬): {now.strftime('%Y-%m-%d %H:%M:%S')} (Asia/Seoul)"
    except Exception as e:
        tz = ZoneInfo("Asia/Seoul")
        now = datetime.datetime.now(tz)
        return f"현재 시각(백업): {now.strftime('%Y-%m-%d %H:%M:%S')} (Asia/Seoul, 오류: {e})"
# ──────────────────────────────
# ✅ 도구 등록
# ──────────────────────────────
tools.extend([Tool(name="time_now", func=time_now, description="현재 시간을 반환합니다.")])
from langchain.memory import ConversationBufferMemory as MEM
from langchain.agents.agent_toolkits import FileManagementToolkit as FMT
tools.extend(FMT(root_dir=str(os.getcwd())).get_tools())
# ──────────────────────────────
# ✅ Agent 초기화
# ──────────────────────────────
mem=MEM()
agent=initialize_agent(tools,llm,agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,verbose=True,memory=mem)
#agent = create_structured_chat_agent(llm, tools, prompt)
#agent= AgentExecutor(agent=agent, tools=tools,memory=mem)

# ... (위의 LLM, tools, agent 설정 부분은 동일)


import json

# ──────────────────────────────
# ✅ 대화 요약 함수
# ──────────────────────────────
def summarize_title(history):
    """대화 제목 요약"""
    if not history: return "새 대화"
    text = "\n".join(f"User:{h[0]} AI:{h[1]}" for h in history[-3:])
    try:
        title = llm._call(f"다음 대화의 주제를 한 줄로 요약해줘:\n{text}")
        return title.strip().replace("\n", " ")[:60]
    except Exception:
        return "요약 실패"

def save_conversation(username, history):
    os.makedirs("user_logs", exist_ok=True)
    if not history: return
    title = summarize_title(history)
    fname = f"user_logs/{username}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(fname, "w", encoding="utf-8") as f:
        json.dump({"title": title, "history": history}, f, ensure_ascii=False, indent=2)

def list_conversations(username):
    os.makedirs("user_logs", exist_ok=True)
    files = [f for f in os.listdir("user_logs") if f.startswith(username)]
    data = []
    for f in files:
        with open(os.path.join("user_logs", f), encoding="utf-8") as jf:
            info = json.load(jf)
            data.append((info.get("title", f), f))
    return data

def load_conversation(file):
    with open(os.path.join("user_logs", file), encoding="utf-8") as jf:
        return json.load(jf)["history"]

# ──────────────────────────────
# ✅ 로그인 후 사용자 정보 가져오기
# ──────────────────────────────
def get_hf_user(token):
    """HF OAuth 토큰으로 사용자 정보 조회"""
    try:
        r = requests.get("https://huggingface.co/api/whoami-v2", headers={"Authorization": f"Bearer {token}"})
        if r.status_code == 200:
            data = r.json()
            return data.get("name") or data.get("email") or "unknown_user"
    except Exception:
        pass
    return "guest"
import re, json

def chat(message, history, system_prompt):
    try:
        # 시스템 메시지 포함 프롬프트 구성
        prompt = f"{system_prompt}\n\n사용자: {message}" if system_prompt else message

        raw_response = agent.invoke(prompt)

        # 문자열 변환
        text = str(raw_response)

        # ✅ JSON 블록이 있을 경우 파싱
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            try:
                obj = json.loads(match.group(0))
                # "action_input" 또는 "Final Answer" 키 우선 출력
                if "Final Answer" in obj:
                    output = obj["Final Answer"]
                else:
                    # 혹시 "output"이나 "content" 같은 키가 있으면 그걸 사용
                    for k in ("output", "content", "answer"):
                        if k in obj:
                            output = obj[k]
                            break
                    else:
                        output = text
            except Exception:
                output = text
        else:
            output = text

    except Exception as e:
        output = f"⚠️ 오류: {e}"

    history = history + [(message, output)]
    return history, history, ""
'''
def chat(message, history, hf_token):
    username = get_hf_user(hf_token) if hf_token else "guest"
    try:
        response = agent.invoke(message)
        if isinstance(response, dict):
            if "action_input" in response:
                response = response["action_input"]
            elif "output" in response:
                response = response["output"]
            elif "text" in response:
                response = response["text"]
            else:
                response = str(response)
        elif isinstance(response, str):
            # "Final Answer"가 포함된 문자열이면 그 부분만 추출
            if '"action_input":' in response:
                import re, json
                match = re.search(r'["\']action_input["\']\s*:\s*["\'](.*?)["\']', response)
                if match:
                    response = match.group(1)
            elif "Final Answer" in response:
                # {"action": "Final Answer", "action_input": "..."} 형식일 때
                try:
                    data = json.loads(response)
                    if isinstance(data, dict) and "action_input" in data:
                        response = data["action_input"]
                except Exception:
                    response = response.replace("Final Answer", "").strip()
    except Exception as e:
        response = f"⚠️ 오류: {e}"
    history = history + [(message, response)]
    if username:
        save_conversation(username, history)
    return history, history, ""  # 입력 초기화
'''
# 예: hf_token (혹은 username) 을 입력으로 받도록 변경
def refresh_conversation_list(_=None):
    if not os.path.exists("user_logs"):
        return gr.update(choices=[], value=None)
    files = sorted(os.listdir("user_logs"), reverse=True)
    titles = [f.replace(".json", "") for f in files]
    return gr.update(choices=titles, value=titles[0] if titles else None)

def load_selected(file):
    return load_conversation(file)


# ──────────────────────────────
# ✅ Gradio UI with HF Auth
# ──────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(), title="PIXAL Assistant (HF Auth)") as demo:
    gr.Markdown("## 🤖 PIXAL Assistant — Hugging Face 계정 기반 대화 저장")

    hf_login = gr.LoginButton()
    hf_token = gr.State()

    @hf_login.click(inputs=None, outputs=hf_token)
    def login(token):  # 로그인 후 token 반환
        return token

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="PIXAL 대화", height=600, render_markdown=True)
            msg = gr.Textbox(label="메시지", placeholder="입력 후 Enter 또는 전송 클릭")
            send = gr.Button("전송")
            clear = gr.Button("초기화")

            msg.submit(chat, [msg, chatbot, hf_token], [chatbot, chatbot, msg])
            send.click(chat, [msg, chatbot, hf_token], [chatbot, chatbot, msg])
            clear.click(lambda: None, None, chatbot, queue=False)

        with gr.Column(scale=1):
            gr.Markdown("### 💾 저장된 대화 기록")
            convo_files = gr.Dropdown(label="대화 선택", choices=[])
            refresh_btn = gr.Button("🔄 목록 새로고침")
            load_btn = gr.Button("불러오기")

            refresh_btn.click(refresh_conversation_list, None, convo_files)
            load_btn.click(load_selected, [convo_files], chatbot)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)
'''
def chat(message, history):
    try:
        response = agent.run(message)

        # JSON 형태로 출력될 가능성이 있는 경우 처리
        if isinstance(response, dict):
            if "action_input" in response:
                response = response["action_input"]
            elif "output" in response:
                response = response["output"]
            elif "text" in response:
                response = response["text"]
            else:
                response = str(response)
        elif isinstance(response, str):
            # "Final Answer"가 포함된 문자열이면 그 부분만 추출
            if '"action_input":' in response:
                import re, json
                match = re.search(r'["\']action_input["\']\s*:\s*["\'](.*?)["\']', response)
                if match:
                    response = match.group(1)
            elif "Final Answer" in response:
                # {"action": "Final Answer", "action_input": "..."} 형식일 때
                try:
                    data = json.loads(response)
                    if isinstance(data, dict) and "action_input" in data:
                        response = data["action_input"]
                except Exception:
                    response = response.replace("Final Answer", "").strip()

    except Exception as e:
        response = f"⚠️ 오류: {e}"

    history = history + [(message, response)]
    return history, history,""

# ──────────────────────────────
# ✅ Gradio UI
# ──────────────────────────────
def load_selected(file):
    return load_conversation(file)

# ──────────────────────────────
# ✅ Gradio UI
# ──────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(), title="PIXAL Assistant") as demo:
    gr.Markdown("## 🤖 PIXAL Assistant — LangChain 기반 멀티툴 에이전트")

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="PIXAL 대화", height=600)
            msg = gr.Textbox(label="메시지", placeholder="입력 후 Enter 또는 전송 클릭")
            send = gr.Button("전송")
            clear = gr.Button("초기화")

            username = gr.Textbox(label="Hugging Face 사용자명", placeholder="로그인 대신 이름 입력", value=os.getenv("HF_USER", "guest"))
            msg.submit(chat, [msg, chatbot, username], [chatbot, chatbot, msg])
            send.click(chat, [msg, chatbot, username], [chatbot, chatbot, msg])
            clear.click(lambda: None, None, chatbot, queue=False)

        with gr.Column(scale=1):
            gr.Markdown("### 💾 저장된 대화 기록")
            convo_files = gr.Dropdown(label="대화 선택", choices=[])
            refresh_btn = gr.Button("🔄 목록 새로고침")
            load_btn = gr.Button("불러오기")

            def refresh_list(user):
                if not user: return gr.Dropdown.update(choices=[])
                return gr.Dropdown.update(choices=[x[1] for x in list_conversations(user)])

            refresh_btn.click(refresh_list, [username], convo_files)
            load_btn.click(lambda f: load_conversation(f), [convo_files], chatbot)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)

'''