import time
import requests
import json
import argparse
import io
import logging
import types
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.json import JSON


# Local imports
from logging_config import setup_logging
import utils
import prompts as prompts

logger = logging.getLogger(__name__)
console = Console()

# Configuration
DEFAULT_AGENT_API = "http://localhost:8080/v1"
DEFAULT_REPL_API = "http://localhost:8090/v1"

DEFAULT_CONTEXT_FILE = "context.txt"
DEFAULT_TASK_FILE = "task.txt"
MAX_REPL_STEPS = 20
MAX_VIRTUAL_CONTEXT_RATIO = 0.85


class LlamaClient:
    def __init__(self, base_url, name="LlamaClient"):
        self.base_url = base_url.rstrip("/")
        self.name = name
        self.model = None
        self.n_ctx = 4096
        self._get_model_info()
        self.max_input_tokens = int(self.n_ctx * MAX_VIRTUAL_CONTEXT_RATIO)
        self.color = self._determine_color()
        if debug: logger.debug(f"Connected to {name} ({base_url}). Model: {self.model}. Context: {self.n_ctx}. Max Input: {self.max_input_tokens}")

    def _determine_color(self):
        if "8080" in self.base_url:
            return "dodger_blue1"
        elif "8090" in self.base_url:
            return "dodger_blue3"
        else:
            return "cyan1"

    def _get_model_info(self):
        try:
            resp = requests.get(f"{self.base_url}/models")
            resp.raise_for_status()
            data = resp.json()
            model_data = data.get("data", [])
            if model_data:
                model = model_data[0]
                self.model = model.get("id", "default")
                meta = model.get("meta", {})
                self.n_ctx = meta.get("n_ctx", 4096)
            else:
                self.model = "default"
        except Exception as e:
            logger.error(f"[{self.name}] Failed to get model info: {e}. Defaulting.")
            self.model = "default"

    def count_tokens(self, messages):
        try:
            resp = requests.post(
                f"{self.base_url}/chat/completions/input_tokens",
                json={"model": self.model, "messages": messages},
                timeout=30.0,
            )
            if resp.status_code == 200:
                return resp.json().get("input_tokens", 0)
        except Exception as e:
            logger.debug(f"[{self.name}] Token count failed: {e}. Using estimate.")
        return sum(len(json.dumps(m)) // 4 + 4 for m in messages)

    def count_text_tokens(self, text):
        return self.count_tokens([{"role": "user", "content": text}])

    def completion(self, messages, schema=None, temperature=0.1):
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
        }
        if schema:
            payload["response_format"] = {
                "type": "json_schema",
                "json_schema": {"name": "response", "schema": schema}
            }
        if debug:
            last_content = messages[-1].get("content", "") if messages else ""
            console.print(Panel(
                last_content[-500:] if len(last_content) > 500 else last_content,
                title=f"Last message to {self.name}",
                title_align="left",
                border_style=self.color
            ))
        try:
            resp = requests.post(f"{self.base_url}/chat/completions", json=payload)
            resp.raise_for_status()
            content = resp.json()["choices"][0]["message"]["content"].strip()
            if debug:
                console.print(Panel(
                    JSON.from_data(content),
                    title=f"{self.name} Response",
                    title_align="left",
                    border_style=self.color
                ))
            return content
        except Exception as e:
            logger.error(f"[{self.name}] Error calling LLM: {e}")
            return f"Error: {e}"

class AgentTools:
    def __init__(self, repl_client: LlamaClient, data_content: str):
        self.client = repl_client
        self.RAW_CORPUS = data_content

    def llm_query(self, content_chunk, query):
        if content_chunk == "RAW_CORPUS":
            return "ERROR: You passed the string 'RAW_CORPUS' You must pass the CONTENT of the variable (e.g., `chunk = RAW_CORPUS[:1000]`, then `llm_query(chunk, ...)`)."

        estimated_tokens = len(content_chunk) // 3
        if estimated_tokens > (self.client.n_ctx * 2):
            return f"ERROR: Chunk is massively too large (approx {estimated_tokens} tokens). Slice strictly."

        chunk_tokens = self.client.count_text_tokens(content_chunk)
        query_tokens = self.client.count_text_tokens(query)
        total = chunk_tokens + query_tokens + 150

        if debug: logger.debug(f"[Sub-LLM] Processing Query with {total} tokens.")

        if total > self.client.n_ctx:
            msg = f"ERROR: Chunk too large ({chunk_tokens} tokens). Limit is {self.client.n_ctx}. Slice smaller."
            logger.warning(msg)
            return msg

        sub_messages = [
            {"role": "system", "content": (
                "You are a strict reading assistant. "
                "Answer the question based ONLY on the provided Context. "
                "Do not use outside training data. "
                "If the answer is not in the text, say 'NULL'."
            )},
            {"role": "user", "content": f"Context:\n{content_chunk}\n\nQuestion: {query}"}
        ]
        results = self.client.completion(sub_messages)
        result_tokens = self.client.count_text_tokens(results)
        if debug: logger.debug(f"[Sub-LLM] Responded with {result_tokens} tokens.")
        return results

class AgentOutputBuffer:
    def __init__(self, max_total_chars=20000, max_len_per_print=1009):
        self._io = io.StringIO()
        self.max_total_chars = max_total_chars
        self.max_len_per_print = max_len_per_print
        self.current_chars = 0
        self.global_truncated = False

    def custom_print(self, *args, **kwargs):
        temp_io = io.StringIO()
        print(*args, file=temp_io, **kwargs)
        text = temp_io.getvalue()

        if len(text) > self.max_len_per_print:
            truncated_text = text[:self.max_len_per_print]
            text = (
                f"{truncated_text}\n"
                f"... [LINE TRUNCATED: Output exceeded {self.max_len_per_print-9} chars. "
                f"Use slicing or llm_query() to inspect data.] ...\n"
            )

        if self.current_chars + len(text) > self.max_total_chars:
            remaining = self.max_total_chars - self.current_chars
            if remaining > 0:
                self._io.write(text[:remaining])

            if not self.global_truncated:
                self._io.write(f"\n... [SYSTEM HALT: Total output limit ({self.max_total_chars}) reached] ...\n")
                self.global_truncated = True

            self.current_chars += len(text)
        else:
            self._io.write(text)
            self.current_chars += len(text)

    def read_and_clear(self):
        value = self._io.getvalue()
        self._io = io.StringIO()
        self.current_chars = 0
        self.global_truncated = False
        return value

def run_agent(agent_client, repl_client, context_text, task_text):
    tools = AgentTools(repl_client, context_text)

    agent_schema = {
        "type": "object",
        "properties": {
            "thought": {"type": "string", "description": "Reasoning about current state and what to do next."},
            "action": {"type": "string", "enum": ["execute_python", "final_answer"]},
            "content": {"type": "string", "description": "Python code or Final Answer text."}
        },
        "required": ["thought", "action", "content"]
    }

    out_buffer = AgentOutputBuffer()

    trace_filepath = utils.init_trace_file(debug)

    exec_env = {
        "RAW_CORPUS": tools.RAW_CORPUS,
        "llm_query": tools.llm_query,
        "re": __import__("re"),
        "math": __import__("math"),
        "json": __import__("json"),
        "collections": __import__("collections"),
        "statistics": __import__("statistics"),
        "random": __import__("random"),
        "datetime": __import__("datetime"),
        "difflib": __import__("difflib"),
        "string": __import__("string"),

        "print": out_buffer.custom_print
    }

    system_instruction = prompts.get_system_prompt()

    messages = [
        {"role": "system", "content": system_instruction},
        {"role": "user", "content": f"USER TASK: {task_text}"}
    ]

    step = 0
    while step < MAX_REPL_STEPS:
        step += 1
        if debug: logger.debug(f"Step {step} of {MAX_REPL_STEPS}")

        modules = []
        functions = []
        variables = []
        ACTIVE_VAR_SNIPPET_LEN = 100

        for name, val in exec_env.items():
            if name.startswith("__"): continue
            if name == "print": continue

            if isinstance(val, types.ModuleType):
                modules.append(name)
            elif callable(val):
                functions.append(name)
            else:
                type_name = type(val).__name__
                s_val = str(val)
                snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_val
                variables.append(f"{name} ({type_name}): {snippet}")

        dynamic_state_msg = (
            f"[SYSTEM STATE REMINDER]\n"
            f"Current Step: {step}/{MAX_REPL_STEPS}\n"
            f"Available Libraries: {', '.join(modules)}\n"
            f"Available Tools: {', '.join(functions)}\n"
            f"Active Variables:\n" + ("\n".join([f"  - {v}" for v in variables]) if variables else "  (None)") + "\n---"
        )

        inference_messages = messages.copy()
        inference_messages.append({"role": "user", "content": dynamic_state_msg})

        usage = agent_client.count_tokens(inference_messages)
        if debug: logger.debug(f"Context Usage: {usage} / {agent_client.max_input_tokens}")

        if usage > agent_client.max_input_tokens:
             if debug: logger.warning("Context limit exceeded. Triggering History Compression.")

             messages = utils.compress_history(debug, agent_client, messages, keep_last_pairs=2)

             new_usage = agent_client.count_tokens(inference_messages)
             if debug: logger.debug(f"Context Usage after compression: {new_usage}")

             if new_usage > agent_client.max_input_tokens:
                 logger.error("Compression insufficient. Forcing hard truncation.")
                 messages.pop(2)

        response_text = agent_client.completion(inference_messages, schema=agent_schema, temperature=0.5)

        try:
            response_json = json.loads(response_text)
        except json.JSONDecodeError:
            logger.error("JSON Parse Error")
            messages.append({"role": "user", "content": "System: Invalid JSON returned. Please retry."})
            continue

        thought = response_json.get("thought", "")
        action = response_json.get("action", "")
        content = response_json.get("content", "")

        if action == "execute_python" and content:
            content = utils.safeguard_and_repair(debug, agent_client, messages, agent_schema, content)

        if debug:
            console.print(Panel(
                    f"[italic]{thought}[/italic]",
                    title="Agent Thought",
                    title_align="left",
                    border_style="magenta"
                ))

        messages.append({"role": "assistant", "content": json.dumps(response_json, indent=2, ensure_ascii=False)})

        if action == "final_answer":
            if debug: logger.debug(f"Raw Agent Output: {content}")

            final_report = utils.generate_final_report(debug, agent_client, task_text, content)

            final_report_md = Markdown(final_report)
            print("\n\n")
            console.print(final_report_md)
            print("\n")
            break

        elif action == "execute_python":
            if debug and content != response_json.get("content"):
                console.print(Panel(content, title="Executing Code via Safeguard", title_align="left", border_style="cyan"))
            elif debug and content == response_json.get("content"):
                console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow"))

            observation = ""
            try:
                out_buffer.read_and_clear()
                exec(content, exec_env)
                observation = out_buffer.read_and_clear()

                if not observation:
                    observation = "Code executed successfully (no output)."
            except Exception as e:
                observation = f"Python Error: {e}"
                logger.error(f"Code Execution Error: {e}")

            if debug:
                console.print(Panel(
                    f"{observation.strip()}",
                    title="Observation",
                    title_align="left",
                    border_style="dark_green"
                ))
            messages.append({"role": "user", "content": f"Observation:\n{observation}"})

        else:
            messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})

        utils.save_agent_trace(trace_filepath, messages)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="""Edge Recursive Language Model

    A sophisticated data extraction and analysis tool that mimics the process of a human data scientist, carefully exploring and structuring a large dataset before performing targeted queries.""")
    parser.add_argument("--context", default=DEFAULT_CONTEXT_FILE, help="Path to text file to process")
    parser.add_argument("--task", default=DEFAULT_TASK_FILE, help="Path to task instruction file")
    parser.add_argument("--override_task", help="Direct string override for the task")
    parser.add_argument("--agent_api", default=DEFAULT_AGENT_API, help="URL for the Main Agent LLM")
    parser.add_argument("--repl_api", default=DEFAULT_REPL_API, help="URL for the Sub-call/REPL LLM")
    parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging")

    args = parser.parse_args()
    debug = args.debug
    log_level=logging.DEBUG if debug else logging.INFO
    setup_logging(level=log_level, debug=debug)

    if debug: logger.info("Starting EdgeRLM...")
    context_content = utils.load_file(args.context)
    if debug: logger.debug(f"Loaded Context: {len(context_content)} characters.")
    task_content = args.override_task if args.override_task else utils.load_file(args.task)

    agent_client = LlamaClient(args.agent_api, "Agent")
    repl_client = LlamaClient(args.repl_api, "REPL")

    run_agent(agent_client, repl_client, context_content, task_content)