#!/usr/bin/env python3
import os
import re
import sys
import json
import glob
from datetime import datetime, timezone
import xml.etree.ElementTree as ET

# Configuration
TEST_SRC_DIR = "app/src/androidTest/java/com/yorvana/regression"
BENCHMARK_SRC_DIR = "macrobenchmark/src/main/java/com/yorvana/macrobenchmark"
HISTORY_JSON = "docs/regression-history.json"
DASHBOARD_MD = "docs/regression-dashboard.md"
ADDITIONAL_OUTPUT_DIR = "app/build/outputs/managed_device_android_test_additional_output"
XML_REPORTS_DIR_PATTERN = "**/build/outputs/androidTest-results/managedDevice/**/*.xml"

# Scenarios dictionary to store parsed test info from sources
# Key: fqcn#method_name, Value: {scenario_id, tier, quarantined_metadata}
SCENARIOS_MAP = {}
SCENARIO_IDS_SEEN = set()
QUARANTINED_INFO = {} # Key: scenario_id, Value: metadata dict
SCENARIO_ID_TO_TIER = {} # Key: scenario_id, Value: tier

def parse_kotlin_sources():
    """Scans Kotlin test files to build the mapping of test methods to scenario IDs and quarantine status."""
    paths = []
    for root, _, files in os.walk(TEST_SRC_DIR):
        for f in files:
            if f.endswith(".kt"):
                paths.append(os.path.join(root, f))
    for root, _, files in os.walk(BENCHMARK_SRC_DIR):
        for f in files:
            if f.endswith(".kt"):
                paths.append(os.path.join(root, f))

    for path in paths:
        with open(path, "r", encoding="utf-8") as f:
            content = f.read()

        package_match = re.search(r"package\s+([a-zA-Z0-9_.]+)", content)
        package = package_match.group(1) if package_match else ""

        # Parse classes (including optional annotations before class declaration)
        class_matches = list(re.finditer(r"(?:^|[\r\n])\s*((?:@[A-Za-z0-9_.\(\)\s=,\":]*?\s*)*)(?:open\s+|abstract\s+|sealed\s+|value\s+|inner\s+)?class\s+([A-Za-z0-9_]+)", content))
        for idx, class_match in enumerate(class_matches):
            class_annotations_raw = class_match.group(1) or ""
            class_name = class_match.group(2)
            fqcn = f"{package}.{class_name}" if package else class_name

            # Determine class-level tier
            class_tier = None
            if "RegressionFull" in class_annotations_raw:
                class_tier = "RegressionFull"
            elif "Regression" in class_annotations_raw:
                class_tier = "Regression"
            elif "BillingSandbox" in class_annotations_raw:
                class_tier = "BillingSandbox"

            # Slice content to this class block (from class start to next class start or end of file)
            start_pos = class_match.start()
            end_pos = class_matches[idx + 1].start() if idx + 1 < len(class_matches) else len(content)
            class_body = content[start_pos:end_pos]

            # Parse methods inside class block
            method_matches = re.finditer(
                r"((?:@[A-Za-z0-9_.\(\)\s=,\"-]*?\s*)*)fun\s+([A-Za-z0-9_]+)\s*\(",
                class_body
            )
            for method_match in method_matches:
                annotations_raw = method_match.group(1) or ""
                method_name = method_match.group(2)

                # Skip non-test functions
                is_test = "@Test" in annotations_raw or "@Benchmark" in annotations_raw
                if not is_test or method_name in ("setUp", "tearDown", "setup"):
                    continue

                # Determine method tier (default to class tier)
                method_tier = class_tier
                if "RegressionFull" in annotations_raw:
                    method_tier = "RegressionFull"
                elif "Regression" in annotations_raw:
                    method_tier = "Regression"
                elif "BillingSandbox" in annotations_raw:
                    method_tier = "BillingSandbox"

                # Parse ScenarioId
                scenario_id = None
                sid_match = re.search(r'@ScenarioId\(\s*"([^"]+)"\s*\)', annotations_raw)
                if sid_match:
                    scenario_id = sid_match.group(1)
                else:
                    # Apply naming convention
                    # Match rV01_ -> R-V01, rA01_ -> R-A01, M03_ -> M-03
                    conv_match = re.match(r"^[rR]([a-zA-Z])(\d{2})_", method_name)
                    if conv_match:
                        scenario_id = f"R-{conv_match.group(1).upper()}{conv_match.group(2)}"
                    else:
                        m_match = re.match(r"^[mM](\d{2})_", method_name)
                        if m_match:
                            scenario_id = f"M-{m_match.group(1)}"
                        else:
                            scenario_id = method_name # Fallback to method name

                # Parse Quarantine status (order-independent)
                quarantined = False
                q_meta = {}
                q_match = re.search(r'@Quarantined\(([^)]+)\)', annotations_raw)
                if q_match:
                    quarantined = True
                    q_args = q_match.group(1)
                    since_m = re.search(r'since\s*=\s*"([^"]+)"', q_args)
                    issue_m = re.search(r'issue\s*=\s*(\d+)', q_args)
                    reason_m = re.search(r'reason\s*=\s*"([^"]+)"', q_args)
                    q_meta = {
                        "since": since_m.group(1) if since_m else "unknown",
                        "issue": int(issue_m.group(1)) if issue_m else 0,
                        "reason": reason_m.group(1) if reason_m else "quarantined"
                    }
                elif "@Quarantined" in annotations_raw or "Quarantined" in class_annotations_raw:
                    quarantined = True
                    # Check class level quarantined metadata if method has none
                    qc_match = re.search(r'@Quarantined\(([^)]+)\)', class_annotations_raw)
                    if qc_match:
                        qc_args = qc_match.group(1)
                        since_m = re.search(r'since\s*=\s*"([^"]+)"', qc_args)
                        issue_m = re.search(r'issue\s*=\s*(\d+)', qc_args)
                        reason_m = re.search(r'reason\s*=\s*"([^"]+)"', qc_args)
                        q_meta = {
                            "since": since_m.group(1) if since_m else "unknown",
                            "issue": int(issue_m.group(1)) if issue_m else 0,
                            "reason": reason_m.group(1) if reason_m else "quarantined"
                        }
                    else:
                        q_meta = {"since": "unknown", "issue": 0, "reason": "quarantined"}

                key = f"{fqcn}#{method_name}"
                SCENARIOS_MAP[key] = {
                    "scenario_id": scenario_id,
                    "tier": method_tier or "Regression",
                    "quarantined": quarantined,
                    "q_meta": q_meta
                }
                SCENARIO_IDS_SEEN.add(scenario_id)

                tier = method_tier or "Regression"
                if scenario_id in SCENARIO_ID_TO_TIER:
                    existing_tier = SCENARIO_ID_TO_TIER[scenario_id]
                    existing_methods = [k for k, v in SCENARIOS_MAP.items() if v["scenario_id"] == scenario_id]
                    print(f"WARNING: Duplicate scenario ID '{scenario_id}' found! "
                          f"Already mapped to: {existing_methods}. Now also mapped to: '{key}'.", file=sys.stderr)
                    if existing_tier != tier:
                        # Prefer more restrictive tier
                        tier_priority = {"BillingSandbox": 3, "Regression": 2, "RegressionFull": 1}
                        if tier_priority.get(tier, 0) > tier_priority.get(existing_tier, 0):
                            SCENARIO_ID_TO_TIER[scenario_id] = tier
                else:
                    SCENARIO_ID_TO_TIER[scenario_id] = tier

                if quarantined:
                    QUARANTINED_INFO[scenario_id] = {
                        "q_meta": q_meta,
                        "tier": method_tier or "Regression",
                        "method": key
                    }

def check_flake_file(classname, methodname):
    """Checks if a flake marker file exists for the given test case within the host additional outputs directory."""
    test_key = f"{classname}#{methodname}"
    sanitized_key = test_key.replace("[", "_").replace("]", "_").replace("=", "_").replace(" ", "_").replace("#", "_")
    # Search both the standard additional output dir (local) and CI download dir (.tmp-reports)
    paths = [
        f"{ADDITIONAL_OUTPUT_DIR}/**/flake-{sanitized_key}.txt",
        f".tmp-reports/**/flake-{sanitized_key}.txt",
        f"tmp-reports/**/flake-{sanitized_key}.txt"
    ]
    for path in paths:
        files = glob.glob(path, recursive=True)
        if len(files) > 0:
            return True
    return False

def parse_xml_reports():
    """Parses JUnit XML report files to extract test outcomes."""
    results = {}
    xml_files = (
        glob.glob("**/regression-core-preserved/**/*.xml", recursive=True)
        + glob.glob("**/managedDevice/**/*.xml", recursive=True)
    )
    if not xml_files:
        xml_files = glob.glob("**/androidTest-results/**/*.xml", recursive=True)

    for xml_file in xml_files:
        try:
            tree = ET.parse(xml_file)  # nosemgrep
            root = tree.getroot()
            for testcase in root.findall(".//testcase"):
                classname = testcase.get("classname")
                name = testcase.get("name")
                if not classname or not name:
                    continue

                # Strip parameterized variant suffix e.g., test[0] -> test
                base_name = re.sub(r'\[.*\]$', '', name)
                key = f"{classname}#{base_name}"
                
                # Suffix for variant
                variant_suffix = ""
                var_match = re.search(r'(\[.*\])$', name)
                if var_match:
                    variant_suffix = var_match.group(1)

                scenario_info = SCENARIOS_MAP.get(key)
                if not scenario_info:
                    # Fallback lookup: match method name suffix if classname has suffix/differs
                    fallback_key = [k for k in SCENARIOS_MAP.keys() if k.endswith(f"#{base_name}")]
                    if fallback_key:
                        scenario_info = SCENARIOS_MAP[fallback_key[0]]

                if scenario_info:
                    scenario_id = scenario_info["scenario_id"]
                    # Suffix scenario ID for parameterized variant
                    final_scenario_id = f"{scenario_id}{variant_suffix}"
                else:
                    # Fallback to base name if not found in sources
                    conv_match = re.match(r"^[rR]([a-zA-Z])(\d{2})_", base_name)
                    if conv_match:
                        scenario_id = f"R-{conv_match.group(1).upper()}{conv_match.group(2)}"
                    else:
                        scenario_id = base_name
                    final_scenario_id = f"{scenario_id}{variant_suffix}"

                # Determine outcome
                outcome = "pass"
                if testcase.find("failure") is not None or testcase.find("error") is not None:
                    outcome = "fail"
                elif testcase.find("skipped") is not None:
                    outcome = "skip"
                else:
                    # Check if it was flaky (passed but has a flake file)
                    if check_flake_file(classname, base_name):
                        outcome = "flake"

                if final_scenario_id not in results:
                    results[final_scenario_id] = outcome
        except Exception as e:
            print(f"Error parsing XML file {xml_file}: {e}", file=sys.stderr)
    return results

def main():
    parse_kotlin_sources()
    current_results = parse_xml_reports()
    infra_failure = len(current_results) == 0
    if infra_failure:
        print("WARNING: No test results parsed! Recording as infrastructure failure.")
        try:
            os.makedirs("tmp-reports", exist_ok=True)
            with open("tmp-reports/infra_failed.txt", "w") as f:
                f.write("true")
        except Exception as e:
            print(f"Error writing infra failure marker: {e}", file=sys.stderr)

    # Load history
    history = {"runs": []}
    if os.path.exists(HISTORY_JSON):
        try:
            with open(HISTORY_JSON, "r") as f:
                history = json.load(f)
        except Exception as e:
            print(f"Error reading history JSON: {e}", file=sys.stderr)

    # Append current run
    current_run = {
        "timestamp": datetime.now(timezone.utc).isoformat() + "Z",
        "results": current_results,
        "infra_failure": infra_failure
    }
    history["runs"].append(current_run)

    # Keep only last 100 runs
    history["runs"] = history["runs"][-100:]

    # Write history JSON
    os.makedirs(os.path.dirname(HISTORY_JSON), exist_ok=True)
    with open(HISTORY_JSON, "w") as f:
        json.dump(history, f, indent=2)

    # Compute Stats
    runs_to_process = [run for run in history["runs"] if not run["infra_failure"]]
    num_runs = len(runs_to_process)

    scenario_stats = {} # Key: base_scenario_id, Value: {passes, flakes, fails, skips, consecutive_green, consecutive_passes}
    
    # Initialize all seen scenarios
    all_known_ids = set(SCENARIO_IDS_SEEN)
    for run in runs_to_process:
        for k in run["results"].keys():
            base_id = re.sub(r'\[.*\]$', '', k)
            all_known_ids.add(base_id)

    for base_id in all_known_ids:
        scenario_stats[base_id] = {
            "pass_count": 0,
            "flake_count": 0,
            "fail_count": 0,
            "skip_count": 0,
            "consecutive_green": 0,
            "consecutive_passes": 0,
            "runs_evaluated": 0,
            "outcomes_history": [] # Order: oldest to newest
        }

    # Evaluate history
    for run in runs_to_process:
        # For each run, roll up results by base_scenario_id
        run_rolled_up = {}
        for final_id, outcome in run["results"].items():
            base_id = re.sub(r'\[.*\]$', '', final_id)
            if base_id not in run_rolled_up:
                run_rolled_up[base_id] = []
            run_rolled_up[base_id].append(outcome)

        for base_id in all_known_ids:
            stats = scenario_stats[base_id]
            outcomes = run_rolled_up.get(base_id, [])
            if not outcomes:
                # Test didn't run in this run (e.g. was excluded from this tier)
                stats["outcomes_history"].append("missing")
                continue

            stats["runs_evaluated"] += 1
            # Roll up logic: if any variant failed -> fail, if any flaked -> flake, if all skipped -> skip, else pass
            if "fail" in outcomes:
                stats["fail_count"] += 1
                stats["outcomes_history"].append("fail")
            elif "flake" in outcomes:
                stats["flake_count"] += 1
                stats["outcomes_history"].append("flake")
            elif all(o == "skip" for o in outcomes):
                stats["skip_count"] += 1
                stats["outcomes_history"].append("skip")
            else:
                stats["pass_count"] += 1
                stats["outcomes_history"].append("pass")

    # Compute streaks according to promotion playbook
    for base_id, stats in scenario_stats.items():
        consec_green = 0
        consec_passes = 0
        for outcome in stats["outcomes_history"]:
            # consecutive_green (for de-quarantine): resets on fail. Increments on pass or flake. Neutral on skip or missing.
            if outcome == "fail":
                consec_green = 0
            elif outcome in ("pass", "flake"):
                consec_green += 1

            # consecutive_passes (for promotion): resets on fail or flake. Increments on pass. Neutral on skip or missing.
            if outcome in ("fail", "flake"):
                consec_passes = 0
            elif outcome == "pass":
                consec_passes += 1

        stats["consecutive_green"] = consec_green
        stats["consecutive_passes"] = consec_passes

    # Compute overall rates
    total_non_infra_runs = num_runs
    total_infra_runs = sum(1 for run in history["runs"] if run["infra_failure"])

    # Check if last run was an infra failure
    last_run_infra_failure = history["runs"][-1]["infra_failure"] if history["runs"] else False

    # Generate Markdown Dashboard
    md = []
    md.append("# Regression Test Dashboard & Flake Audit")
    
    if last_run_infra_failure:
        md.append("\n> [!CAUTION]")
        md.append("> **Infrastructure Failure Detected**: The most recent scheduled run failed due to system/infrastructure issues (e.g., emulator or runner crash) and did not produce test results.")

    md.append(f"\nGenerated on: **{datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}**")
    md.append(f"\nTotal historical runs tracked: **{len(history['runs'])}** | Non-infra runs: **{total_non_infra_runs}** | Infra noise failures: **{total_infra_runs}**")

    # Overall Summary Table
    md.append("\n## Suite Health Summary")
    md.append("| Metric | Value | Description |")
    md.append("|---|---|---|")
    
    # Calculate global flake budget
    total_evals = 0
    total_flakes = 0
    total_fails = 0
    for base_id, stats in scenario_stats.items():
        total_evals += stats["pass_count"] + stats["flake_count"] + stats["fail_count"]
        total_flakes += stats["flake_count"]
        total_fails += stats["fail_count"]
    
    overall_pass_rate = 100.0 * (total_evals - total_fails) / max(total_evals, 1)
    overall_flake_rate = 100.0 * total_flakes / max(total_evals, 1)

    md.append(f"| **Overall Pass Rate** | {overall_pass_rate:.2f}% | Evaluated tests passing or flaking successfully |")
    md.append(f"| **Suite Flake Rate** | {overall_flake_rate:.2f}% | Tests requiring 1 or more retries (Budget: ≤1.0%) |")
    md.append(f"| **Infra Noise Runs** | {total_infra_runs} | Build runs failing due to runner/emulator infrastructure |")

    # Quarantined Tests Section
    md.append("\n## Quarantined Tests")
    md.append("> [!WARNING]")
    md.append("> Quarantined tests are excluded from gating PR builds to protect developer velocity, but continue to run in nightly regression runs.")
    md.append("\n| Scenario ID | Test Name / File | Tier | Since | Issue | Reason | Pass Rate (Last 100) | Streak | Eligible for Promotion |")
    md.append("|---|---|---|---|---|---|---|---|---|")

    quarantine_count = 0
    repo = os.environ.get("GITHUB_REPOSITORY", "yorvana/android")
    for base_id, info in QUARANTINED_INFO.items():
        quarantine_count += 1
        stats = scenario_stats.get(base_id, {"pass_count": 0, "flake_count": 0, "fail_count": 0, "runs_evaluated": 0, "consecutive_green": 0})
        total_valid = stats["pass_count"] + stats["flake_count"] + stats["fail_count"]
        pass_rate_str = "0%"
        if total_valid > 0:
            pass_rate_str = f"{100.0 * (stats['pass_count'] + stats['flake_count']) / total_valid:.1f}%"
        
        consec = stats["consecutive_green"]
        eligible = "✅ Yes (≥30 Green Runs)" if consec >= 30 else f"❌ No ({consec}/30)"
        
        since = info["q_meta"].get("since", "unknown")
        issue = info["q_meta"].get("issue", 0)
        reason = info["q_meta"].get("reason", "N/A")
        issue_url = f"https://github.com/{repo}/issues/{issue}"
        
        md.append(f"| **{base_id}** | `{info['method']}` | `{info['tier']}` | {since} | [#{issue}]({issue_url}) | {reason} | {pass_rate_str} ({stats['pass_count'] + stats['flake_count']}/{total_valid}) | {consec} | {eligible} |")

    if quarantine_count == 0:
        md.append("| - | No tests currently quarantined | - | - | - | - | - | - | - |")

    # Active Tests Section
    md.append("\n## Active Regression Tests")
    md.append("| Scenario ID | Tier | Pass Rate | Flake Rate | Consecutive Passes | Status |")
    md.append("|---|---|---|---|---|---|")

    active_scenarios = sorted([sid for sid in all_known_ids if sid not in QUARANTINED_INFO and sid in SCENARIO_IDS_SEEN])
    for base_id in active_scenarios:
        stats = scenario_stats[base_id]
        total_valid = stats["pass_count"] + stats["flake_count"] + stats["fail_count"]
        pass_rate_str = "N/A"
        flake_rate_str = "N/A"
        status = "🟢 Healthy"
        
        if total_valid > 0:
            pass_rate = 100.0 * (stats['pass_count'] + stats['flake_count']) / total_valid
            flake_rate = 100.0 * stats['flake_count'] / total_valid
            pass_rate_str = f"{pass_rate:.1f}%"
            flake_rate_str = f"{flake_rate:.1f}%"
            if flake_rate > 5.0:
                status = "🟡 Flaky (Needs Quarantine)"
            if pass_rate < 95.0:
                status = "🔴 Failing (Needs Attention)"

        tier = SCENARIO_ID_TO_TIER.get(base_id, "Regression")

        md.append(f"| **{base_id}** | `{tier}` | {pass_rate_str} | {flake_rate_str} | {stats['consecutive_passes']} | {status} |")

    with open(DASHBOARD_MD, "w") as f:
        f.write("\n".join(md))

    print(f"Aggregation complete. Updated {HISTORY_JSON} and {DASHBOARD_MD}.")

if __name__ == "__main__":
    main()