"""
Requirement Agent for Task Manager MCP Servers

Adapted from the original project's RequirementAgent for use in MCP servers.
This agent converts user ideas into structured requirement documents.
"""

import os
import json
import requests
from typing import Optional


class RequirementAgent:
    """
    需求说明Agent - 负责将用户模糊的想法扩展成结构化的需求文档。
    它的核心是通过一次交互，引导用户补充关键信息。
    """

    SYSTEM_PROMPT_TEMPLATE = """
你是资深的AI产品经理。你的工作是接收用户模糊的想法，并将其扩展成一个结构化的项目需求文档。

**你的原则:**
1.  **一次交互:** 为了最大限度地减少来回沟通，你的主要目标是引导用户在一次回应中提供所有必要信息，而不是问很多问题。
2.  **提供模板并引导输入:** 基于用户的想法，你将尽可能地填写下面的"需求文档模板"。对于信息缺失的部分，你将使用清晰的占位符 `[需用户提供]` 来向用户展示需要补充的内容。
3.  **教育用户:** 你的回应本身也是一个自然地教导用户"什么是好的需求文档"的机会。
4.  **语言一致性:** 你必须使用与用户初始想法相同的语言进行回应。如果用户使用中文，你的整个文档也必须是中文。

**需求文档模板 (Markdown 格式):**
```markdown
# 项目需求: {project_name}

## 1. 愿景与目标
[描述本项目旨在解决的核心问题以及期望达成的最终目标。]

## 2. 核心功能
- 功能 A: [详细描述此功能。]
- 功能 B: [详细描述此功能。]
- ...

## 3. 目标用户
[描述本产品的主要用户群体。例如：忙碌的职场人士、学生等。]

## 4. 技术限制与非功能性需求
- **技术栈:** [是否有任何期望或限制使用的技术？例如：Python后端, React前端。]
- **其他:** [是否有任何关于性能、安全性或其他非功能性的需求？]
```

**你的任务:**
分析用户的初始想法: `{initial_idea}`，并遵循以上模板创建一个需求文档。
对于所有信息缺失、需要用户输入的部分，你必须使用格式：`[需用户提供: 请在这里说明需要的信息]`。
"""

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        """
        从环境变量加载API密钥和基础URL。
        """
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
        self.base_url = base_url or os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
        if not self.api_key:
            raise ValueError("OPENAI_API_KEY environment variable not set.")

    def generate_requirements_document(self, user_input: str, existing_doc: Optional[str] = None) -> str:
        """
        Generates a structured requirements document based on user input.
        """
        
        system_prompt = ""
        if existing_doc:
            # 提示 AI 在现有文档上进行迭代
            system_prompt = """
你是一个资深的AI产品经理。你的工作是与用户进行对话，共同探索想法，并将讨论结果最终沉淀为一份结构化的项目需求文档。

**你的交互模式:**
你必须根据用户的输入，从以下三种响应模式中选择一种：

1.  **提问 (QUESTION):** 当用户的想法模糊、宽泛，或需要更多细节才能进行下一步时使用。
    - **触发场景:** 用户说 "我想做一个类似小红书的应用" 或 "我们来聊聊社交功能吧"。
    - **你的行为:** 不要直接生成文档。而是提出探索性的问题来引导用户。例如："这个想法很棒！小红书的功能很多，您是对它的社交分享、内容推荐还是电商部分更感兴趣呢？您的项目有什么独特的切入点吗？"
    - **输出格式:** 以 `QUESTION:::` 作为开头，后面跟上你的问题。

2.  **讨论 (DISCUSSION):** 当用户提出一个开放性问题，寻求策略、建议或知识时使用。
    - **触发场景:** 用户问 "小红书的社交平台策略是怎样的？" 或 "你对这个功能有什么建议吗？"。
    - **你的行为:** 像一个真正的产品专家一样，提供有见地的分析和建议。
    - **输出格式:** 以 `DISCUSSION:::` 作为开头，后面跟上你的见解。

3.  **更新文档 (DOC_UPDATE):** 当用户的指令明确，提供了足够的信息来创建或修改需求文档时使用。
    - **触发场景:** 用户说 "把登录方式改为手机号登录" 或提供了足够细节的新功能描述。
    - **你的行为:** 生成或更新一份完整的、结构化的Markdown需求文档。
    - **输出格式:** 以 `DOC_UPDATE:::` 作为开头，后面跟上完整的Markdown文档内容。

**上下文:**
- **现有文档:**
---
{existing_doc}
---
- **用户最新输入:** {user_input}

**重要规则:**
- 你的每一次回复都必须严格以 `QUESTION:::`、`DISCUSSION:::` 或 `DOC_UPDATE:::` 中的一个作为开头。
- 如果现有文档为空，而用户输入的信息不足以生成一份有意义的文档，你应该选择 `QUESTION` 模式。
- 语言必须与用户保持一致。

---
"""
        else:
            # 当没有现有文档时，提示AI从头开始创建，但也遵循新的交互模式
            system_prompt = """
你是一个资深的AI产品经理。你的工作是与用户进行对话，将一个模糊的想法逐步细化，并最终形成一份结构化的项目需求文档。

**你的交互模式:**
你必须根据用户的输入，从以下三种响应模式中选择一种：

1.  **提问 (QUESTION):** 当用户的想法模糊、宽泛，或需要更多细节才能生成文档时使用。
    - **触发场景:** 用户说 "我想做一个新的APP" 或 "我们来做一个电商网站吧"。
    - **你的行为:** 不要直接生成文档。而是提出探索性的问题来引导用户，收集必要信息。例如："听起来很有趣！这个APP的核心目标是什么？主要为哪些用户解决什么问题呢？"
    - **输出格式:** 以 `QUESTION:::` 作为开头，后面跟上你的问题。

2.  **讨论 (DISCUSSION):** 当用户提出一个开放性问题，寻求策略、建议或知识时使用。
    - **触发场景:** 用户问 "你对现在市场上的社交APP有什么看法？" 或 "你对这个功能有什么建议吗？"。
    - **你的行为:** 像一个真正的产品专家一样，提供有见地的分析和建议。
    - **输出格式:** 以 `DISCUSSION:::` 作为开头，后面跟上你的见解。

3.  **更新文档 (DOC_UPDATE):** 当用户的指令明确，提供了足够的信息来创建一份初步的需求文档时使用。
    - **触发场景:** 用户说 "我想做一个任务管理工具，核心功能包括任务创建、分配和状态跟踪"。
    - **你的行为:** 生成一份结构化的Markdown需求文档。如果信息不足，可以在文档中用 `[需用户提供]` 标记。
    - **输出格式:** 以 `DOC_UPDATE:::` 作为开头，后面跟上完整的Markdown文档内容。

**上下文:**
- **用户最新输入:** {user_input}

**重要规则:**
- 你的每一次回复都必须严格以 `QUESTION:::`、`DISCUSSION:::` 或 `DOC_UPDATE:::` 中的一个作为开头。
- 如果用户输入的信息不足以生成一份有意义的文档，你应该选择 `QUESTION` 模式。
- 语言必须与用户保持一致。
"""

        final_prompt = system_prompt.format(
            user_input=user_input, 
            existing_doc=existing_doc if existing_doc else "无"
        )

        api_url = f"{self.base_url}/chat/completions"
        if not self.base_url.endswith('/v1'):
            api_url = f"{self.base_url}/v1/chat/completions"
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # Ensure proper encoding for the payload
        try:
            encoded_prompt = final_prompt.encode('utf-8', errors='replace').decode('utf-8')
        except (UnicodeEncodeError, UnicodeDecodeError):
            encoded_prompt = final_prompt

        payload = {
            "model": "gpt-4-turbo",
            "messages": [
                {"role": "system", "content": encoded_prompt}
            ],
            "temperature": 0.7
        }

        try:
            # Add retry mechanism for network issues
            import time
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    # Use json parameter instead of data for proper encoding
                    response = requests.post(
                        api_url,
                        headers=headers,
                        json=payload,  # This automatically handles encoding
                        timeout=60,  # Increased timeout
                        verify=True  # Ensure SSL verification
                    )
                    response.raise_for_status()
                    break
                except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                    if attempt < max_retries - 1:
                        print(f"Network error on attempt {attempt + 1}, retrying in 2 seconds...")
                        time.sleep(2)
                        continue
                    else:
                        raise Exception(f"Network connection failed after {max_retries} attempts: {str(e)}")

            response_data = response.json()
            content = response_data.get("choices", [{}])[0].get("message", {}).get("content", "")

            # Ensure proper encoding for the response
            try:
                return content.encode('utf-8', errors='replace').decode('utf-8')
            except (UnicodeEncodeError, UnicodeDecodeError):
                return content

        except requests.exceptions.RequestException as e:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            print(f"RequirementAgent (requests) 调用失败: {error_msg}")
            return f"生成需求文档时出错: {error_msg}"
        except Exception as e:
            error_msg = str(e).encode('utf-8', errors='replace').decode('utf-8')
            print(f"RequirementAgent 在处理流时发生未知错误: {error_msg}")
            return f"处理需求文档时发生未知错误: {error_msg}"
