# ============================================================================
# test_agent_mode.py - Unified Automated Test for feat/agent-mode
# ============================================================================
#
# Single-command execution that:
#   1. Validates environment (Python, cobc, Java, jcl-cobol-data-create path)
#   2. Fixes JCL_ROOT in orchestrator_jcl.py if needed
#   3. Installs Python dependencies if missing
#   4. Runs unit tests (pytest)
#   5. Runs dry-run validation
#   6. Runs full agent mode pipeline with ZAN04MAT
#
# Usage:
#   python scripts/test_agent_mode.py
#   python scripts/test_agent_mode.py --skip-tests
#   python scripts/test_agent_mode.py --skip-pipeline
#   python scripts/test_agent_mode.py --dry-run-only
#
# ============================================================================

import argparse
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path


# ============================================================================
# Constants
# ============================================================================

PROJECT_ROOT = Path(__file__).resolve().parent.parent
JCL_ROOT_CANDIDATES = [
    Path(r"C:\Users\marye\Desktop\2026技术大赛\jcl-cobol-data-create"),
    Path(r"D:\jcl-cobol-data-create"),
]
COBOL_PROJECT = Path(r"C:\Users\marye\Desktop\2026技术大赛\cobol-tna-system")
PROGRAM = "ZAN04MAT"

TEST_INPUTS = {
    "design": COBOL_PROJECT / "詳細設計書" / f"詳細設計書_{PROGRAM}.md",
    "cobol-src": COBOL_PROJECT / "src" / f"{PROGRAM}.cbl",
    "file-db-md": COBOL_PROJECT / "詳細設計書" / "COPY句定義書.md",
    "cpy": COBOL_PROJECT / "cpy",
    "db-md": COBOL_PROJECT / "詳細設計書" / "DB定義書.md",
}

DEFAULT_API_KEY = "sk-6156cccdc9c14d949cf5bfc5afc67a03"
DEFAULT_API_MODEL = "deepseek-v4-flash"

ORCHESTRATOR_JCL = PROJECT_ROOT / "orchestrator_jcl.py"


# ============================================================================
# Helpers
# ============================================================================

class Colors:
    RESET  = "\033[0m"
    RED    = "\033[91m"
    GREEN  = "\033[92m"
    YELLOW = "\033[93m"
    BLUE   = "\033[94m"
    CYAN   = "\033[96m"
    BOLD   = "\033[1m"


def header(text):
    print(f"\n{Colors.CYAN}{Colors.BOLD}{'=' * 70}")
    print(f"  {text}")
    print(f"{'=' * 70}{Colors.RESET}\n")


def step(text):
    print(f"  {Colors.GREEN}[STEP]{Colors.RESET} {text}")


def info(text):
    print(f"  {Colors.BLUE}[INFO]{Colors.RESET} {text}")


def warn(text):
    print(f"  {Colors.YELLOW}[WARN]{Colors.RESET} {text}")


def fail(text):
    print(f"  {Colors.RED}[FAIL]{Colors.RESET} {text}")


def ok(text):
    print(f"  {Colors.GREEN}[ OK ]{Colors.RESET} {text}")


def run_cmd(cmd, cwd=None, env=None, timeout=120, capture=True):
    return subprocess.run(
        cmd,
        cwd=cwd or str(PROJECT_ROOT),
        env=env,
        timeout=timeout,
        capture_output=capture,
        text=True,
        encoding="utf-8",
        errors="replace",
    )


# ============================================================================
# Phase 1: Environment Check
# ============================================================================

def check_environment():
    header("Phase 1: Environment Check")
    all_ok = True

    # --- Python version ---
    step("Python version")
    ver = sys.version_info
    if ver >= (3, 9):
        ok(f"Python {ver.major}.{ver.minor}.{ver.micro}")
    else:
        fail(f"Python {ver.major}.{ver.minor}.{ver.micro} (need >= 3.9)")
        all_ok = False

    # --- cobc (GnuCOBOL) ---
    step("GnuCOBOL (cobc)")
    cobc_path = shutil.which("cobc")
    if cobc_path:
        ok(f"cobc: {cobc_path}")
        try:
            r = run_cmd(["cobc", "--version"], timeout=10)
            first_line = (r.stdout or r.stderr or "").strip().splitlines()[:1]
            if first_line:
                info(first_line[0])
        except Exception:
            pass
    else:
        fail("cobc not found in PATH")
        all_ok = False

    # --- Java ---
    step("Java")
    java_path = shutil.which("java")
    if java_path:
        ok(f"java: {java_path}")
    else:
        warn("java not found in PATH (COBOL-only mode, Java comparison skipped)")

    # --- jcl-cobol-data-create ---
    step("jcl-cobol-data-create")
    jcl_root = _find_jcl_root()
    if jcl_root:
        ok(f"Found at: {jcl_root}")
    else:
        fail("jcl-cobol-data-create not found in any candidate location")
        all_ok = False

    # --- Test input files ---
    step("Test input files for ZAN04MAT")
    for label, path in TEST_INPUTS.items():
        if path.exists():
            ok(f"{label}: {path}")
        else:
            fail(f"{label}: {path} (NOT FOUND)")
            all_ok = False

    # --- Project structure ---
    step("Project structure")
    for d in ["cobol_testgen", "runners", "comparator", "agents", "tests", "config"]:
        p = PROJECT_ROOT / d
        if p.exists():
            ok(f"{d}/")
        else:
            fail(f"{d}/ (missing)")
            all_ok = False

    return all_ok


def _find_jcl_root():
    for candidate in JCL_ROOT_CANDIDATES:
        if candidate.exists() and (candidate / "agent" / "__init__.py").exists():
            return candidate
    return None


# ============================================================================
# Phase 2: Fix JCL_ROOT Path
# ============================================================================

def fix_jcl_root():
    header("Phase 2: Fix JCL_ROOT in orchestrator_jcl.py")

    if not ORCHESTRATOR_JCL.exists():
        fail(f"orchestrator_jcl.py not found at {ORCHESTRATOR_JCL}")
        return False

    content = ORCHESTRATOR_JCL.read_text(encoding="utf-8")

    # Find current JCL_ROOT value
    match = re.search(r'^JCL_ROOT\s*=\s*r?"([^"]*)"', content, re.MULTILINE)
    if not match:
        warn("Could not parse JCL_ROOT from orchestrator_jcl.py")
        return False

    current_path_str = match.group(1)
    current_path = Path(current_path_str)

    if current_path.exists() and (current_path / "agent" / "__init__.py").exists():
        ok(f"JCL_ROOT already correct: {current_path}")
        return True

    # Find the actual path
    actual_root = _find_jcl_root()
    if not actual_root:
        fail("Cannot find jcl-cobol-data-create to fix the path")
        return False

    actual_str = str(actual_root).replace("\\", "\\\\")
    new_line = f'JCL_ROOT = r"{actual_str}"'

    # Replace in content
    new_content = re.sub(
        r'^JCL_ROOT\s*=.*$',
        new_line,
        content,
        count=1,
        flags=re.MULTILINE,
    )

    ORCHESTRATOR_JCL.write_text(new_content, encoding="utf-8")
    ok(f"Fixed JCL_ROOT: {current_path_str} -> {actual_str}")
    return True


# ============================================================================
# Phase 3: Install Dependencies
# ============================================================================

def install_dependencies():
    header("Phase 3: Install Python Dependencies")

    req_main = PROJECT_ROOT / "requirements.txt"
    req_jcl = _find_jcl_root()
    if req_jcl:
        req_jcl = req_jcl / "requirements.txt"

    # Check what's installed
    step("Checking installed packages")
    required = {
        "httpx": "httpx",
        "pyyaml": "pyyaml",
        "pytest": "pytest",
        "fastapi": "fastapi",
        "uvicorn": "uvicorn",
        "lark": "lark",
        "requests": "requests",
    }

    missing = []
    for pkg, import_name in required.items():
        try:
            __import__(import_name)
            ok(f"{pkg}")
        except ImportError:
            warn(f"{pkg} -- MISSING")
            missing.append(pkg)

    if missing:
        step(f"Installing missing packages: {', '.join(missing)}")
        # Install from requirements.txt
        if req_main.exists():
            r = run_cmd(
                [sys.executable, "-m", "pip", "install", "-r", str(req_main)],
                timeout=120,
            )
            if r.returncode == 0:
                ok("Main requirements installed")
            else:
                warn(f"pip install had issues: {(r.stderr or '')[:200]}")

        # Install lark separately if missing (not in requirements.txt)
        if "lark" in missing:
            r = run_cmd(
                [sys.executable, "-m", "pip", "install", "lark>=1.1.0"],
                timeout=60,
            )
            if r.returncode == 0:
                ok("lark installed")

        # Install jcl-cobol-data-create requirements
        if req_jcl and req_jcl.exists():
            r = run_cmd(
                [sys.executable, "-m", "pip", "install", "-r", str(req_jcl)],
                timeout=60,
            )
            if r.returncode == 0:
                ok("jcl-cobol-data-create requirements installed")

        # Verify again
        step("Verifying installation")
        still_missing = []
        for pkg, import_name in required.items():
            try:
                __import__(import_name)
                ok(f"{pkg}")
            except ImportError:
                fail(f"{pkg} still missing")
                still_missing.append(pkg)

        if still_missing:
            fail(f"Could not install: {', '.join(still_missing)}")
            return False
    else:
        ok("All required packages are installed")

    return True


# ============================================================================
# Phase 4: Run Unit Tests (pytest)
# ============================================================================

def run_unit_tests():
    header("Phase 4: Unit Tests (pytest)")

    step("Running pytest...")
    r = run_cmd(
        [sys.executable, "-m", "pytest", "tests/", "-v", "--tb=short", "-x",
         "--ignore=tests/test_biz_e2e.py", "--ignore=tests/test_web_e2e.py"],
        timeout=300,
    )

    print(r.stdout[-3000:] if len(r.stdout) > 3000 else r.stdout)

    if r.returncode == 0:
        ok("All unit tests passed")
        return True
    else:
        # Count pass/fail from output
        output = r.stdout + (r.stderr or "")
        pass_count = output.count(" PASSED")
        fail_count = output.count(" FAILED")
        error_count = output.count(" ERROR")
        warn(f"Results: {pass_count} passed, {fail_count} failed, {error_count} errors")
        if fail_count <= 3:
            warn("Minor failures detected -- continuing with pipeline")
            return True
        else:
            fail("Too many test failures")
            return False


# ============================================================================
# Phase 5: Dry-Run Validation
# ============================================================================

def run_dry_run():
    header("Phase 5: Dry-Run Validation")

    step("Checking all inputs via main.py --dry-run...")
    r = run_cmd(
        [
            sys.executable, "main.py",
            "--mode", "agent",
            "--dry-run",
            "--design", str(TEST_INPUTS["design"]),
            "--cobol-src", str(TEST_INPUTS["cobol-src"]),
            "--file-db-md", str(TEST_INPUTS["file-db-md"]),
            "--cpy", str(TEST_INPUTS["cpy"]),
            "--db-md", str(TEST_INPUTS["db-md"]),
        ],
        timeout=30,
    )

    output = (r.stdout or "") + "\n" + (r.stderr or "")
    print(output)

    if "DRY-RUN: all inputs OK" in output or r.returncode == 0:
        ok("Dry-run validation passed")
        return True
    else:
        fail("Dry-run validation failed")
        return False


# ============================================================================
# Phase 6: Full Agent Mode Pipeline
# ============================================================================

def run_full_pipeline():
    header("Phase 6: Full Agent Mode Pipeline (ZAN04MAT)")

    output_dir = PROJECT_ROOT / "output" / "test_agent_mode"
    if output_dir.exists():
        shutil.rmtree(str(output_dir))
    output_dir.mkdir(parents=True, exist_ok=True)

    step(f"Output directory: {output_dir}")
    step(f"Program: {PROGRAM}")
    step(f"API model: {DEFAULT_API_MODEL}")
    info("DeepSeek API calls will be made -- ensure network connectivity")

    t0 = time.time()

    cmd = [
        sys.executable, "main.py",
        "--mode", "agent",
        "--design", str(TEST_INPUTS["design"]),
        "--cobol-src", str(TEST_INPUTS["cobol-src"]),
        "--file-db-md", str(TEST_INPUTS["file-db-md"]),
        "--cpy", str(TEST_INPUTS["cpy"]),
        "--db-md", str(TEST_INPUTS["db-md"]),
        "--output", str(output_dir),
        "--api-key", DEFAULT_API_KEY,
        "--api-model", DEFAULT_API_MODEL,
        "--verbose",
    ]

    step("Executing pipeline (this may take several minutes)...")
    info(f"Command: {' '.join(cmd)}")

    try:
        # Use a larger timeout since API calls + compilation can be slow
        r = run_cmd(cmd, timeout=600)
        elapsed = time.time() - t0

        print(r.stdout[-5000:] if len(r.stdout) > 5000 else r.stdout)
        if r.stderr:
            print(r.stderr[-2000:] if len(r.stderr) > 2000 else r.stderr)

        info(f"Pipeline completed in {elapsed:.0f}s")

        # Check results
        if r.returncode == 0:
            ok("Pipeline completed successfully")
        elif r.returncode == 2:
            warn("Pipeline completed with BLOCKED status (compile or input issue)")
        elif r.returncode == 3:
            warn("Pipeline completed with ERROR status")
        else:
            warn(f"Pipeline returned exit code {r.returncode}")

        # Check output files
        step("Checking output artifacts")
        _check_output_artifacts(output_dir)

        # Check reports
        reports_dir = PROJECT_ROOT / "reports" / PROGRAM
        if reports_dir.exists():
            step("Reports generated")
            for f in reports_dir.rglob("*"):
                if f.is_file():
                    info(f"  {f.relative_to(PROJECT_ROOT)}")
        else:
            info("No reports directory found (reports may be in output dir)")

        return True

    except subprocess.TimeoutExpired:
        elapsed = time.time() - t0
        fail(f"Pipeline timed out after {elapsed:.0f}s")
        return False
    except Exception as e:
        fail(f"Pipeline error: {e}")
        return False


def _check_output_artifacts(output_dir):
    prog_dir = output_dir / PROGRAM
    if not prog_dir.exists():
        info(f"No output for {PROGRAM} at {prog_dir}")
        return

    for item in sorted(prog_dir.iterdir()):
        if item.is_dir():
            info(f"  {item.name}/")
            for f in item.rglob("*"):
                if f.is_file():
                    size = f.stat().st_size
                    info(f"    {f.name} ({size} bytes)")
        else:
            info(f"  {item.name} ({item.stat().st_size} bytes)")


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Unified automated test for feat/agent-mode",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--skip-tests", action="store_true",
                        help="Skip unit tests (pytest)")
    parser.add_argument("--skip-pipeline", action="store_true",
                        help="Skip full agent mode pipeline")
    parser.add_argument("--dry-run-only", action="store_true",
                        help="Only check environment and dry-run")
    parser.add_argument("--api-key", default=DEFAULT_API_KEY,
                        help="DeepSeek API key")
    parser.add_argument("--api-model", default=DEFAULT_API_MODEL,
                        help="API model name")
    args = parser.parse_args()

    global DEFAULT_API_KEY, DEFAULT_API_MODEL
    DEFAULT_API_KEY = args.api_key
    DEFAULT_API_MODEL = args.api_model

    # Make sure we're in the project root
    os.chdir(str(PROJECT_ROOT))

    print(f"\n{Colors.BOLD}{'#' * 70}")
    print(f"#  COBOL-Java Agent Mode -- Unified Test Script")
    print(f"#  Branch: feat/agent-mode")
    print(f"#  Program: {PROGRAM}")
    print(f"#  Project: {PROJECT_ROOT}")
    print(f"{'#' * 70}{Colors.RESET}\n")

    results = {}
    t_start = time.time()

    # Phase 1: Environment
    results["env"] = check_environment()
    if not results["env"]:
        fail("\nEnvironment check failed. Please fix the issues above and retry.")
        sys.exit(1)

    # Phase 2: Fix JCL_ROOT
    results["jcl_root"] = fix_jcl_root()

    # Phase 3: Dependencies
    results["deps"] = install_dependencies()
    if not results["deps"]:
        fail("\nDependency installation failed.")
        sys.exit(1)

    # Phase 4: Unit tests
    if not args.skip_tests and not args.dry_run_only:
        results["tests"] = run_unit_tests()
    else:
        info("\nSkipping unit tests (--skip-tests or --dry-run-only)")

    # Phase 5: Dry-run
    results["dry_run"] = run_dry_run()

    # Phase 6: Full pipeline
    if not args.skip_pipeline and not args.dry_run_only:
        results["pipeline"] = run_full_pipeline()
    else:
        info("\nSkipping full pipeline (--skip-pipeline or --dry-run-only)")

    # Summary
    elapsed_total = time.time() - t_start
    header("Summary")
    for phase, passed in results.items():
        status = f"{Colors.GREEN}PASS{Colors.RESET}" if passed else f"{Colors.RED}FAIL{Colors.RESET}"
        print(f"  {phase:<15} {status}")

    print(f"\n  Total time: {elapsed_total:.0f}s")
    print(f"  Output dir: {PROJECT_ROOT / 'output' / 'test_agent_mode'}")
    print(f"  Reports:    {PROJECT_ROOT / 'reports' / PROGRAM}")
    print()

    all_pass = all(results.values())
    if all_pass:
        ok("All phases completed successfully!")
    else:
        failed_phases = [p for p, v in results.items() if not v]
        warn(f"Failed phases: {', '.join(failed_phases)}")

    sys.exit(0 if all_pass else 1)


if __name__ == "__main__":
    main()