# agent.py
from pathlib import Path
from typing import Optional, Any
from tiktoken import encoding_for_model, get_encoding
from .schema import AgentConfig
from .llm_client import create_llm_client, LLMClient
import json
from typing import Optional, Callable, Any, Dict, List, Union
import importlib.util
import sys
import logging
import os
from datetime import datetime
import hashlib

logger = logging.getLogger(__name__)

for noisy in ["chromadb", "posthog", "urllib3", "backoff", "httpx"]:
    logging.getLogger(noisy).setLevel(logging.CRITICAL)


class Agent:
    def __init__(self, base_agents_dir: Path, config: AgentConfig, api_key: str, base_url: Optional[str] = None) -> None:
        self.config = config
        self.client: LLMClient = create_llm_client(
            model=config.model,
            api_key=api_key,
            base_url=base_url,
            provider=config.provider,
        )
        self.base_agent_dir = base_agents_dir / 'agents' / config.name

        # load prompts
        self._instruction_prompt = self._load_prompt(config.instruction_path)
        self._think_prompt = self._load_prompt(config.think_path)
        self._example_prompt = self._load_prompt(config.example_path)
        self._next_model_rule_prompt = self._load_prompt(config.next_model_rule_path)
        self._return_format_prompt = self._load_prompt(config.return_path)

        # preprocess function (runs on input BEFORE sending to LLM)
        if config.preprocess_path and config.preprocess_function_name:
            self._preprocess_function = self._function_loader(
                self.base_agent_dir / config.preprocess_path,
                config.preprocess_function_name
            )
        else:
            self._preprocess_function = None

        # postprocess function (runs on LLM output BEFORE returning)
        if config.postprocess_path and config.postprocess_function_name:
            self._postprocess_function = self._function_loader(
                self.base_agent_dir / config.postprocess_path,
                config.postprocess_function_name
            )
        else:
            self._postprocess_function = None

        # legacy return format function (kept for backward compat)
        if config.return_format_function_path and config.return_format_function_name:
            self._output_function = self._function_loader(
                self.base_agent_dir / config.return_format_function_path,
                config.return_format_function_name
            )
        else:
            self._output_function = self._default_output

        self._json_schema = self._load_prompt(config.json_schema_path)
        self._return_format = config.return_format

        # hyper params
        self._temperature = config.temperature
        self._max_tokens = config.max_tokens

        self.full_system_prompt = self._build_system_prompt()

        self.last_user_prompt: Optional[Any] = None
        self.last_rag_context: Optional[Any] = None
        self.least_answer: Optional[Any] = None

        # Setup per-agent log directories
        self._io_log_dir = self.base_agent_dir / "logs" / "input_output"
        self._audit_log_dir = self.base_agent_dir / "logs" / "audit"
        self._io_log_dir.mkdir(parents=True, exist_ok=True)
        self._audit_log_dir.mkdir(parents=True, exist_ok=True)

        # Audit logger (tokens, errors) — one file per date
        self._audit_logger = self._create_logger(
            f"audit_{config.name}",
            self._audit_log_dir,
        )

        # Log system prompt if changed
        self._log_system_prompt()

    # ── File Loading ─────────────────────────────────────────

    def _load_prompt(self, relative_path: Optional[Path]) -> Any:
        """Load a file relative to the agent's directory."""
        if not relative_path:
            return None
        return self._load_file(self.base_agent_dir / relative_path)

    def _load_file(self, path: Optional[Path]) -> Any:
        if not path:
            return None

        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")

        suffix = path.suffix.lower()
        if suffix == ".json":
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        elif suffix in [".md", ".txt", ".text"]:
            with open(path, "r", encoding="utf-8") as f:
                return f.read()
        else:
            raise ValueError(f"Unsupported file type: {suffix}")

    # ── System Prompt ────────────────────────────────────────

    def _build_system_prompt(self) -> str:
        parts: list[str] = []

        if self._instruction_prompt:
            parts.append(str(self._instruction_prompt))

        if self._think_prompt:
            parts.append("### THINKING GUIDELINES\n" + str(self._think_prompt))

        if self._example_prompt:
            parts.append("### EXAMPLE\n" + str(self._example_prompt))

        if self._next_model_rule_prompt:
            parts.append("### NEXT MODEL HANDOFF RULES\n" + str(self._next_model_rule_prompt))

        if self._return_format_prompt:
            parts.append("### OUTPUT FORMAT\n" + str(self._return_format_prompt))

        return "\n\n".join(parts)

    # ── Output Processing ────────────────────────────────────

    def _default_output(self, output: str) -> Any:
        try:
            return json.loads(output)
        except json.JSONDecodeError:
            return output

    def _function_loader(
        self,
        module_path: Union[str, Path],
        func_name: Optional[str],
    ) -> Callable[..., Any]:
        if not func_name:
            raise ValueError("func_name is required for return_format_function_name")

        if isinstance(module_path, Path):
            module_path = module_path.resolve()
            spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
            if spec is None or spec.loader is None:
                raise ImportError(f"Cannot load module from file '{module_path}'")

            module = importlib.util.module_from_spec(spec)
            sys.modules[module.__name__] = module
            spec.loader.exec_module(module)
        else:
            module = importlib.import_module(module_path)

        func = getattr(module, func_name, None)
        if func is None:
            raise AttributeError(
                f"Module '{module.__name__}' has no attribute '{func_name}'"
            )
        if not callable(func):
            raise TypeError(
                f"Attribute '{func_name}' in module '{module.__name__}' is not callable"
            )

        return func

    # ── Run ──────────────────────────────────────────────────

    def build_messages(self, user_prompt: str) -> list[dict[str, str]]:
        messages: list[dict[str, str]] = []

        if self.full_system_prompt:
            messages.append({"role": "system", "content": self.full_system_prompt})

        messages.append({"role": "user", "content": user_prompt})
        return messages

    def run(self, user_prompt: str) -> Any:
        self.last_user_prompt = user_prompt

        # ── Step 1: Preprocess (if defined) ──
        preprocess_output = None
        if self._preprocess_function:
            preprocess_output = self._preprocess_function(user_prompt)
            llm_input = str(preprocess_output)
        else:
            llm_input = user_prompt

        # ── Step 2: Call LLM ──
        messages = self.build_messages(llm_input)

        response_format = None
        if self._json_schema:
            response_format = {
                "type": "json_schema",
                "json_schema": {
                    "name": self.config.name,
                    "schema": self._json_schema,
                    "strict": True
                }
            }
        elif self._return_format == "json_object":
            response_format = {"type": "json_object"}

        llm_raw_output = self.client.chat(
            messages=messages,
            temperature=self._temperature,
            max_tokens=self._max_tokens,
            response_format=response_format,
        )

        # Parse LLM output (json decode / legacy return format function)
        llm_output = self._output_function(llm_raw_output)

        # ── Step 3: Postprocess (if defined) ──
        postprocess_output = None
        if self._postprocess_function:
            postprocess_output = self._postprocess_function(llm_output)
            final_output = postprocess_output
        else:
            final_output = llm_output

        self.least_answer = final_output

        # ── Logging ──
        self._log_io(
            input_text=user_prompt,
            preprocess_output=preprocess_output,
            llm_output=llm_output,
            postprocess_output=postprocess_output,
            llm_input=llm_input,
        )
        self._log_tokens_summary()

        return self.least_answer

    # ── Getters ──────────────────────────────────────────────

    def get_agent_config(self) -> AgentConfig:
        return self.config

    # ── Token Tracking ───────────────────────────────────────

    def _count_tokens(self, text: str) -> int:
        if text is None:
            return 0
        try:
            encoding = encoding_for_model(self.config.model)
        except KeyError:
            encoding = get_encoding("o200k_base")
        return len(encoding.encode(text))

    def get_tokens_summary(self) -> dict:
        return {
            "model_name": self.config.name,
            "model_model": self.config.model,
            "prompt_system_tokens": self._count_tokens(self.full_system_prompt),
            "least_user_prompt": self._count_tokens(self.last_user_prompt),
            "last_rag_tokens": self._count_tokens(self.last_rag_context),
            "last_full_prompt_tokens": self._count_tokens(
                str(self.full_system_prompt) + str(self.last_user_prompt) + str(self.last_rag_context)
            ),
            "last_answer_tokens": self._count_tokens(
                str(self.least_answer) if self.least_answer else ""
            )
        }

    @staticmethod
    def _create_logger(name: str, log_dir: Path) -> logging.Logger:
        """Create a logger that writes to a date-based file inside log_dir."""
        log_file = log_dir / f"{datetime.now().strftime('%Y-%m-%d')}.log"
        _logger = logging.getLogger(name)
        _logger.setLevel(logging.INFO)
        # Avoid duplicate handlers if Agent is instantiated multiple times
        if not _logger.handlers:
            handler = logging.FileHandler(log_file, encoding="utf-8")
            handler.setFormatter(logging.Formatter(
                "%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
            ))
            _logger.addHandler(handler)
        _logger.propagate = False
        return _logger

    def _log_system_prompt(self) -> None:
        """
        Log system prompt to agents_dir/{agent_name}/logs/system_prompt.json.
        Only appends a new entry if the hash changed since the last entry.
        """
        try:
            log_file = self.base_agent_dir / "logs" / "system_prompt.json"
            current_hash = hashlib.sha256(
                self.full_system_prompt.encode("utf-8")
            ).hexdigest()

            # Check if hash changed
            entries = []
            if log_file.exists():
                with open(log_file, "r", encoding="utf-8") as f:
                    entries = json.load(f)
                if entries and entries[-1].get("system_prompt_hash") == current_hash:
                    return  # No change, skip

            entry = {
                "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "system_prompt_hash": current_hash,
                "instruction_prompt": str(self._instruction_prompt) if self._instruction_prompt else None,
                "think_prompt": str(self._think_prompt) if self._think_prompt else None,
                "return_prompt": str(self._return_format_prompt) if self._return_format_prompt else None,
                "next_model_rule_prompt": str(self._next_model_rule_prompt) if self._next_model_rule_prompt else None,
                "example_prompt": str(self._example_prompt) if self._example_prompt else None,
            }

            entries.append(entry)

            with open(log_file, "w", encoding="utf-8") as f:
                json.dump(entries, f, ensure_ascii=False, indent=2)

        except Exception as e:
            logger.warning(f"Failed to log system prompt: {e}")

    def _log_io(
        self,
        input_text: str,
        preprocess_output: Any,
        llm_output: Any,
        postprocess_output: Any,
        llm_input: str,
    ) -> None:
        """Log full pipeline to agents_dir/{agent_name}/logs/input_output/{date}.json"""
        try:
            def _serialize(val: Any) -> Any:
                if val is None:
                    return None
                return val if isinstance(val, (dict, list, str, int, float, bool)) else str(val)

            log_entry = {
                "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "system_prompt_hash": hashlib.sha256(
                    self.full_system_prompt.encode("utf-8")
                ).hexdigest(),
                "input": input_text,
            }

            # Only include preprocess if it was used
            if preprocess_output is not None:
                log_entry["preprocess_output"] = _serialize(preprocess_output)

            log_entry["llm_output"] = _serialize(llm_output)

            # Only include postprocess if it was used
            if postprocess_output is not None:
                log_entry["postprocess_output"] = _serialize(postprocess_output)

            log_entry["token_input"] = self._count_tokens(llm_input)
            log_entry["token_output"] = self._count_tokens(
                str(llm_output) if llm_output else ""
            )

            log_file = self._io_log_dir / f"{datetime.now().strftime('%Y-%m-%d')}.json"

            entries = []
            if log_file.exists():
                with open(log_file, "r", encoding="utf-8") as f:
                    entries = json.load(f)

            entries.append(log_entry)

            with open(log_file, "w", encoding="utf-8") as f:
                json.dump(entries, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"Failed to log IO: {e}")

    def _log_tokens_summary(self) -> None:
        """Log token summary to agents_dir/{agent_name}/logs/audit/{date}.log"""
        try:
            summary = self.get_tokens_summary()
            self._audit_logger.info(
                f"Token usage: {json.dumps(summary, indent=2)}"
            )
        except Exception as e:
            self._audit_logger.warning(f"Failed to log token summary: {e}")