Alibabacloud Data Agent Skill

# Alibaba Cloud Credentials
# This SDK uses the Alibaba Cloud default credential chain.
# Credentials are resolved automatically from:
#   1. Environment variables (handled by SDK)
#   2. ~/.aliyun/config.json
#   3. ECS instance role / OIDC role

# Alternative: API Key Authentication (for file analysis scenarios)
# DATA_AGENT_API_KEY=your_api_key

# Region Configuration (Optional)
DATA_AGENT_REGION=cn-hangzhou

# DMSUnit Configuration (Optional)
# Leave unset for automatic GetActiveRouteUnit resolution.
# Set when your tenant must use a specific DMSUnit instead of the region fallback.
# DATA_AGENT_DMS_UNIT=cn-hangzhou

# Custom Endpoint (Optional, auto-generated if not set)
# DATA_AGENT_ENDPOINT=dms.cn-hangzhou.aliyuncs.com

# Timeout Settings (Optional)
DATA_AGENT_TIMEOUT=300
DATA_AGENT_MAX_RETRY=3

# Polling Settings (Optional)
DATA_AGENT_POLL_INTERVAL=2
DATA_AGENT_MAX_POLL_COUNT=60

# Qoder
.qoder/

# macOS
.DS_Store
.AppleDouble
.LSOverride
Thumbs.db

# Node.js
node_modules/
npm-debug.log*
yarn-debug.log*
yarn-error.log*
.npm
.yarn-integrity
dist/
build/
.cache/

# Python
__pycache__/
*.py[cod]
*$py.class
*.egg-info/
.eggs/
dist/
build/
.venv/
venv/
env/
.env

# Java
*.class
*.jar
*.war
*.ear
target/
.mvn/

# IDE
.idea/
.vscode/
*.swp
*.swo
*~

# Environment & Secrets
.env
.env.local
.env.*.local
*.pem
*.key

# Logs
logs/
*.log

# OS
*.tmp
*.bak


.agents/
skills-lock.json

.claude

.qwen

# 忽略所有以 validation-result- 开头、以 .json 结尾的文件
validation-result-*.json
tests/
sessions/
CLAUDE.md
evals/
.omc/
ua-fix-report.md

                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. For the purposes
      of this License, Derivative Works shall not include works that remain
      separable from, or merely link (or bind by name) to the interfaces of,
      the Work and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.

      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "[]"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.

   Copyright [yyyy] [name of copyright owner]

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

"""CLI package for Data Agent unified command-line tool.

Author: Tinker
Created: 2026-03-04
"""

from cli.formatters import (
    _fmt_jupyter_cell,
    _fmt_task_finish,
    _fmt_insights,
    _fmt_table_summaries,
    _extract_json_objects,
    _format_data_event,
    _format_parsed_json,
    _fmt_plan_progress,
    _fmt_status_change,
    _fmt_output_conclusion,
    _fmt_recommended_questions,
    _SKIP_DATA_CATEGORIES,
)
from cli.streaming import (
    StreamState,
    _print_event,
    _stream_response,
    _finalize_stream,
    _is_user_confirmation_event,
)
from cli.cmd_db import cmd_db, _build_data_source, _db_attach, _db_batch, _db_single
from cli.cmd_ls import cmd_ls, _extract_list
from cli.cmd_file import cmd_file, _print_generated_files
from cli.cmd_attach import cmd_attach
from cli.parser import build_parser, main

__all__ = [
    # Formatters
    "_fmt_jupyter_cell",
    "_fmt_task_finish",
    "_fmt_insights",
    "_fmt_table_summaries",
    "_extract_json_objects",
    "_format_data_event",
    "_format_parsed_json",
    "_fmt_plan_progress",
    "_fmt_status_change",
    "_fmt_output_conclusion",
    "_fmt_recommended_questions",
    "_SKIP_DATA_CATEGORIES",
    # Streaming
    "StreamState",
    "_print_event",
    "_stream_response",
    "_finalize_stream",
    "_is_user_confirmation_event",
    # Commands
    "cmd_db",
    "_build_data_source",
    "_db_attach",
    "_db_batch",
    "_db_single",
    "cmd_ls",
    "_extract_list",
    "cmd_file",
    "_print_generated_files",
    "cmd_attach",
    # Parser / entry
    "build_parser",
    "main",
]

"""Custom Agent management subcommand (agent).

Author: Tinker
Created: 2026-04-16
"""

import argparse
import sys

from data_agent import DataAgentConfig, DataAgentClient


def _get_field(obj: dict, *names: str, default=""):
    """Get a field value trying multiple possible key names."""
    for name in names:
        if name in obj:
            return obj[name]
    return default


def _cmd_list(client: DataAgentClient, args: argparse.Namespace) -> None:
    """List custom agents."""
    workspace_id = getattr(args, "workspace_id", None)
    search = getattr(args, "search", None)
    page_number = getattr(args, "page_number", 1)
    page_size = getattr(args, "page_size", 20)

    print(f"Region: {client.config.region}")
    print(f"Fetching custom agents (status=RELEASED)...")

    try:
        resp = client.list_custom_agents(
            workspace_id=workspace_id,
            search_key=search,
            page_number=page_number,
            page_size=page_size,
        )
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    # Extract agent list from response
    data = resp.get("data") or resp.get("Data") or {}
    items = []
    if isinstance(data, dict):
        items = data.get("content") or data.get("Content") or []
    elif isinstance(data, list):
        items = data

    total_elements = 0
    total_pages = 0
    if isinstance(data, dict):
        total_elements = data.get("totalElements") or data.get("TotalElements") or len(items)
        total_pages = data.get("totalPages") or data.get("TotalPages") or 1

    if not items:
        print("No custom agents found.")
        return

    sep = "=" * 60
    print(f"\n{sep}")
    print(f"  Custom Agents ({len(items)})  Total: {total_elements}")
    print(f"{sep}")

    for agent in items:
        name = _get_field(agent, "agentName", "AgentName", "name", "Name")
        agent_id = _get_field(agent, "customAgentId", "CustomAgentId")
        status = _get_field(agent, "status", "Status", default="unknown")
        ws_id = _get_field(agent, "workspaceId", "WorkspaceId")
        creator = _get_field(agent, "creator", "Creator")
        desc = _get_field(agent, "description", "Description")

        print(f"\n  {name}  [{status}]")
        print(f"    CustomAgentId : {agent_id}")
        if ws_id:
            print(f"    WorkspaceId   : {ws_id}")
        if creator:
            print(f"    Creator       : {creator}")
        if desc:
            print(f"    Description   : {desc}")

    if total_pages > 1:
        print(f"\n  Page {page_number}/{total_pages} (use --page-number to navigate)")

    # Print usage hint
    print(f"\n{'-' * 60}")
    print("  Tip: 使用 db --custom-agent-id <ID> ... 来指定自定义Agent进行分析")
    print(f"{'-' * 60}")
    print()


def _cmd_describe(client: DataAgentClient, args: argparse.Namespace) -> None:
    """Describe a custom agent in detail."""
    custom_agent_id = getattr(args, "custom_agent_id", None)
    if not custom_agent_id:
        print("Error: --custom-agent-id is required for describe action", file=sys.stderr)
        sys.exit(1)

    workspace_id = getattr(args, "workspace_id", None)

    print(f"Region: {client.config.region}")
    print(f"Fetching custom agent details: {custom_agent_id}...")

    try:
        resp = client.describe_custom_agent(
            custom_agent_id=custom_agent_id,
            workspace_id=workspace_id,
        )
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    # Extract agent detail from response
    data = resp.get("data") or resp.get("Data") or {}

    if not data:
        print("No agent details returned.")
        return

    name = _get_field(data, "agentName", "AgentName", "name", "Name")
    agent_id = _get_field(data, "customAgentId", "CustomAgentId")
    status = _get_field(data, "status", "Status", default="unknown")
    ws_id = _get_field(data, "workspaceId", "WorkspaceId")
    creator = _get_field(data, "creator", "Creator")
    desc = _get_field(data, "description", "Description")
    instruction = _get_field(data, "instruction", "Instruction")
    knowledge = _get_field(data, "knowledge", "Knowledge")

    sep = "=" * 60
    print(f"\n{sep}")
    print(f"  Custom Agent Detail")
    print(f"{sep}")
    print(f"  Name            : {name}")
    print(f"  CustomAgentId   : {agent_id}")
    print(f"  Status          : {status}")
    if ws_id:
        print(f"  WorkspaceId     : {ws_id}")
    if creator:
        print(f"  Creator         : {creator}")
    if desc:
        print(f"  Description     : {desc}")
    if instruction:
        print(f"  Instruction     : {instruction}")
    if knowledge:
        print(f"  Knowledge       : {knowledge}")

    # Print any additional fields not covered above
    known_keys = {
        "agentName", "AgentName", "name", "Name",
        "customAgentId", "CustomAgentId",
        "status", "Status",
        "workspaceId", "WorkspaceId",
        "creator", "Creator",
        "description", "Description",
        "instruction", "Instruction",
        "knowledge", "Knowledge",
    }
    extra = {k: v for k, v in data.items() if k not in known_keys and v}
    if extra:
        print(f"\n  -- Additional Fields --")
        for k, v in extra.items():
            print(f"  {k:16s}: {v}")

    print(f"\n{'-' * 60}")
    print(f"  Tip: db --custom-agent-id {agent_id} ... 来使用此Agent")
    print(f"{'-' * 60}")
    print()


def cmd_agent(args: argparse.Namespace) -> None:
    """Handle agent subcommand."""
    config = DataAgentConfig.from_env()
    client = DataAgentClient(config)

    action = getattr(args, "action", "list")

    if action == "describe":
        _cmd_describe(client, args)
    else:
        _cmd_list(client, args)

"""Attach to existing session subcommand (attach).

Author: Tinker
Created: 2026-03-04
"""

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

from cli.streaming import _stream_response, StreamState, _print_event
from cli.cmd_db import _db_attach
from cli.log_handler import StructuredLogHandler
# from cli.notify import push_notification
from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion
from cli.streaming_utils import run_worker_with_handler, execute_single_query
from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock
from data_agent import (
    DataAgentConfig,
    DataAgentClient,
    SessionManager,
    MessageHandler,
    FileManager,
    SSEClient,
)


def cmd_attach(args: argparse.Namespace) -> None:
    """Connect to an existing session for continuing analysis or confirming plan."""
    session_id = args.session_id
    is_worker = os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1"
    async_run = getattr(args, "async_run", True)

    # Initialize components
    config = DataAgentConfig.from_env()
    client = DataAgentClient(config)
    session_manager = SessionManager(client)
    message_handler = MessageHandler(client)
    sse_client = SSEClient(config)
    file_manager = FileManager(client)

    # Async mode is determined by the --async-run/--no-async-run flag, regardless of query presence
    if async_run and not is_worker:
        # PARENT PROCESS: spawn background worker
        print(f"Connecting to session: {session_id}")
        try:
            session = client.describe_session(session_id=session_id)
        except Exception as e:
            print(f"Error: Failed to connect to session: {e}", file=sys.stderr)
            sys.exit(1)

        # Manually create a temporary session object to pass to the common utility
        class TempSession:
            def __init__(self, sess_obj):
                self.session_id = sess_obj.session_id
                self.agent_id = sess_obj.agent_id

        temp_session = TempSession(session)

        # Use common async worker setup
        setup_async_worker(args, temp_session)

    elif is_worker:
        # WORKER PROCESS: run the attach operation in background - using common utility
        def attach_query_executor(message_handler, session, args):
            output_mode = getattr(args, "output", "summary")
            session_dir = Path(f"sessions/{session.session_id}")

            query = args.query

            if query:
                print(f"\n> User Query: {query}\n", flush=True)
                got_content, need_confirm, output_text = _stream_response(
                    message_handler, session, query,
                    output_mode=output_mode, output_dir=session_dir,
                    process_log_handler=None,  # In worker process, no process log handler needed for this path
                    is_attach=True  # Indicate this is an attach operation
                )

                # Check if the query was for confirmation and if there are more steps to process
                is_confirmation_query = query.strip() in ["确认执行", "confirm", "execute", "同意后续所有SQL执行"]

                # If the query was a confirmation, there might be more steps to execute
                if is_confirmation_query:
                    # After confirmation, the process has handled the confirmation
                    # The subsequent steps will be tracked by the server and can be viewed
                    # by attaching to the session again
                    print(f"\n[Worker] Confirmation processed. Subsequent steps will continue in the background.", flush=True)

                # Determine final status
                final_status = "waiting_input" if need_confirm else "completed"

                with open(session_dir / "status.txt", "w", encoding="utf-8") as f:
                    f.write(final_status)
                with open(session_dir / "result.json", "w", encoding="utf-8") as f:
                    json.dump({
                        "status": final_status,
                        "session_id": session.session_id,
                    }, f, ensure_ascii=False, indent=2)

                return got_content, need_confirm
            else:
                # For attach without query, just monitor or replay
                from_start = getattr(args, "from_start", False)
                checkpoint = getattr(args, "checkpoint", None)

                # In worker process, we need to initialize components locally
                _config = DataAgentConfig.from_env()
                _client = DataAgentClient(_config)
                _sse_client = SSEClient(_config)
                _file_manager = FileManager(_client)

                _db_attach(_sse_client, _file_manager, session, from_start=from_start, checkpoint=checkpoint, output_mode=output_mode)

                return True, False  # got_content=True, need_confirm=False

        # Since the attach worker has complex setup logic that differs significantly from the other commands,
        # we'll just call the run_worker_with_handler without a data_source
        run_worker_with_handler(args, query_execution_func=attach_query_executor)

    else:
        # SYNCHRONOUS MODE (--no-async-run)
        print(f"Connecting to session: {session_id}")
        print(f"  Region: {config.region}")
        try:
            # First, try to get the session info to discover the actual agent_id
            # For some sessions, the API may not return the agent_id if an empty agent_id is provided
            # We need to try with empty agent_id first, and if it fails, we'll try the create_or_reuse approach
            session = None
            try:
                session = client.describe_session(agent_id="", session_id=session_id)

                # Special case: If we get CREATING status with no agent ID, this means the session doesn't actually exist
                # for the current user, but the API returns a default response. We should exit immediately.
                if not session.agent_id and session.status.value == "CREATING":
                    print(f"Error: Session {session_id} does not exist or does not belong to the current user.", file=sys.stderr)
                    sys.exit(1)

                # Only print session info if it passes the above check
                print(f"Initial session check - agent: '{session.agent_id}', status: {session.status.value}")

            except Exception as e:
                # This could be an authentication error or the session doesn't exist
                print(f"Error connecting to session: {e}", file=sys.stderr)
                print(f"This could be due to invalid credentials or the session doesn't exist.", file=sys.stderr)

                # If it's an authentication error, we should exit immediately
                from data_agent.exceptions import AuthenticationError
                if isinstance(e, AuthenticationError):
                    print(f"Authentication failed. Please check your credentials in .env file.", file=sys.stderr)
                    sys.exit(1)

                # For other errors, continue with fallback approaches
                print(f"Trying alternative methods to connect to session...", file=sys.stderr)

            # If session retrieval failed or agent_id is still empty, try using session manager
            if not session or not session.agent_id:
                print("Agent ID not found in session info, this may indicate the session doesn't exist or isn't ready yet.")

                # Check if the session exists but is still in CREATING state
                if session and session.status.value == "CREATING":
                    print("Session exists but is still being created. Waiting for it to be ready...")
                    # Since session is in CREATING state but no agent_id was returned,
                    # we need to wait until the agent_id becomes available
                    import time
                    max_retries = 10  # Reduce wait time (10*30 = 5 minutes max)
                    retry_count = 0

                    while retry_count < max_retries:
                        time.sleep(30)  # Wait 30 seconds between checks
                        try:
                            session = client.describe_session(agent_id="", session_id=session_id)
                            print(f"Checking session again - agent: '{session.agent_id}', status: {session.status.value}")

                            # If we now have an agent_id, break out of the loop
                            if session.agent_id:
                                print(f"Got agent ID: {session.agent_id}. Session is ready.")
                                break

                            # If status is no longer CREATING but we still don't have an agent_id,
                            # the session may have failed to create properly or doesn't belong to the user
                            if session.status.value != "CREATING":
                                print(f"Session status changed to {session.status.value} but still no agent ID.")
                                print(f"This may indicate the session does not belong to the current user.", file=sys.stderr)
                                break

                        except Exception as api_error:
                            print(f"Check session retry failed: {api_error}", file=sys.stderr)
                            from data_agent.exceptions import AuthenticationError
                            if isinstance(api_error, AuthenticationError):
                                print(f"Authentication failed. Please check your credentials in .env file.", file=sys.stderr)
                                sys.exit(1)

                        retry_count += 1

                    if not session.agent_id:
                        print(f"Session is still not ready after waiting. Status: {session.status.value if session else 'unknown'}", file=sys.stderr)
                        print(f"If status is not CREATING, this may mean the session does not belong to your account.", file=sys.stderr)

                if not session or not session.agent_id:
                    # At this point, if we still don't have an agent_id, the session either:
                    # 1. Does not exist
                    # 2. Does not belong to the current user
                    # 3. Is still being created but taking too long
                    # 4. Has been cancelled or failed to create properly
                    print(f"Error: Could not retrieve agent ID for session {session_id}.", file=sys.stderr)
                    print(f"This may mean the session does not exist, does not belong to your account,", file=sys.stderr)
                    print(f"is still being created (waited too long), or was cancelled/failed to create.", file=sys.stderr)
                    sys.exit(1)

            rid = f", request_id: {session.request_id}" if session.request_id else ""
            print(f"Session connected: {session.session_id} (agent: {session.agent_id}, status: {session.status.value}{rid})")

            # If status is CREATING and a query is provided, we need to wait for it to be ready
            # But for IDLE state, we can proceed without waiting since it's a valid state
            if session.status.value == "CREATING" and getattr(args, "query", None):
                print("Waiting for session to be ready...")
                session = session_manager.create_or_reuse(
                    session_id=session_id, agent_id=session.agent_id, wait_for_running=True,
                )
                rid = f", request_id: {session.request_id}" if session.request_id else ""
                print(f"Session ready: {session.session_id} (status: {session.status.value}{rid})")

        except Exception as e:
            print(f"Error: Failed to connect to session: {e}", file=sys.stderr)
            sys.exit(1)

        output_mode = getattr(args, "output", "summary")
        session_dir = Path(f"sessions/{session.session_id}")

        # Create process log handler for attach command
        with StructuredLogHandler(session_dir, log_prefix="process") as log_handler:
            # Log initial connection info
            log_text = f"Connected to session: {session.session_id} (agent: {session.agent_id}, status: {session.status.value})"
            log_handler.write_both(log_text + "\n")

            if args.query:
                query = args.query
                print(f"\n> User Query: {query}\n")

                # Log the query
                query_log = f"Processing query: {query}\n"
                log_handler.write_both(query_log)

                try:
                    got_content, need_confirm, _ = _stream_response(
                        message_handler, session, query,
                        output_mode=output_mode, output_dir=session_dir,
                        process_log_handler=log_handler,  # Pass the process log handler
                    is_attach=True  # Indicate this is an attach operation
                    )

                    # Check if the query was for confirmation and if there are more steps to process
                    is_confirmation_query = query.strip() in ["确认执行", "confirm", "execute", "同意后续所有SQL执行"]

                    # Log the response completion
                    response_log = f"Query processed, got_content: {got_content}, need_confirm: {need_confirm}, is_confirmation: {is_confirmation_query}\n"
                    log_handler.write_both(response_log)

                    if not got_content:
                        print("(No response received, please retry)")
                    elif need_confirm:
                        print("\n⚠️  Agent has created an execution plan. User confirmation required.")
                        print(f"   To continue: python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'your input'")
                    elif is_confirmation_query:
                        # After confirmation, continue monitoring for additional steps
                        print(f"\n[Sync mode] Post-confirmation: Monitoring for additional steps...")

                        # Refresh session to get the latest status
                        updated_session = client.describe_session(agent_id=session.agent_id, session_id=session.session_id)

                        # If session is still running after confirmation, continue monitoring
                        if updated_session.status.value in ["RUNNING", "WAIT_INPUT"]:
                            print(f"Session status after confirmation: {updated_session.status.value}")

                            # Continue to attach and monitor the session
                            print(f"\nContinuing to monitor session progress...")

                            # Use already imported _db_attach and initialized clients
                            _db_attach(sse_client, file_manager, updated_session, from_start=False, checkpoint=None, output_mode=output_mode)
                    else:
                        try:
                            updated_session = client.describe_session(agent_id=session.agent_id, session_id=session.session_id)
                            if updated_session.status.value == "WAIT_INPUT":
                                print("\n⚠️  Agent has created an execution plan and is waiting for confirmation.")
                                print("   Use -q option to confirm the plan or provide feedback:")
                                print(f"     python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'confirm'")
                                print(f"     python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'modify the plan'")
                        except Exception:
                            pass
                except Exception as e:
                    error_msg = f"Request failed: {e}"
                    print(error_msg)
                    log_handler.write_both(error_msg + "\n")
            else:
                from_start = getattr(args, "from_start", False)
                checkpoint = getattr(args, "checkpoint", None)

                # Log the attachment without query
                attach_log = f"Attaching to session without query (from_start: {from_start}, checkpoint: {checkpoint})\n"
                log_handler.write_both(attach_log)

                # Additional verification: if session agent_id is still empty, try to retrieve it directly
                if not session.agent_id:
                    print(f"No agent ID in current session object, attempting to refresh session info...", file=sys.stderr)
                    try:
                        # Try to get fresh session data directly from the API
                        refreshed_session = client.describe_session(agent_id="", session_id=session_id)
                        if refreshed_session.agent_id:
                            session = refreshed_session
                            print(f"Successfully retrieved agent ID: {refreshed_session.agent_id}")
                        else:
                            print(f"Error: Could not retrieve agent ID for session {session_id}. Session may not exist.", file=sys.stderr)
                            sys.exit(1)
                    except Exception as e:
                        print(f"Error retrieving session info: {e}", file=sys.stderr)
                        sys.exit(1)

                # Final check before calling _db_attach
                if not session.agent_id:
                    print(f"Error: Cannot attach to session {session_id} without a valid agent ID.", file=sys.stderr)
                    print(f"The session might not exist or may be corrupted.", file=sys.stderr)
                    sys.exit(1)

                _db_attach(sse_client, file_manager, session, from_start=from_start, checkpoint=checkpoint, output_mode=output_mode)

"""Database analysis subcommand (db).

Author: Tinker
Created: 2026-03-04
"""

import argparse
import sys
import os
import json
import subprocess
from pathlib import Path
from typing import Optional

# from cli.notify import push_notification
from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock
from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion
from cli.streaming_utils import run_worker_with_handler, execute_single_query, execute_query_batch
from cli.streaming import _stream_response, _finalize_stream, StreamState, _print_event
from data_agent import (
    DataAgentConfig,
    DataAgentClient,
    SessionManager,
    MessageHandler,
    FileManager,
    DataSource,
    SSEClient,
)


def _build_data_source(args: argparse.Namespace) -> DataSource:
    """Build DataSource from command-line arguments."""
    tables = [t.strip() for t in args.tables.split(",")] if args.tables else []
    table_ids = [t.strip() for t in args.table_ids.split(",")] if args.table_ids else []

    return DataSource(
        dms_instance_id=getattr(args, 'dms_instance_id', None),
        dms_database_id=args.dms_db_id,
        instance_name=getattr(args, 'instance_name', None) or "",
        db_name=args.db_name,
        tables=tables,
        table_ids=table_ids,
        engine=args.engine,
        region_id=args.region,
    )


def _db_attach(
    sse_client: SSEClient,
    file_manager: FileManager,
    session,
    from_start: bool = False,
    checkpoint: Optional[int] = None,
    output_mode: str = "summary",
) -> None:
    """Attach to an existing session's SSE stream without sending a message.

    Streams all incoming events to stdout in real-time using the same
    formatting as _stream_response.  Useful for:
      - Watching an ongoing analysis
      - Replaying the last round to review a plan before confirming

    After the stream ends, calls ListFileUpload to download any
    agent-generated report files to sessions/<session_id>/reports/.
    """
    session_id = session.session_id
    session_dir = Path(f"sessions/{session_id}")

    if checkpoint is None:
        checkpoint = 0 if from_start else None

    label = "watching live stream"
    if from_start:
        label = "replaying from start"
    elif checkpoint is not None:
        label = f"resuming from checkpoint {checkpoint}"

    print(f"\nAttaching to session {session_id} ({label})...")
    print(f"Session status: {session.status.value}")

    # Special reminder for WAIT_INPUT status
    if session.status.value == "WAIT_INPUT":
        print("\n⚠️  Session is in WAIT_INPUT state.")
        print("   The agent has generated SQL and is waiting for your confirmation.")
        print()
        print("   To view the SQL and confirm:")
        print(f"     python3 skill/data_agent_cli.py attach --session-id {session_id} --from-start")
        print()
        print("   To confirm and execute ONLY this SQL (DO NOT create a new session):")
        print(f"     python3 skill/data_agent_cli.py attach --session-id {session_id} -q '确认执行当前SQL'")
        print()
        print("   To agree to execute all subsequent SQL automatically:")
        print(f"     python3 skill/data_agent_cli.py attach --session-id {session_id} -q '同意后续所有SQL执行'")
        print()
        print("   To modify the query:")
        print(f"     python3 skill/data_agent_cli.py attach --session-id {session_id} -q 'your new question'")

    print(f"Output directory: {session_dir.resolve()}")
    print("Press Ctrl+C to detach.")
    print("\n---\n")

    got_content = False
    last_checkpoint = checkpoint if checkpoint else 0
    last_progress_time = 0
    import time

    state = StreamState(output_mode=output_mode)
    state.output_dir = session_dir
    state.session_id = session_id
    state.is_attach = True
    state.session_status = getattr(session, 'status', None)
    if hasattr(state.session_status, 'value'):
        state.session_status = state.session_status.value

    try:
        # Initialize structured logging for attach mode
        from cli.streaming import init_structured_logging, close_structured_logging
        init_structured_logging(session_dir)

        for event in sse_client.stream_chat_content(
            agent_id=session.agent_id,
            session_id=session_id,
            checkpoint=checkpoint,
        ):
            if event.event_type == "SSE_FINISH":
                break

            # Track checkpoint progress (only in detail/raw mode)
            if output_mode != "summary" and event.checkpoint is not None and event.checkpoint > last_checkpoint:
                current_time = time.time()
                # Show checkpoint progress every 30 seconds or every 50 checkpoints
                if current_time - last_progress_time >= 30 or event.checkpoint - last_checkpoint >= 50:
                    print(f"  [checkpoint: {event.checkpoint}]", flush=True)
                    last_progress_time = current_time
                last_checkpoint = event.checkpoint

            c, _ = _print_event(event, output_mode, state=state)
            if c:
                got_content = True

    except KeyboardInterrupt:
        print("\n\nDetached.")
        return
    except Exception as e:
        # Extract request_id from HTTP error response if available
        request_id = ""
        resp = getattr(e, "response", None)
        if resp is not None:
            request_id = resp.headers.get("x-acs-request-id", "")
        rid_str = f" (Request-Id: {request_id})" if request_id else ""
        print(f"\nError: {e}{rid_str}", file=sys.stderr)
        return
    finally:
        # Close structured logging
        close_structured_logging()

    _finalize_stream(state)
    if state.got_content:
        got_content = True

    if got_content:
        print()  # trailing newline
    else:
        print("(No new events -- session may be waiting for your input)")

    # -- Download agent-generated files via ListFileUpload --
    # Server needs a few seconds to finalize report files after session completes
    print("\nFetching agent-generated files (waiting for server to finalize)...")
    time.sleep(5)
    report_dir = session_dir / "reports"
    total_reports = 0
    for category in ("WebReport", "TextReport", "DefaultArtifact"):
        try:
            files = file_manager.list_files(session_id, file_category=category)
        except Exception as e:
            print(f"  Warning: could not list {category}: {e}", file=sys.stderr)
            continue
        if not files:
            continue
        report_dir.mkdir(parents=True, exist_ok=True)
        for rf in files:
            if not rf.download_url:
                print(f"  Skipping {rf.filename} ({category}): no download URL")
                continue
            save_path = report_dir / (rf.filename or f"{rf.file_id}.bin")
            try:
                file_manager.download_from_url(rf.download_url, str(save_path))
                print(f"  [{category}] saved \u2192 {save_path.resolve()}")
                total_reports += 1
            except Exception as e:
                print(f"  Failed to download {rf.filename} ({category}): {e}", file=sys.stderr)

    if total_reports == 0:
        print("  No report files found for this session.")

    print("\n---\n")
    print(f'> \U0001f4a1 To continue conversation:')
    print(f'>    python3 skill/data_agent_cli.py attach --session-id {session_id} -q "your message"')


def _db_batch(
    message_handler: MessageHandler,
    session,
    data_source: DataSource,
    queries: list,
    output_mode: str = "summary",
    output_dir: Optional[Path] = None,
) -> tuple[bool, bool, str]:
    """Execute batch preset queries."""
    got_content, need_confirm = False, False
    full_text = ""
    for query in queries:
        print(f"\n{'=' * 60}")
        print(f"Query: {query}")
        print("=" * 60)
        c, nc, t = _stream_response(message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir)
        if c: got_content = True
        if t: full_text += f"\n### Query: {query}\n" + t + "\n"
        if nc:
            need_confirm = True
            break
    return got_content, need_confirm, full_text


def _db_single(
    message_handler: MessageHandler,
    session,
    data_source: DataSource,
    query: str,
    output_mode: str = "summary",
    output_dir: Optional[Path] = None,
) -> tuple[bool, bool, str]:
    """Execute a single query with streaming output."""
    print(f"\nAnalyzing...\n")
    got_content, need_confirm, full_text = _stream_response(
        message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir
    )
    if not got_content:
        print("(No response received, please retry)")
    elif need_confirm:
        print("\n\u26a0\ufe0f  \u9700\u8981\u7528\u6237\u786e\u8ba4\uff0c\u7a0b\u5e8f\u5c06\u9000\u51fa\u3002\u8bf7\u5b8c\u6210\u786e\u8ba4\u540e\u4f7f\u7528\u4f1a\u8bddID\u7ee7\u7eed\u5bf9\u8bdd\u3002")
    return got_content, need_confirm, full_text


def cmd_db(args: argparse.Namespace) -> None:
    """Handle db subcommand."""
    # Validate required database parameters
    missing = []
    for attr, name in [
        ("dms_db_id", "--dms-db-id"),
        ("db_name", "--db-name"),
        ("tables", "--tables"),
    ]:
        if not getattr(args, attr, None):
            missing.append(name)
    if missing:
        print(f"Error: Missing required parameters: {', '.join(missing)}", file=sys.stderr)
        sys.exit(1)

    # Initialize components
    config = DataAgentConfig.from_env()
    client = DataAgentClient(config)
    session_manager = SessionManager(client)
    message_handler = MessageHandler(client)

    # Create new session
    session_mode = args.session_mode.upper()
    data_source = _build_data_source(args)
    is_worker = is_worker_process()

    if getattr(args, "async_run", False) and not is_worker:
        # PARENT PROCESS LOGIC
        enable_search = getattr(args, 'enable_search', False)
        print(f"Creating session for async execution...")
        workspace_id = getattr(args, 'workspace_id', None)
        custom_agent_id = getattr(args, 'custom_agent_id', None)
        session = session_manager.create_or_reuse(mode=session_mode, database_id=str(args.dms_db_id), enable_search=enable_search, workspace_id=workspace_id, custom_agent_id=custom_agent_id)

        # Use common async worker setup
        setup_async_worker(args, session)

    elif is_worker:
        # WORKER PROCESS LOGIC - using common utility
        def db_query_executor(message_handler, session, args):
            output_mode = getattr(args, "output", "summary")
            data_source = _build_data_source(args)

            if args.query:
                _, need_confirm, _ = _db_single(message_handler, session, data_source, args.query, output_mode=output_mode, output_dir=Path(f"sessions/{session.session_id}"))
            else:
                if session_mode == "ANALYSIS":
                    default_queries = [
                        f"Analyze the overall data structure and table relationships of {data_source.db_name} database",
                        "Identify key metrics and distribution characteristics in the data",
                    ]
                else:
                    default_queries = [
                        f"What tables exist in {data_source.db_name} database?",
                        "Who has the highest sales?",
                    ]
                _, need_confirm, _ = _db_batch(message_handler, session, data_source, default_queries, output_mode=output_mode, output_dir=Path(f"sessions/{session.session_id}"))
            return True, need_confirm  # got_content=True, need_confirm

        run_worker_with_handler(args, data_source=_build_data_source(args), query_execution_func=db_query_executor)

    # NORMAL SYNCHRONOUS LOGIC
    mode_desc = {
        "ASK_DATA": "ASK_DATA mode (SQL query + natural language response)",
        "ANALYSIS": "ANALYSIS mode (deep analysis + report generation)",
        "INSIGHT": "INSIGHT mode",
        "CLAW": "CLAW mode (agentic)",
    }.get(session_mode, session_mode)

    print(f"Creating session: {mode_desc}...")
    print(f"  Region: {config.region}")
    enable_search = getattr(args, 'enable_search', False)
    workspace_id = getattr(args, 'workspace_id', None)
    custom_agent_id = getattr(args, 'custom_agent_id', None)
    session = session_manager.create_or_reuse(mode=session_mode, database_id=str(args.dms_db_id), enable_search=enable_search, workspace_id=workspace_id, custom_agent_id=custom_agent_id)
    print(f"Session ready: {session.session_id}")
    print(f"\n💡 Tip: To continue this session later, use: python3 scripts/data_agent_cli.py attach --session-id {session.session_id}")

    # Execute query
    output_mode = getattr(args, "output", "summary")
    session_dir = Path(f"sessions/{session.session_id}")
    session_dir.mkdir(parents=True, exist_ok=True)

    # Initialize structured logging for sync mode
    from cli.streaming import init_structured_logging, close_structured_logging
    init_structured_logging(session_dir)

    try:
        if args.query:
            _, _, output_text = _db_single(message_handler, session, data_source, args.query, output_mode=output_mode, output_dir=session_dir)
        else:
            # Default batch preset queries
            if session_mode == "ANALYSIS":
                default_queries = [
                    f"Analyze the overall data structure and table relationships of {data_source.db_name} database",
                    "Identify key metrics and distribution characteristics in the data",
                ]
            else:
                default_queries = [
                    f"What tables exist in {data_source.db_name} database?",
                    "Who has the highest sales?",
                ]
            print(f"\nNo query specified, running preset queries ({len(default_queries)} total)...")
            _, _, output_text = _db_batch(message_handler, session, data_source, default_queries, output_mode=output_mode, output_dir=session_dir)
    finally:
        close_structured_logging()

    # Write result status
    # if output_text:
    #     with open(session_dir / "output.md", "w", encoding="utf-8") as f:
    #         f.write(output_text)
    #     with open(session_dir / "result.json", "w", encoding="utf-8") as f:
    #         json.dump({"status": "completed", "output_file": "output.md"}, f)
    with open(session_dir / "result.json", "w", encoding="utf-8") as f:
        json.dump({"status": "completed"}, f)

"""DMS tools subcommand (dms).

Author: Tinker
Created: 2026-03-05
"""

import argparse
import sys

from data_agent import DataAgentConfig, DmsMcpTools


def cmd_dms(args: argparse.Namespace) -> None:
    """Handle dms subcommand for DMS tools."""
    config = DataAgentConfig.from_env()
    mcp_tools = DmsMcpTools(config)

    tool = getattr(args, "tool", None)

    if tool == "list-instances":
        _list_instances(mcp_tools, args)
    elif tool == "search-database":
        _search_database(mcp_tools, args)
    elif tool == "list-tables":
        _list_tables(mcp_tools, args)
    else:
        print("Error: Unknown tool. Use 'list-instances', 'search-database', or 'list-tables'.", file=sys.stderr)
        sys.exit(1)


def _list_instances(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None:
    """List DMS instances using MCP tool."""
    search_key = getattr(args, "search", None)
    db_type = getattr(args, "db_type", None)
    env_type = getattr(args, "env_type", None)
    page_number = getattr(args, "page_number", 1)
    page_size = getattr(args, "page_size", 50)

    print(f"Region: {mcp_tools._config.region}")
    print("Fetching instances from DMS...")
    print("-" * 60)

    try:
        result = mcp_tools.list_instances(
            search_key=search_key,
            db_type=db_type,
            env_type=env_type,
            page_number=page_number,
            page_size=page_size,
        )

        if not result.items:
            print("No instances found.")
            return

        print(f"Page {result.page_number}/{result.total_pages} (Total: {result.total_count} instances)\n")

        for inst in result.items:
            print(f"  {inst.instance_alias}  [{inst.instance_type}]")
            print(f"    Instance ID   : {inst.instance_id}")
            print(f"    Host:Port     : {inst.host}:{inst.port}")
            print(f"    State         : {inst.state}")
            print(f"    Env Type      : {inst.env_type}")
            print(f"    Source        : {inst.instance_source}")
            if inst.instance_resource_id:
                print(f"    Resource ID   : {inst.instance_resource_id}")
            print()

        # Show pagination hint
        print("-" * 60)
        if result.has_next:
            print(f"📄 Next page: --page-number {result.page_number + 1}")
        else:
            print("📄 This is the last page.")

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


def _search_database(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None:
    """Search databases by schema name."""
    search_key = getattr(args, "search_key", None)
    page_number = getattr(args, "page_number", 1)
    page_size = getattr(args, "page_size", 200)

    if not search_key:
        print("Error: --search-key is required for search-database tool.", file=sys.stderr)
        sys.exit(1)

    print(f"Region: {mcp_tools._config.region}")
    print(f"Search Key: {search_key}")
    print("Searching databases in DMS...")
    print("-" * 60)

    try:
        databases = mcp_tools.search_database(
            search_key=search_key,
            page_number=page_number,
            page_size=page_size,
        )

        if not databases:
            print("No databases found.")
            print("\nTroubleshooting tips:")
            print("  1. The database may not be registered in DMS")
            print("  2. The search_key is case-sensitive, try different cases")
            print("  3. Try using 'dms list-instances' to see available instances first")
            return

        print(f"Found {len(databases)} database(s):\n")

        for db in databases:
            print(f"  {db.schema_name}  [{db.db_type}]")
            print(f"    Database ID   : {db.database_id}")
            print(f"    Host:Port     : {db.host}:{db.port}")
            print(f"    Instance      : {db.instance_alias} (ID: {db.instance_id})")
            print(f"    Env Type      : {db.env_type}")
            print()

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


def _list_tables(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None:
    """List tables in a database."""
    database_id = getattr(args, "database_id", None)
    search_name = getattr(args, "search_name", None)
    page_number = getattr(args, "page_number", 1)
    page_size = getattr(args, "page_size", 200)

    if not database_id:
        print("Error: --database-id is required for list-tables tool.", file=sys.stderr)
        sys.exit(1)

    print(f"Region: {mcp_tools._config.region}")
    print(f"Database ID: {database_id}")
    if search_name:
        print(f"Search Name: {search_name}")
    print("Fetching tables from DMS...")
    print("-" * 60)

    try:
        tables = mcp_tools.list_tables(
            database_id=database_id,
            search_name=search_name,
            page_number=page_number,
            page_size=page_size,
        )

        if not tables:
            print("No tables found.")
            return

        print(f"Found {len(tables)} table(s):\n")

        for table in tables:
            print(f"  {table.table_name}")
            print(f"    Table ID      : {table.table_id}")
            print(f"    Table GUID    : {table.table_guid}")
            print(f"    Schema        : {table.schema_name}")
            if table.engine:
                print(f"    Engine        : {table.engine}")
            if table.table_comment:
                print(f"    Comment       : {table.table_comment}")
            print()

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

"""File analysis subcommand (file).

Author: Tinker
Created: 2026-03-04
"""

import argparse
import sys
import os
import json
import subprocess
from pathlib import Path

# from cli.notify import push_notification
from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion
from cli.streaming_utils import run_worker_with_handler, execute_single_query, execute_query_batch
from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock
from cli.dual_logger import run_with_dual_logging
from cli.streaming import _stream_response
from data_agent import (
    DataAgentConfig,
    DataAgentClient,
    SessionManager,
    MessageHandler,
    FileManager,
    DataSource,
)


def cmd_file(args: argparse.Namespace) -> None:
    """Handle file subcommand."""
    file_id = getattr(args, 'file_id', None)
    file_path = args.file_path

    # Validate arguments
    if not file_id and not file_path:
        print("Error: Either FILE path or --file-id must be provided", file=sys.stderr)
        sys.exit(1)
    if file_id and file_path:
        print("Error: Cannot specify both FILE path and --file-id", file=sys.stderr)
        sys.exit(1)

    # Initialize components
    config = DataAgentConfig.from_env()
    client = DataAgentClient(config)
    session_manager = SessionManager(client)
    message_handler = MessageHandler(client)
    file_manager = FileManager(client)

    is_worker = is_worker_process()
    session_mode = args.session_mode.upper()

    if getattr(args, "async_run", False) and not is_worker:
        # PARENT PROCESS LOGIC
        enable_search = getattr(args, 'enable_search', False)
        print(f"Creating session for async file analysis...")

        if file_id:
            # Using Data Center file ID
            print(f"Using Data Center file ID: {file_id}")

            # Create data source with file ID for query execution
            file_data_source = DataSource(
                data_source_type="FILE",
                file_id=file_id,
                region_id=config.region
            )

            # Create regular session without binding to specific data source initially
            workspace_id = getattr(args, 'workspace_id', None)
            custom_agent_id = getattr(args, 'custom_agent_id', None)
            session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_id, workspace_id=workspace_id, custom_agent_id=custom_agent_id)
        else:
            # Local file workflow: upload first
            if not Path(file_path).exists():
                print(f"Error: File not found: {file_path}", file=sys.stderr)
                sys.exit(1)

            # Validate file type
            if not file_manager.is_supported_file(file_path):
                supported = ".csv, .xlsx, .xls"
                print(f"Error: Unsupported file type. Supported formats: {supported}", file=sys.stderr)
                sys.exit(1)

            # Upload file and get file info
            print(f"Uploading file: {file_path}")
            try:
                file_info = file_manager.upload_file(file_path)
                print(f"File uploaded successfully!")
                print(f"  File ID : {file_info.file_id}")
                print(f"  Filename: {file_info.filename}")
                print(f"  Size    : {file_info.size} bytes")

                # Create data source with uploaded file ID
                file_data_source = DataSource(
                    data_source_type="FILE",
                    file_id=file_info.file_id,
                )
            except Exception as e:
                print(f"Upload failed: {e}", file=sys.stderr)
                sys.exit(1)

            # Create session
            workspace_id = getattr(args, 'workspace_id', None)
            custom_agent_id = getattr(args, 'custom_agent_id', None)
            session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_info.file_id, workspace_id=workspace_id, custom_agent_id=custom_agent_id)

        # Store file_data_source as a dict for async worker to use (avoid serialization issues)
        args.file_data_source = file_data_source.to_api_dict()

        # Use common async worker setup
        setup_async_worker(args, session)

        print(f"\n✅ Async task started. Session ID: {session.session_id}")
        print(f"Check progress at: sessions/{session.session_id}/progress.log")

        sys.exit(0)

    elif is_worker:
        # WORKER PROCESS LOGIC - using common utility
        def file_query_executor(message_handler, session, args):
            if hasattr(args, 'file_data_source'):
                # If file_data_source is a dict (async mode), reconstruct DataSource
                if isinstance(args.file_data_source, dict):
                    # Create a new DataSource from the dict representation
                    # For file analysis, we need to use data_source_type="FILE" and the correct file_id
                    # The DataSource.to_api_dict() method expects data_source_type="FILE" to trigger FileId inclusion
                    file_data_source = DataSource(
                        data_source_type="FILE",  # Use "FILE" type for file analysis
                        file_id=args.file_data_source.get('FileId', ''),
                        region_id=args.file_data_source.get('RegionId', config.region)
                    )
                else:
                    # If file_data_source is already a DataSource object (sync mode)
                    file_data_source = args.file_data_source
            else:
                # Try to load from input.json if not available in args
                # In worker process, we can determine the session dir and load input.json
                import json
                from pathlib import Path
                session_dir = Path(f"sessions/{session.session_id}")
                input_file = session_dir / "input.json"

                if input_file.exists():
                    try:
                        with open(input_file, 'r', encoding='utf-8') as f:
                            saved_args = json.load(f)

                        if 'file_data_source' in saved_args:
                            file_data_source_dict = saved_args['file_data_source']
                            # Create DataSource for file analysis using data_source_type="FILE"
                            file_data_source = DataSource(
                                data_source_type="FILE",  # Use "FILE" type for file analysis
                                file_id=file_data_source_dict.get('FileId', ''),
                                region_id=file_data_source_dict.get('RegionId', config.region)
                            )
                    except Exception as e:
                        print(f"Warning: Could not load file_data_source from input.json: {e}", file=sys.stderr)

            if file_data_source is None:
                # Fallback: reconstruct based on available info
                if hasattr(args, 'file_id') and args.file_id:
                    file_data_source = DataSource(
                        data_source_type="FILE",
                        file_id=args.file_id,
                        region_id=getattr(args, 'region_id', 'cn-hangzhou')
                    )
                else:
                    print("Error: No file data source available in worker", file=sys.stderr)
                    return False, False

            output_mode = getattr(args, "output", "summary")
            session_dir = Path(f"sessions/{session.session_id}")

            if args.query:
                queries = [args.query]
            else:
                queries = [
                    "请分析上传文件的数据结构",
                    "数据的关键统计指标和分布情况是什么？",
                    "数据中是否存在异常值或离群点？",
                ]

            print(f"\n{'=' * 60}")
            print("File Analysis")
            print("=" * 60)

            output_text = ""
            for query in queries:
                print(f"\nQuery: {query}")
                print("-" * 50)
                got_content, need_confirm, t = _stream_response(
                    message_handler, session, query,
                    data_source=file_data_source, output_mode=output_mode,
                    output_dir=session_dir,
                )
                if t:
                    output_text += f"\n### Query: {query}\n" + t + "\n"

                if need_confirm:
                    return got_content, need_confirm  # got_content, need_confirm

            if args.list_generated_files:
                _print_generated_files(file_manager, session.session_id)

            return True, False  # got_content=True, need_confirm=False

        # Pass args as-is since we've attached file_data_source to it
        run_worker_with_handler(args, query_execution_func=file_query_executor)

    # NORMAL SYNCHRONOUS LOGIC
    if file_id:
        # Using Data Center file ID
        print(f"Using Data Center file ID: {file_id}")

        # Create data source with file ID
        file_data_source = DataSource(
            data_source_type="FILE",
            file_id=file_id,
            region_id=config.region
        )
    else:
        # Local file workflow: validate and upload
        if not Path(file_path).exists():
            print(f"Error: File not found: {file_path}", file=sys.stderr)
            sys.exit(1)

        # Validate file type
        if not file_manager.is_supported_file(file_path):
            supported = ".csv, .xlsx, .xls"
            print(f"Error: Unsupported file type. Supported formats: {supported}", file=sys.stderr)
            sys.exit(1)

        # Upload file
        print(f"Uploading file: {file_path}")
        try:
            file_info = file_manager.upload_file(file_path)
        except Exception as e:
            print(f"Upload failed: {e}", file=sys.stderr)
            sys.exit(1)

        print(f"File uploaded successfully!")
        print(f"  File ID : {file_info.file_id}")
        print(f"  Filename: {file_info.filename}")
        print(f"  Size    : {file_info.size} bytes")

        # Create data source for uploaded file
        file_data_source = DataSource(
            data_source_type="FILE",
            file_id=file_info.file_id,
        )

    # Create session
    session_mode = args.session_mode.upper()
    mode_desc = {
        "ASK_DATA": "ASK_DATA mode",
        "ANALYSIS": "ANALYSIS mode (recommended for file analysis)",
        "INSIGHT": "INSIGHT mode",
        "CLAW": "CLAW mode (agentic)",
    }.get(session_mode, session_mode)

    print(f"\nCreating session: {mode_desc}...")
    print(f"  Region: {config.region}")
    enable_search = getattr(args, 'enable_search', False)
    workspace_id = getattr(args, 'workspace_id', None)
    custom_agent_id = getattr(args, 'custom_agent_id', None)
    # For file analysis, pass file_id to create_or_reuse so the session is bound to the file
    session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_data_source.file_id if 'file_data_source' in locals() else None, workspace_id=workspace_id, custom_agent_id=custom_agent_id)
    print(f"Session ready: {session.session_id}")
    print(f"\n💡 Tip: To continue this session later, use: python3 data_agent_cli.py attach --session-id {session.session_id}")

    # Get output mode
    output_mode = getattr(args, "output", "summary")
    session_dir = Path(f"sessions/{session.session_id}")
    session_dir.mkdir(parents=True, exist_ok=True)

    # Initialize structured logging for sync mode
    from cli.streaming import init_structured_logging, close_structured_logging
    init_structured_logging(session_dir)

    try:
        # Determine queries to execute
        if args.query:
            queries = [args.query]
        else:
            # Default preset analysis questions
            queries = [
                "请分析上传文件的数据结构",
                "数据的关键统计指标和分布情况是什么？",
                "数据中是否存在异常值或离群点？",
            ]

        # Execute queries with the file data source
        print(f"\n{'=' * 60}")
        print("File Analysis")
        print("=" * 60)

        output_text = ""
        for query in queries:
            print(f"\nQuery: {query}")
            print("-" * 50)
            try:
                got_content, need_confirm, t = _stream_response(
                    message_handler, session, query,
                    data_source=file_data_source, output_mode=output_mode,
                    output_dir=session_dir,
                )
                if t:
                    output_text += f"\n### Query: {query}\n" + t + "\n"

                if not got_content:
                    print("(No response received, please retry)")
                elif need_confirm:
                    print("\n⚠️  Agent has created an execution plan. User confirmation required.")
                    print(f"   To continue: python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'User Input' ")
                    if output_text:
                        with open(session_dir / "output.md", "w", encoding="utf-8") as f:
                            f.write(output_text)
                        with open(session_dir / "result.json", "w", encoding="utf-8") as f:
                            json.dump({"status": "waiting_input", "output_file": "output.md"}, f)
                    return
            except Exception as e:
                print(f"Request failed: {e}")

        # List generated files if requested
        if args.list_generated_files:
            _print_generated_files(file_manager, session.session_id)

    finally:
        # Close structured logging
        close_structured_logging()

    # Write result status
    with open(session_dir / "result.json", "w", encoding="utf-8") as f:
        json.dump({"status": "completed"}, f)


def _print_generated_files(file_manager: FileManager, session_id: str) -> None:
    """Print list of files generated by the Agent."""
    print(f"\n{'=' * 60}")
    print("Generated Files")
    print("=" * 60)
    try:
        generated = file_manager.list_files(session_id)
        if generated:
            for f in generated:
                print(f"  - {f.filename} ({f.file_type}, {f.size} bytes)")
                if f.download_url:
                    print(f"    Download: {f.download_url}")
        else:
            print("  No generated files.")
    except Exception as e:
        print(f"  Failed to get file list: {e}")