import sys
import os

if __name__ == "__main__":
    # Add project root to sys.path to allow imports when running as a script
    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
    from tools.trace_view.perfetto_tool import SliceInfoTool, SliceFinderTool, SqlQueryTool
    from tools.trace_view.connection_manager import ConnectionManager
else:
    from .perfetto_tool import SliceInfoTool, SliceFinderTool, SqlQueryTool
    from .connection_manager import ConnectionManager


class TraceViewAnalyzeTool:
    """Tool for analyzing trace data (specifically `trace_view.json`)."""

    def __init__(self):
        connection_manager = ConnectionManager()
        self.slice_finder_tool = SliceFinderTool(connection_manager)
        self.sql_query_tool = SqlQueryTool(connection_manager)

    def find_slices(
        self,
        trace_path: str,
        pattern: str,
        process_name: str | None = None,
        match_mode: str = "contains",
        limit: int = 100,
        main_thread_only: bool = False,
        time_range: dict | None = None,
    ) -> str:
        """
        Discover slices by name in a `trace_view.json` file (Chrome Trace Event format).

        SCOPE:
        - Strictly for `trace_view.json` timeline files (generated by Ascend Profiler/MSProf).
        - Used to search for specific operators (e.g., 'MatMul'), Python functions, or other trace events.

        WHY USE THIS:
        - Explore unknown slice names and hot paths fast (no manual SQL).
        - See frequency and duration stats (min/avg/max and p50/p90/p99 when available) per slice name.
        - Get linkable examples (id, ts, dur, track_id) to jump in UI or correlate with other tools.
        - Filter by process, main thread, and time range to narrow investigations.

        PARAMETERS:
        - pattern: String to match against slice names.
        - match_mode: 'contains' (default), 'exact', or 'glob'.
        - process_name: Optional filter; supports '*' wildcard.
        - main_thread_only: Limit to process main threads.
        - time_range: {'start_ms': X, 'end_ms': Y}.
        - limit: Max example slices to return (default 100).

        OUTPUT:
        - aggregates: Per-slice-name counts and duration stats (min/avg/max, p50/p90/p99 when available).
        - examples: Top slices by duration with thread/process context and track id for linking.
        - notes: Capability or fallback notices.
        """
        return self.slice_finder_tool.find_slices(
            trace_path,
            pattern,
            process_name,
            match_mode,
            limit,
            main_thread_only,
            time_range,
        )

    def execute_sql_query(self, trace_path: str, sql_query: str, process_name: str | None = None) -> str:
        """
        Execute PerfettoSQL scripts on a `trace_view.json` file.

        SCOPE:
        - Strictly for `trace_view.json` timeline files (Chrome Trace Event format).
        - Allows advanced custom analysis not covered by other methods.

        USE THIS WHEN: Other tools don't provide what you need, you need complex filtering/joins, 
        or you want to correlate data across multiple tables. This is your power tool for custom 
        analysis - use it when pre-built tools are too limiting.

        CAPABILITIES: Full SQL access to all trace tables including:
        - slice: All trace slices with timing
        - thread/process: Thread and process metadata
        - counter: Performance counters over time
        - sched_slice: CPU scheduling information
        - flow: Flow events connecting slices

        SECURITY: Accepts full PerfettoSQL/SQLite scripts. No automatic LIMIT is applied; large
        queries may return many rows. The script is executed verbatim by TraceProcessor.

        COMMON PATTERNS:
        - Duration analysis: "SELECT name, dur/1e6 as ms FROM slice WHERE dur > 10e6"
        - Aggregation: "SELECT name, COUNT(*), AVG(dur)/1e6 FROM slice GROUP BY name"
        - Time filtering: "SELECT * FROM slice WHERE ts BETWEEN 1e9 AND 2e9"
        - Process filtering: "SELECT * FROM thread WHERE upid IN (SELECT upid FROM process WHERE name LIKE '%python%')"

        POWER USER TIP: Use `INCLUDE PERFETTO MODULE ...` statements to load standard library
        modules. You can also use `CREATE PERFETTO TABLE`/
        `VIEW`/`FUNCTION`/`MACRO`/`INDEX` where supported by TraceProcessor.

        References:
        - PerfettoSQL Syntax: https://perfetto.dev/docs/analysis/perfetto-sql-syntax
        - Standard Library (Prelude): https://perfetto.dev/docs/analysis/stdlib-docs#package-prelude
        """
        return self.sql_query_tool.execute_sql_query(trace_path, sql_query, process_name)

    def analyze_overlap(self, trace_path: str) -> str:
        """
        Analyze the 'Overlap Analysis' process to provide a global performance overview.

        USE THIS WHEN: You are analyzing a `trace_view.json` file for the first time. 
        This is the primary method to obtain global performance information, specifically 
        the proportion of time spent on Computing, Communication, and Scheduling.

        CAPABILITIES: Calculates detailed statistics for the 'Overlap Analysis' process, including:
        - Computing: Time spent on actual computation.
        - Communication: Total communication time.
        - Communication(Not Overlapped): Communication time exposed (not overlapped with computation).
        - Free: Scheduling wait time or idle time.

        OUTPUT:
        Returns a JSON string with the total duration (ms) and percentage breakdown for each state.
        This helps identify if the workload is Compute-bound, Communication-bound, or has Scheduling overhead.
        """
        query = """
        SELECT
            s.name,
            CAST(SUM(s.dur) AS FLOAT) / 1e6 AS duration_ms
        FROM slice s
        JOIN track t ON s.track_id = t.id
        LEFT JOIN process_track pt ON t.id = pt.id
        LEFT JOIN process p1 ON pt.upid = p1.upid
        LEFT JOIN thread_track tt ON t.id = tt.id
        LEFT JOIN thread th ON tt.utid = th.utid
        LEFT JOIN process p2 ON th.upid = p2.upid
        WHERE
            (p1.name = 'Overlap Analysis' OR p2.name = 'Overlap Analysis')
            AND s.name IN ('Computing', 'Communication', 'Communication(Not Overlapped)', 'Free')
        GROUP BY s.name
        """
        
        # Execute query
        import json
        result_json = self.sql_query_tool.execute_sql_query(trace_path, query)
        result_data = json.loads(result_json)
        
        if not result_data.get("success", False):
            return result_json

        rows = result_data.get("result", {}).get("rows", [])
        
        # Calculate totals and percentages
        total_duration = sum(row.get("duration_ms", 0) for row in rows)
        
        analysis_result = {
            "process": "Overlap Analysis",
            "total_duration_ms": total_duration,
            "breakdown": []
        }
        
        for row in rows:
            name = row.get("name")
            duration = row.get("duration_ms", 0)
            percentage = (duration / total_duration * 100) if total_duration > 0 else 0
            analysis_result["breakdown"].append({
                "name": name,
                "duration_ms": duration,
                "percentage": f"{percentage:.2f}%"
            })
            
        return json.dumps(analysis_result, indent=2)


if __name__ == "__main__":
    tool = TraceViewAnalyzeTool()
    trace_path = "/Users/weizhang/Downloads/kv_cache_type_page_seqlen_1024_bs_1_profile_count_0/g340-cd51-4900-7a3f-dc8e-d804-1072_348921_20251105072538961_ascend_pt/ASCEND_PROFILER_OUTPUT/trace_view.json"
    
    print("--- Overlap Analysis ---")
    res = tool.analyze_overlap(trace_path)
    print(res)