"""
Report builder for generating reports based on query filters.

This module provides the main interface for generating reports using the
flexible query filter system.
"""

from typing import Dict, List, Optional

from src.database.db import get_db
from src.models.models import TestResult, TestTask
from src.services.workload_templates import is_premium_workload_config

from .aggregation import AggregationStrategy, aggregate_runs, group_runs_by_thread_count
from .data_aggregator import extract_cluster_config_from_task, extract_runs_from_result
from .models import ClusterConfig, SPUReportData, WorkflowData, WorkloadRunData
from .query_executor import query_tasks_by_filter
from .query_filter import QueryFilter


def generate_report_from_filter(
    query_filter: QueryFilter,
    aggregation: AggregationStrategy = AggregationStrategy.LATEST,
) -> Optional[SPUReportData]:
    """Generate a performance report based on query filter.

    Args:
        query_filter: Filter conditions for selecting test tasks
        aggregation: Strategy for aggregating multiple matching experiments

    Returns:
        SPUReportData with aggregated results, or None if no matching tasks

    Example:
        >>> # Generate report for AWS G2 machines with >= 64 TiDB cores
        >>> filter = QueryFilter() \\
        ...     .add_condition(FilterDimension.CLOUD_PROVIDER, value="AWS") \\
        ...     .add_condition(FilterDimension.MACHINE_GENERATION, value="G2") \\
        ...     .add_condition(FilterDimension.TIDB_TOTAL_CORES,
        ...                    operator=FilterOperator.GREATER_THAN_OR_EQUAL,
        ...                    value=64)
        >>> report = generate_report_from_filter(filter, AggregationStrategy.AVERAGE)
    """
    # Execute everything within a single database session
    with get_db() as db:
        # Query matching tasks (session-bound)
        tasks = query_tasks_by_filter(query_filter, db_session=db)

        if not tasks:
            return None

        # Extract cluster config from first task
        cluster_config = extract_cluster_config_from_task(tasks[0])

        # Generate SPU identifier
        spu_identifier = _generate_spu_identifier(cluster_config)

        # Extract workflow data from tasks (within session)
        workflows = _extract_workflow_data_in_session(db, tasks, aggregation)

    # Create report
    report = SPUReportData(
        spu_identifier=spu_identifier,
        cluster_config=cluster_config,
        workflows=workflows,
    )

    return report


def generate_comparison_report(
    filter1: QueryFilter,
    filter2: QueryFilter,
    aggregation: AggregationStrategy = AggregationStrategy.LATEST,
) -> tuple[Optional[SPUReportData], Optional[SPUReportData]]:
    """Generate comparison report between two filter conditions.

    Args:
        filter1: First filter condition
        filter2: Second filter condition
        aggregation: Strategy for aggregating multiple matching experiments

    Returns:
        Tuple of (report1, report2), either may be None if no matching tasks
    """
    report1 = generate_report_from_filter(filter1, aggregation)
    report2 = generate_report_from_filter(filter2, aggregation)

    return report1, report2


def _extract_cluster_config(task: TestTask) -> ClusterConfig:
    """Extract ClusterConfig from TestTask.

    Args:
        task: TestTask from database

    Returns:
        ClusterConfig object
    """
    return ClusterConfig(
        cloud_provider=task.cloud_provider,
        region=task.region,
        cluster_version=task.cluster_version,
        tidb_spu=task.tidb_spu,
        tidb_count=task.tidb_count,
        tidb_config=task.tidb_config,
        tikv_spu=task.tikv_spu,
        tikv_count=task.tikv_count,
        tikv_storage_type=task.tikv_storage_type,
        tikv_storage_size_gi=task.tikv_storage_size_gi,
        tikv_config=task.tikv_config,
        pd_spu=task.pd_spu,
        pd_count=task.pd_count,
        pd_config=task.pd_config,
        tiflash_spu=task.tiflash_spu,
        tiflash_count=task.tiflash_count,
        tiflash_storage_type=task.tiflash_storage_type,
        tiflash_storage_size_gi=task.tiflash_storage_size_gi,
        tiflash_config=task.tiflash_config,
        tiproxy_spu=task.tiproxy_spu,
        tiproxy_count=task.tiproxy_count,
        tiproxy_config=task.tiproxy_config,
        workload_config=task.workload_config,
    )


def _generate_spu_identifier(config: ClusterConfig) -> str:
    """Generate SPU identifier string from cluster config.

    Args:
        config: Cluster configuration

    Returns:
        String identifier like "2C8Gx1_2C8Gx3_AWS_us-west-2_v8.5.4"
        For Premium: "5000_20000_AWS_us-west-2_v8.5.7"
    """
    # Premium: use RCU min_max instead of SPU
    wlc = config.workload_config or {}
    if is_premium_workload_config(wlc):
        rcu = wlc.get("rcu", {})
        min_rcu = rcu.get("min", "")
        max_rcu = rcu.get("max", "")
        tidb_part = f"{min_rcu}_{max_rcu}"
    else:
        tidb_part = f"{config.tidb_spu}x{config.tidb_count}"
        tikv_part = f"{config.tikv_spu}x{config.tikv_count}"

    cloud_part = config.cloud_provider
    region_part = config.region
    version_part = config.cluster_version

    if is_premium_workload_config(wlc):
        return f"{tidb_part}_{cloud_part}_{region_part}_{version_part}"
    else:
        return f"{tidb_part}_{tikv_part}_{cloud_part}_{region_part}_{version_part}"


def _extract_workflow_data_in_session(
    db, tasks: List[TestTask], aggregation: AggregationStrategy
) -> List[WorkflowData]:
    """Extract workflow data from tasks within an active session.

    Args:
        db: Active database session
        tasks: List of matching tasks
        aggregation: Aggregation strategy

    Returns:
        List of WorkflowData with aggregated runs
    """
    # Group tasks by workflow template
    workflows_dict: Dict[str, List[TestTask]] = {}

    for task in tasks:
        template = task.workload_template
        if template not in workflows_dict:
            workflows_dict[template] = []
        workflows_dict[template].append(task)

    # Convert to WorkflowData
    workflows = []

    for template, template_tasks in workflows_dict.items():
        # Extract all runs from these tasks
        all_runs = []
        for task in template_tasks:
            # Get test result for this task
            result = (
                db.query(TestResult)
                .filter_by(task_id=task.id)
                .filter(TestResult.deleted == False)
                .order_by(TestResult.id.desc())
                .first()
            )

            if not result or not result.result_data:
                continue

            # Extract runs
            runs = extract_runs_from_result(result, task)
            all_runs.extend(runs)

        # Apply aggregation based on strategy
        if aggregation == AggregationStrategy.ALL:
            # Keep all runs without aggregation
            aggregated_runs = all_runs
        else:
            # Group runs by thread count and aggregate
            grouped_runs = group_runs_by_thread_count(all_runs)
            aggregated_runs = []

            for thread_count, runs in sorted(grouped_runs.items()):
                aggregated_run = aggregate_runs(runs, aggregation, template)
                if aggregated_run:
                    aggregated_runs.append(aggregated_run)

        # Create WorkflowData
        if template_tasks:
            workflow = WorkflowData(
                workflow_type=template,
                workflow_execution_id=f"aggregated_{len(template_tasks)}_tasks",
                task_id=template_tasks[0].id,  # Use first task ID
                status="SUCCESS" if aggregated_runs else "PARTIAL",
                runs=aggregated_runs,
            )
            workflows.append(workflow)

    return workflows


def _extract_workflow_data(
    tasks: List[TestTask], aggregation: AggregationStrategy
) -> List[WorkflowData]:
    """Extract workflow data from tasks and apply aggregation.

    Args:
        tasks: List of matching tasks
        aggregation: Aggregation strategy

    Returns:
        List of WorkflowData with aggregated runs
    """
    # Group tasks by workflow template
    workflows_dict: Dict[str, List[TestTask]] = {}

    for task in tasks:
        template = task.workload_template
        if template not in workflows_dict:
            workflows_dict[template] = []
        workflows_dict[template].append(task)

    # Convert to WorkflowData
    workflows = []

    for template, template_tasks in workflows_dict.items():
        # Extract all runs from these tasks
        all_runs = []
        for task in template_tasks:
            runs = _extract_runs_from_task(task)
            all_runs.extend(runs)

        # Apply aggregation based on strategy
        if aggregation == AggregationStrategy.ALL:
            # Keep all runs without aggregation
            aggregated_runs = all_runs
        else:
            # Group runs by thread count and aggregate
            grouped_runs = group_runs_by_thread_count(all_runs)
            aggregated_runs = []

            for thread_count, runs in sorted(grouped_runs.items()):
                aggregated_run = aggregate_runs(runs, aggregation, template)
                if aggregated_run:
                    aggregated_runs.append(aggregated_run)

        # Create WorkflowData
        if template_tasks:
            workflow = WorkflowData(
                workflow_type=template,
                workflow_execution_id=f"aggregated_{len(template_tasks)}_tasks",
                task_id=template_tasks[0].id,  # Use first task ID
                status="SUCCESS" if aggregated_runs else "PARTIAL",
                runs=aggregated_runs,
            )
            workflows.append(workflow)

    return workflows


def _extract_runs_from_task(task: TestTask) -> List[WorkloadRunData]:
    """Extract WorkloadRunData from a TestTask.

    This function queries the TestResult table and extracts run data.

    Args:
        task: TestTask from database

    Returns:
        List of WorkloadRunData objects
    """
    runs = []

    with get_db() as db:
        # Get the latest result for this task
        result = (
            db.query(TestResult)
            .filter_by(task_id=task.id)
            .filter(TestResult.deleted == False)
            .order_by(TestResult.id.desc())
            .first()
        )

        if not result or not result.result_data:
            return runs

        # Use existing extract_runs_from_result function
        runs = extract_runs_from_result(result, task)

    return runs


if __name__ == "__main__":
    # Test report builder
    from .query_filter import FilterDimension

    print("Testing Report Builder:\n")

    # Test 1: Generate report for AWS G2 machines
    filter1 = (
        QueryFilter()
        .add_condition(FilterDimension.CLOUD_PROVIDER, value="AWS")
        .add_condition(FilterDimension.MACHINE_GENERATION, value="G2")
    )

    print("Test 1: Generate report for AWS G2 machines")
    print(f"  Filter: {filter1.to_description()}")
    report1 = generate_report_from_filter(filter1, AggregationStrategy.LATEST)

    if report1:
        print(f"  ✓ Report generated: {report1.spu_identifier}")
        print(
            f"  Cluster: {report1.cluster_config.tidb_spu}x{report1.cluster_config.tidb_count}"
        )
        print(f"  Workflows: {len(report1.workflows)}")
        if report1.workflows:
            print(f"  Sample workflow: {report1.workflows[0].workflow_type}")
            print(f"  Runs in first workflow: {len(report1.workflows[0].runs)}")
    else:
        print("  ✗ No report generated")

    print()

    # Test 2: Generate comparison report
    filter2 = (
        QueryFilter()
        .add_condition(FilterDimension.CLOUD_PROVIDER, value="AWS")
        .add_condition(FilterDimension.MACHINE_GENERATION, value="G4")
    )

    print("Test 2: Generate comparison report (G2 vs G4)")
    print(f"  Filter 1: {filter1.to_description()}")
    print(f"  Filter 2: {filter2.to_description()}")

    report_v2, report_v4 = generate_comparison_report(
        filter1, filter2, AggregationStrategy.AVERAGE
    )

    if report_v2 and report_v4:
        print(f"  ✓ Both reports generated")
        print(f"    G2: {report_v2.spu_identifier}")
        print(f"    G4: {report_v4.spu_identifier}")
    else:
        print(f"  Partial: G2={report_v2 is not None}, G4={report_v4 is not None}")

    print("\n✓ Report builder tests completed!")