import json
import os
import time
import loguru
from loguru._logger import Logger

from typing import Optional, Literal, List, Dict, Any
from pydantic import BaseModel

logger = loguru.logger


class AgentMeta(BaseModel):
    name: str
    desc: str
    parameters: Optional[Dict[str, Any]] = None


def save_rule_config(source: str, rule_config: "RuleConfig"):
    """
    Save rule config to file
    """
    # filename: <source>_generated_<name>_<timestamp>.json
    # content: rule_config.model_dump_json()
    # path: logs/rule_configs/
    # create directory if not exists
    os.makedirs("logs/rule_configs", exist_ok=True)
    with open(
        f"logs/rule_configs/{source}_generated_{rule_config.name}_{time.strftime('%Y%m%d_%H%M%S')}.json",
        "w",
    ) as f:
        f.write(rule_config.model_dump_json())


def find_json_in_str(s: str) -> str:
    """
    Find the first JSON object in the string
    """
    start = s.find("{")
    end = s.rfind("}")
    if start != -1 and end != -1:
        return s[start : end + 1]
    return None


class RuleConfig(BaseModel):
    """
    Rule configuration

    Parameters:
        name: Rule name
        depend_ctx_key: The key of the data to be retrieved from the context
        match_condition: Condition to be satisfied, will be executed by eval,
                        for example: "力量训练" in ctx.get("intent_analysis_result.intent")

        prompt: Custom prompt for this rule (will be appended to system_prompt)
        prompt_detail: Additional prompt information for this rule, will be appended to system_prompt + prompt

        tool: The name of the tool to be executed
        tool_params: Parameters of the tools to be used
                    ```json
                    {
                        "ctx": ["user_input"], # dynamically generated by other Rule or Agent
                        "config": ["prompt", "prompt_detail"], # fixed configuration
                        "extra": {}, # additional constants
                    }
                    ```
        tool_result_target: Where the AI output of this rule should be sent to:
                            - AS_CONTEXT store in context
                            - DIRECT_RETURN return to user directly
                            - NEW_RULES generate new rules
        tool_result_key: If AS_CONTEXT, the key to store the AI output in ctx
        auto_generated: Whether this rule is automatically generated
        priority: Priority
    """

    name: str
    depend_ctx_key: List[str]
    match_condition: Optional[str] = None

    prompt: str
    prompt_detail: Optional[str] = ""

    tool: Optional[str] = None
    tool_params: Optional[Dict[str, Any]] = {}
    tool_result_target: Literal["AS_CONTEXT", "DIRECT_RETURN", "NEW_RULES"]
    tool_result_key: Optional[str] = None
    auto_generated: bool = False
    priority: int = 0

    @classmethod
    def parse_and_gen(
        cls, source: str, tool_result_full: str, save: bool = False
    ) -> List["RuleConfig"]:
        """
        Parse llm_response_full and generate new rules
        """
        try:
            # Try to parse the full JSON
            json_data = json.loads(tool_result_full, strict=False)
        except json.JSONDecodeError:
            # Try to extract JSON from the text
            try:
                json_str = find_json_in_str(tool_result_full)
                if not json_str:
                    logger.error(f"Failed to extract JSON from result: {tool_result_full[:100]}...")
                    return []

                json_data = json.loads(json_str, strict=False)
            except json.JSONDecodeError as e:
                logger.error(
                    f"Failed to parse JSON: {e}, original content: {tool_result_full[:100]}..."
                )
                return []
        result = []
        try:
            if isinstance(json_data, list):
                result = [cls(**item) for item in json_data]
            elif isinstance(json_data, dict):
                result = [cls(**json_data)]
            else:
                raise ValueError(f"Invalid JSON data: {tool_result_full}")
        except (TypeError, ValueError) as e:
            logger.error(f"Invalid JSON data: {tool_result_full}")
            return []

        if save and len(result) > 0:
            for rule_config in result:
                save_rule_config(source=source, rule_config=rule_config)

        return result


class AgentConfig(BaseModel):
    """
    Agent configuration
    """
    name: str
    meta: AgentMeta
    system_prompt: str
    rules: List[RuleConfig]
    max_concurrent_tasks: int = 1
    timeout_detection_sec: int = 180
    stream_registration_timeout: int = 1
    
