import json
import logging
import os
import re
from typing import List, Optional, Tuple

from core.processors.sub.llm.BaseLLMClient import BaseLLMClient, LLMResponse
from core.processor import Processor
from core.task import Task
from core.loop import Loop
from i18n.translations import t

SYSTEM_PROMPT_ZH = """你是 PETP 执行流程生成助手。你可以：
1. 根据用户的自然语言描述，生成 PETP Execution 的 Task 序列
2. 回答用户关于可用 Processor 的咨询问题（推荐哪些 Processor、解释用法等）

请用中文回复。

## 可用 Processor 列表：

{processor_context}

## 输出规则：

当用户询问建议、咨询问题时，直接用自然语言回复即可（无需 JSON 格式）。

当用户要求生成或修改流程时，使用以下 JSON 格式：

- 如果信息不足以生成完整流程，用 JSON 格式追问：
  {{"action": "ask", "questions": ["问题1", "问题2"]}}

- 如果信息充足，生成 Task 序列：
  {{"action": "generate", "tasks": [
    {{"type": "INITIAL_PARAMS", "input": {{"key": "value"}}}},
    {{"type": "GO_TO_PAGE", "input": {{"url": "{{url}}"}}}},
    ...
  ], "loops": []}}

- 用户要求修改时，用增量修改格式：
  {{"action": "modify", "operations": [
    {{"op": "insert", "after": 3, "tasks": [{{"type": "WAIT_SECONDS", "input": {{"wait_seconds": "3"}}}}]}},
    {{"op": "delete", "index": 5}},
    {{"op": "replace", "index": 2, "task": {{"type": "...", "input": {{...}}}}}}
  ]}}

- 参数值中使用 {{variable_name}} 引用 data_chain 中的变量
- 密码类参数使用占位符 "__ENCRYPTED__"，后续由用户在 GUI 中加密填写
- loops 格式：{{"task_start": 5, "task_end": 10, "loop_key": "collection_key", "loop_times": "0", "loop_index_key": "loop_idx", "item_key": "loop_item"}}
- task index 从 1 开始计数
"""

SYSTEM_PROMPT_EN = """You are a PETP Execution generator assistant. You can:
1. Generate Task sequences based on natural language descriptions
2. Answer questions about available Processors (recommend which to use, explain usage, etc.)

Reply in English.

## Available Processors:

{processor_context}

## Output Rules:

When the user asks for advice or has questions, reply in plain natural language (no JSON needed).

When the user asks to generate or modify a flow, use these JSON formats:

- If you need more information, ask in JSON format:
  {{"action": "ask", "questions": ["question1", "question2"]}}

- If you have enough information, generate the Task sequence:
  {{"action": "generate", "tasks": [
    {{"type": "INITIAL_PARAMS", "input": {{"key": "value"}}}},
    {{"type": "GO_TO_PAGE", "input": {{"url": "{{url}}"}}}},
    ...
  ], "loops": []}}

- For modifications, use incremental format:
  {{"action": "modify", "operations": [
    {{"op": "insert", "after": 3, "tasks": [{{"type": "WAIT_SECONDS", "input": {{"wait_seconds": "3"}}}}]}},
    {{"op": "delete", "index": 5}},
    {{"op": "replace", "index": 2, "task": {{"type": "...", "input": {{...}}}}}}
  ]}}

- Use {{variable_name}} in parameter values to reference data_chain variables
- Use "__ENCRYPTED__" as placeholder for password values
- loops format: {{"task_start": 5, "task_end": 10, "loop_key": "collection_key", "loop_times": "0", "loop_index_key": "loop_idx", "item_key": "loop_item"}}
- task index starts from 1
"""


def resolve_api_key(raw: str) -> str:
    if not raw:
        return ''
    m = re.fullmatch(r'\$\{(.+)\}', raw.strip())
    if m:
        return os.environ.get(m.group(1), '')
    return raw


class ExecutionGenerator:

    def __init__(self, processors_or_categories: List[str], locale: str):
        self._filter = set(processors_or_categories) if processors_or_categories else None
        self._locale = locale
        self._messages = []
        self._total_prompt_tokens = 0
        self._total_completion_tokens = 0
        self._client: Optional[BaseLLMClient] = None
        self._model = ''
        self._available_processors = set(Processor.get_processors())

    def init_client(self, provider: str, api_key: str, base_url: str, model: str):
        from core.processors.AI_LLM_SETUPProcessor import AI_LLM_SETUPProcessor
        defaults = AI_LLM_SETUPProcessor.PROVIDER_DEFAULTS.get(provider, {})

        if not defaults:
            missing = []
            if not api_key:
                missing.append('ai_api_key')
            if not base_url:
                missing.append('ai_base_url')
            if not model:
                missing.append('ai_model')
            if missing:
                supported = ', '.join(AI_LLM_SETUPProcessor.PROVIDER_DEFAULTS.keys())
                raise ValueError(
                    t("ai_gen_unknown_provider", provider=provider,
                      missing=', '.join(missing), supported=supported)
                )

        if api_key and not api_key.startswith('${'):
            resolved_key = api_key
        else:
            if not api_key and defaults.get('api_key_env'):
                api_key = f"${{{defaults['api_key_env']}}}"
            resolved_key = resolve_api_key(api_key)

        if not base_url:
            base_url = defaults.get('base_url', '')
        if not model:
            model = defaults.get('model', '')

        kwargs = {'api_key': resolved_key}
        if base_url:
            kwargs['base_url'] = base_url
        self._client = BaseLLMClient.get_client_by_provider(provider, **kwargs)
        self._model = model
        self._messages = []

    def update_filter(self, processors: List[str]):
        self._filter = set(processors) if processors else None

    def validate_connection(self) -> bool:
        if self._client is None:
            return False
        messages = [{"role": "user", "content": "hi"}]
        response = self._client.chat(messages=messages, model=self._model, temperature=0)
        return bool(response.content or response.reasoning_content)

    def chat(self, user_message: str, current_tasks: Optional[List[Task]] = None) -> dict:
        if self._client is None:
            return {"action": "text", "content": t("ai_gen_no_config")}

        if not self._messages:
            self._messages = [{"role": "system", "content": self._build_system_prompt()}]

        if current_tasks and len(current_tasks) > 0:
            context = self._build_context_message(current_tasks)
            self._messages.append({"role": "user", "content": f"{context}\n\n{user_message}"})
        else:
            self._messages.append({"role": "user", "content": user_message})

        try:
            response: LLMResponse = self._client.chat(
                messages=self._messages,
                model=self._model,
                temperature=0.7,
            )
            self._total_prompt_tokens += response.prompt_tokens
            self._total_completion_tokens += response.completion_tokens

            answer = response.content or response.reasoning_content or ''
            self._messages.append({"role": "assistant", "content": answer})

            return self._parse_response(answer)

        except Exception as e:
            logging.error(f'AI Generator LLM error: {e}')
            return {"action": "text", "content": str(e)}

    def get_token_usage(self) -> Tuple[int, int]:
        return self._total_prompt_tokens, self._total_completion_tokens

    def build_tasks(self, task_dicts: list) -> List[Task]:
        tasks = []
        for td in task_dicts:
            t_type = td.get('type', '')
            t_input = td.get('input', {})
            if isinstance(t_input, dict):
                t_input = json.dumps(t_input, ensure_ascii=False)
            tasks.append(Task(t_type, t_input, False))
        return tasks

    def build_loops(self, loop_dicts: list) -> List[Loop]:
        import time
        loops = []
        for ld in loop_dicts:
            loop_code = f"loop-{time.strftime('%Y%m%d%H%M%S')}"
            loop_attr = json.dumps(ld, ensure_ascii=False)
            loops.append(Loop(loop_code, loop_attr))
        return loops

    def apply_modifications(self, current_tasks: List[Task], operations: list) -> List[Task]:
        tasks = list(current_tasks)
        offset = 0
        for op_dict in operations:
            op = op_dict.get('op')
            if op == 'insert':
                after_idx = op_dict.get('after', 0) - 1 + offset
                new_tasks = self.build_tasks(op_dict.get('tasks', []))
                for i, nt in enumerate(new_tasks):
                    tasks.insert(after_idx + 1 + i, nt)
                offset += len(new_tasks)
            elif op == 'delete':
                idx = op_dict.get('index', 0) - 1 + offset
                if 0 <= idx < len(tasks):
                    tasks.pop(idx)
                    offset -= 1
            elif op == 'replace':
                idx = op_dict.get('index', 0) - 1 + offset
                if 0 <= idx < len(tasks):
                    td = op_dict.get('task', {})
                    t_type = td.get('type', '')
                    t_input = td.get('input', {})
                    if isinstance(t_input, dict):
                        t_input = json.dumps(t_input, ensure_ascii=False)
                    tasks[idx] = Task(t_type, t_input, False)
        return tasks

    def validate_tasks(self, tasks: List[Task]) -> List[str]:
        warnings = []
        for i, task in enumerate(tasks):
            if task.type not in self._available_processors:
                warnings.append(t("ai_gen_invalid_type", ptype=task.type))
        return warnings

    def generate_mcp_desc(self, tasks: List[Task]) -> Optional[str]:
        if self._client is None:
            return None

        task_summary = "\n".join(
            f"{i+1}. {tk.type}: {tk.input}" for i, tk in enumerate(tasks)
        )

        initial_params = {}
        for tk in tasks:
            if tk.type == 'INITIAL_PARAMS':
                try:
                    initial_params = json.loads(tk.input)
                    initial_params.pop('skipped', None)
                except (json.JSONDecodeError, TypeError):
                    pass
                break

        output_keys = []
        for tk in tasks:
            if tk.type == 'HTTP_RESPONSE_KEY':
                try:
                    d = json.loads(tk.input)
                    if 'data_key' in d:
                        output_keys.append(d['data_key'])
                except (json.JSONDecodeError, TypeError):
                    pass
            elif tk.type in ('WRITE_TO_EXCEL', 'DATA_COLLECT', 'DATA_CONVERT'):
                try:
                    d = json.loads(tk.input)
                    for k in ('data_key', 'target_key', 'value_key'):
                        if k in d:
                            output_keys.append(d[k])
                            break
                except (json.JSONDecodeError, TypeError):
                    pass

        locale_hint = "请用中文生成 desc 描述。" if self._locale == 'zh' else "Generate desc in English."

        prompt = f"""{locale_hint}
Generate a mcp_desc JSON for publishing this PETP Execution as an MCP Tool.

## Task list:
{task_summary}

## Extracted info:
- INITIAL_PARAMS keys (tool inputs): {json.dumps(initial_params, ensure_ascii=False) if initial_params else 'None found'}
- Output keys found: {output_keys if output_keys else 'None found'}

## Rules:
1. "desc": Write a clear, concise description (1-2 sentences) of what this tool does. It should be understandable by an AI agent deciding whether to call this tool.
2. "inputSchema": For each INITIAL_PARAMS key, infer the type (string/number/boolean), write a helpful title and description. Mark keys without default values as required.
3. "outParams": List the output keys that the caller can expect in the result.

## Required JSON structure:
{{"desc": "...", "inputSchema": {{"type": "object", "properties": {{"param1": {{"type": "string", "title": "Param 1", "description": "What this param is for"}}}}, "required": ["param1"]}}, "outParams": ["output_key"]}}

Return ONLY the JSON, no markdown wrapping."""

        messages = [{"role": "user", "content": prompt}]
        try:
            response = self._client.chat(messages=messages, model=self._model, temperature=0.3)
            self._total_prompt_tokens += response.prompt_tokens
            self._total_completion_tokens += response.completion_tokens
            return response.content or ''
        except Exception as e:
            logging.error(f'AI MCP desc generation error: {e}')
            return None

    def _build_system_prompt(self) -> str:
        template = SYSTEM_PROMPT_ZH if self._locale == 'zh' else SYSTEM_PROMPT_EN
        context = self._load_processor_context()
        return template.format(processor_context=context)

    def _load_processor_context(self) -> str:
        from i18n.translations import get_localized_desc
        context_parts = []
        for ptype in Processor.get_processors():
            if self._filter is not None and ptype not in self._filter:
                if Processor.get_processor_by_type(ptype).get_category() != 'General':
                    continue
            try:
                clazz = Processor.get_processor_by_type(ptype)
                desc = get_localized_desc(clazz, self._locale)
                tpl = clazz.TPL if hasattr(clazz, 'TPL') else '{}'
                context_parts.append(f"### {ptype}\nTPL: {tpl}\n{desc}")
            except Exception:
                continue
        return "\n\n".join(context_parts)

    def _build_context_message(self, tasks: List[Task]) -> str:
        lines = []
        for i, task in enumerate(tasks):
            lines.append(f"{i+1}. {task.type}: {task.input}")
        header = "当前 Task 列表：" if self._locale == 'zh' else "Current Task List:"
        return header + "\n" + "\n".join(lines)

    def _parse_response(self, text: str) -> dict:
        text = text.strip()
        if text.startswith('```'):
            text = re.sub(r'^```\w*\n?', '', text)
            text = re.sub(r'\n?```$', '', text)
            text = text.strip()

        try:
            parsed = json.loads(text)
            action = parsed.get('action', '')
            if action in ('ask', 'generate', 'modify'):
                return parsed
            return {"action": "text", "content": text}
        except json.JSONDecodeError:
            return {"action": "text", "content": text}
