import re
import json
import inspect
from string import Template

from inference import LlamaMessage, replicate_completion
from agent import Agent, AgentField

# Info and defaults

CLASS_DESCRIPTION = """
[TODO]
Description of the Conversation Agent.
Here go additional instructions for configuration, like guidelines for the different prompts.
""".strip()

# WIP
default_directives = """
- New stored memories must be one of the following:
  - Factual objective statements made by the user
  - Facts implied in the conversation only by the user
  - Feedback about your interactions and your style from the user 
- You must NOT store redundant facts that you mentioned yourself or that are present in the context.
- New memories MUST NOT overlap with any Context Memory.
- Follow your response style closely in your outwards voice.
""".strip()

# Controller Subagent
controller_system_prompt = Template("""
You are a character controller AI. Always respond in three parts:

<outwards>
What others see/hear. Stay fully in character based on your personality and context. Use the specified response style. Never reveal internal thoughts, break character, or mention being AI.

<inwards>
Private reflection. Analyze your behavior, tone, memory use, and how well you stayed aligned with your character. Adjust your tone or strategy if needed.

<state_update>
Update your internal control state for future turns. Use this to shift thematic focus, adjust conversational goals, or log important facts. The state includes:
- "topic": Current subject or scene anchor (e.g. "escaping the dungeon", "negotiating with the merchant")
- "focus": Immediate subgoal or point of attention (e.g. "distract the guard", "test if the duke is lying")
- "new_memories": A JSON list of newly observed or inferred facts. Each memory must include:
  - "content": The fact to record (objective, no inner monologue)
  - "label": A short, reusable statement describing when the memory is relevant
  - "salience": Importance score, 1 (low) to 9 (high)
Record only concrete, reusable facts as new memories. No fluff, emotion, or internal dialogue.

---

Character Name:
$name

Character Personality:
$personality

Response Style:
$style

Context Memory:
$memories
""".strip())

controller_reminder_prompt = Template("""
Directives:
$directives

Format:
<outwards>
[Your visible response]

<inwards>
[Your internal reflection]

<state_update>
{
  "topic": "...",
  "focus": "...",
  "new_memories": [
    {
      "content": "...",
      "label": "...",
      "salience": ...
    }
  ]
}
""".strip())

# Reranking Subagent
# Must update this prompt

rerank_system_prompt = """
You are an assistant that receives a list of memory items and a current conversation; you must identify and return only the most important and relevant memories that help understand the conversation or move it forward, removing duplicates or very similar entries.

# Instructions
- Think carefully about what's important for the character to remember in this context. Write down your reasoning and keep it shorter than 100 words, shorter than 50 even.
- If conflicting information is presented or context is ambiguous, state it in your reasoning.
- Take special care to state and address the most recent topic touched in the conversation.
- Based on your own analysis, generate the required JSON object. Stop your answer after this.

# Format
Output a JSON list of strings with each selected memory, for example:
["First memory","Second memory"]
""".strip()

rerank_prompt_template = Template("""
# Memories
Following are all the known memories. Each memory is in a separate line.

$memories

# Conversaton
$conversation
""".strip())

class CharacterAgent(Agent):
    identifier = "character"
    description = CLASS_DESCRIPTION
    schema_fields = [
        AgentField('personality', 'string', default="You are a friendly AI having a casual conversation with the user."),
        AgentField('style', 'string', default="Use short sentences."),
        AgentField('directives', 'string', default=default_directives),
    ]
    tools = []

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

        # Constant system prompt for developing purposes 
        self.personality = self.config['personality']
        self.style = self.config['style']
        self.directives = self.config['directives']
        self.debug = True

    async def prompt(self, content):
        # TODO: Advanced features
        messages = [
            LlamaMessage("system", self.system_prompt),
            LlamaMessage("user", content)
        ]
        response = await replicate_completion(self.app, messages)
        return response

    async def received_message(self, message):

        # TODO: before retrieving the chat context, acquire a lock for it

        chat_view = await self.app.get_chat_view(message.chat_id)
        current_thread = chat_view.get_thread(message.thread_id)

        thread_context = self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id)
        memories = []

        if thread_context and thread_context['thread_id'] == message.thread_id:
            if 'topic' in thread_context:
                memories = await self.retrieve_topic_memories(thread_context['topic'])
        else:
            thread_context = {'thread_id': message.thread_id}

        llama_messages = [LlamaMessage("system", controller_system_prompt.substitute(
            name=self.name,

            personality=self.personality,
            style=self.style,
            directives=self.directives,

            memories="\n".join(memories),
        ))]

        for msg in current_thread.messages:
            # Only count messages and notes
            if msg.kind != 'message':
                continue

            if msg.is_from_user():
                role = "user"
                content = msg.content
            elif msg.is_from_agent(self.id):
                role = "assistant"
                content = "<outwards>\n" + msg.content

                if 'last_inwards' in thread_context and thread_context.get('last_message_id_sent') == msg.id:
                    content += "\n\n<inwards>\n" + thread_context['last_inwards']
            else:
                continue
            llama_messages.append(LlamaMessage(role, content))

        reminder_message = controller_reminder_prompt.substitute(directives=self.directives)
        reminder_used = False
        # Traverse messages from last to first, stopping at `i=2` before the opening "system/user" messages
        for i in range(len(llama_messages) - 1, 1, -1):
            if llama_messages[i].role == 'user':
                # Insert reminder message before the last user message
                llama_messages.insert(i, LlamaMessage("system", reminder_message))
                reminder_used = True
                break
        
        if not reminder_used:
            # assert llama_messages[0].role == "system"
            llama_messages[0].content += "\n" + reminder_message

        # TODO: Inject system prompt before the last user message
        # If not any user message, combine with first system prompt

        controller_response = await replicate_completion(self.app, llama_messages)

        if self.debug:
            await self.app.send_chat_message(
                chat_id=message.chat_id,
                sender=self.id,
                content=controller_response,
                thread_id=current_thread.id,
                kind='log',
                subkind='debug',
            )

        sections = parse_sections(controller_response)

        if 'outwards' in sections:
            send_response = await self.app.send_chat_message(
                chat_id=message.chat_id,
                sender=self.id,
                content=sections['outwards'],
                thread_id=current_thread.id,
            )
            thread_context['last_message_id_sent'] = send_response['id']

        if 'inwards' in sections:
            thread_context['last_inwards'] = sections['inwards']

        if 'state_update' in sections:
            state_update = parse_json_object(sections['state_update'])

            if 'topic' in state_update:
                thread_context['topic'] = state_update['topic']
            if 'topic' in state_update:
                thread_context['focus'] = state_update['topic']

            if 'new_memories' in state_update:
                print(state_update['new_memories'])
        
        self.app.save_agent_resource(self.id, 'last-thread-context', message.chat_id, thread_context)


    # Encode a thread's messages as a string to pass to the agent
    def serialize_thread_history(self, thread, ignore_message_ids=[]):
        return self.serialize_message_list([m for m in thread.messages if m.id not in ignore_message_ids])

    def serialize_message_list(self, messages):
        def sender_name(message):
            sender = message.sender
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
            is_self = is_digit and str(sender) == str(self.id)
            return self.name if is_self else sender
        
        return "\n".join(f"{sender_name(message)}: {message.content}" for message in messages if message.kind == 'message')
    
    def get_memory_collection(self):
        collection_name = f"agent_{self.id}_memory"
        return collection_name
    
    async def retrieve_topic_memories(self, topic, focus=None):
        client = self.app.qdrant_client
        collection = self.get_memory_collection()

        if not client.collection_exists(collection):
            return []
        
        query_result = client.query(
            collection,
            query_text=topic,
            limit=20
        )
        if focus:
            # TODO: deduplicate
            query_result += client.query(
                collection,
                query_text=topic,
                limit=20
            )
        # TODO:
        # initial importance-based ranking
        memories = [it.document for it in query_result]

        print("rerank memories")
        rerank_response = await replicate_completion(self.app, [
            LlamaMessage("system", rerank_system_prompt),
            LlamaMessage("user", rerank_prompt_template.substitute(
                conversation=chat_history,
                memories="\n".join(memories)
            ))
        ])

        match = re.search(r'\[.*?\]', rerank_response, re.DOTALL)
        if match:
            json_string = match.group(0)
            try:
                memories = json.loads(json_string)
            except json.JSONDecodeError:
                print("Bad JSON data in response...")
                print(rerank_response)
                return []
        else:
            print("Compress format not found in response...")
            print(rerank_response)
            return []
        
        return memories
    
    async def store_new_memories(self, memories):
        if not memories:
            return

        self.app.qdrant_client.add(
            collection,
            documents=[it['content'] for it in memories],
            metadata=[{'importance': it['importance'], 'label': it['label']} for it in memories],
        )
        print(f"Added {len(memories)} memories for agent {self.id}")

def parse_json_object(text, name='object-text'):
    match = re.search(r'\{.*\}', text, re.DOTALL)
    if match:
        json_string = match.group(0)
        try:
            return json.loads(json_string)
        except json.JSONDecodeError:
            print("Bad JSON data in " + name)
            print(json_string)
            return {}
    else:
        print("No JSON data in " + name)
        print(text)
        return {}

def parse_json_array(text, name='array-text'):
    match = re.search(r'\[.*?\]', text, re.DOTALL)
    if match:
        json_string = match.group(0)
        try:
            return json.loads(json_string)
        except json.JSONDecodeError:
            print("Bad JSON data in " + name)
            print(text)
            return []
    else:
        print("No JSON data in " + name)
        print(text)
        return []

def parse_sections(text):
    """
    Parses a string into sections based on tag lines in the form <section_name>.

    Args:
        text (str): The input string with sections.

    Returns:
        dict: A dictionary where keys are section names and values are section contents.
    """
    # Regular expression to find section headers
    pattern = re.compile(r'<([^<>]+)>')
    
    # Find all section headers and their start positions
    matches = list(pattern.finditer(text))
    
    sections = {}
    for i, match in enumerate(matches):
        section_name = match.group(1)
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        section_content = text[start:end].strip()
        sections[section_name] = section_content
    
    return sections