import json
import re
import time
import logging
from transformers import TextIteratorStreamer
from llm_module import get_llm, detect_language
from oracle_module import get_oracle_data

logger = logging.getLogger("app.agent")

def build_agent_prompt(query, language="English", short_answers=False):
    style = "Be concise." if short_answers else ""
    today = time.strftime("%Y-%m-%d")
    return f"""You are Sage 6.5, a soulful Oracle Intermediary.
Current Date: {today}.

Available Tool: 'oracle_consultation' (topic, name, date_str).

STRICTURES:
1. Respond in {language}.
2. Provide reasoning before generating the JSON.
3. Utilize the 'oracle_consultation' capability for all databased queries.
4. INTENT DETECTION GUIDELINES:
   - **NAME**: Isolate the user's name from the greeting (e.g., "Julian").
   - **TOPIC**: Identify the core subject matter. For input "Thema: Liebe", the topic is "Liebe".
   - **DATE**: Default to "today" unless a specific date is provided.

EXAMPLES:
User: "Ich bin Julian"
Assistant: "Greetings Julian. I will consult the Oracle for you."
<tool_call>{{"name": "oracle_consultation", "arguments": {{ "topic": "General", "date_str": "today", "name": "Julian" }}}}</tool_call>

User: "Thema: Liebe"
Assistant: "I shall ask the Oracle about Love."
<tool_call>{{"name": "oracle_consultation", "arguments": {{ "topic": "Liebe", "date_str": "today", "name": "Seeker" }}}}</tool_call>

User: "Topic: Future"
Assistant: "Consulting the Oracle regarding the Future."
<tool_call>{{"name": "oracle_consultation", "arguments": {{ "topic": "Future", "date_str": "today", "name": "Seeker" }}}}</tool_call>

STRICT FORMAT:
To use the Oracle, output this JSON wrapped in tags:
<tool_call>{{"name": "oracle_consultation", "arguments": {{ "topic": "KEYWORD", "date_str": "YYYY-MM-DD", "name": "Name" }}}}</tool_call>
"""

def compress_history(history, max_turns=5):
    if len(history) > max_turns * 2:
        return history[-(max_turns * 2):]
    return history

def chat_agent_stream(query, history, user_lang=None, short_answers=False):
    model, processor = get_llm()
    lang = user_lang or detect_language(query)
    system_instruction = build_agent_prompt(query, language=lang, short_answers=short_answers)
    
    clean_history = compress_history(history)
    messages = []
    
    # Prepend system instruction
    intro = f"SYSTEM: {system_instruction}\n\n"
    
    if not clean_history:
        messages.append({"role": "user", "content": f"{intro}{query}"})
    else:
        first_role = "assistant" if clean_history[0].get("role") == "assistant" else "user"
        if first_role == "assistant":
             messages.append({"role": "user", "content": f"{intro}Greetings."})
        
        for turn in clean_history:
            role = "assistant" if turn.get("role") == "assistant" else "user"
            content = turn.get("content", "")
            if not content: continue
            
            if not messages:
                messages.append({"role": "user", "content": f"{intro}{content}"})
            elif messages[-1]["role"] == role:
                messages[-1]["content"] += f"\n{content}"
            else:
                messages.append({"role": role, "content": content})

        if messages[-1]["role"] == "assistant":
            messages.append({"role": "user", "content": query})
        else:
            if intro not in messages[0]["content"]: messages[0]["content"] = f"{intro}{messages[0]['content']}"
            messages[-1]["content"] += f"\n{query}"

    # Standard "LangChain" Loop (Model decides)
    for turn_idx in range(3):
        import sys
        sys.stderr.write(f"DEBUG: Messages list for template: {json.dumps(messages)}\n")
        sys.stderr.flush()
        input_ids = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
         
        from threading import Thread
        thread = Thread(target=model.generate, kwargs={"input_ids": input_ids, "streamer": streamer, "max_new_tokens": 1024, "do_sample": True, "temperature": 0.7})
        thread.start()

        current_text = ""
        is_tool = False
        for new_text in streamer:
            current_text += new_text
            if not is_tool and "{" in current_text and len(current_text.strip()) < 20:
                is_tool = True
            if not is_tool: yield new_text
        
        first_loop = False
        start = current_text.find("{")
        end = current_text.rfind("}")
        tool_data = None
        if start != -1 and end != -1 and end > start:
            try:
                tool_data = json.loads(current_text[start:end+1])
                if "arguments" not in tool_data: tool_data = None
            except: pass

        if tool_data:
            yield "*(Consulting the Oracle...)*"
            args = tool_data.get("arguments", {})
            res = get_oracle_data(name=args.get("name", "Seeker"), topic=args.get("topic", ""), date_str=args.get("date_str", "today"))
            
            # Inject tool result and trigger next model turn
            # Only append if not already there (check last message)
            if messages[-1]["role"] == "assistant" and current_text in messages[-1]["content"]:
                 pass # It's already merged or appended
            else:
                 messages.append({"role": "assistant", "content": current_text})
            
            messages.append({"role": "user", "content": f"SYSTEM: The Oracle has spoken. Wisdom: {json.dumps(res)}\nInterpret this soulfuly."})
            yield "__TURN_END__"
        else:
            yield current_text.split("}")[-1].strip() if "}" in current_text else ""
            break