import os
import sys
import argparse
import tomllib
import re
import subprocess

# ANSI color codes for specific coloring
RED = "\033[91m"
ORANGE = "\033[38;5;208m"
GREEN = "\033[92m"
RESET = "\033[0m"

# Resolve project root to the directory containing this script,
# so validation works regardless of the caller's working directory.
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))


def _proj_path(*parts: str) -> str:
    """Return an absolute path relative to the project root."""
    return os.path.join(PROJECT_ROOT, *parts)


# ── checks ────────────────────────────────────────────────────────────────

def check_uv_lock():
    if not os.path.exists(_proj_path("uv.lock")):
        return "- Missing uv.lock - run 'uv lock' to generate it"
    return None


def check_project_scripts():
    toml_path = _proj_path("pyproject.toml")
    if not os.path.exists(toml_path):
        return "- Missing pyproject.toml"

    with open(toml_path, "rb") as f:
        try:
            data = tomllib.load(f)
            scripts = data.get("project", {}).get("scripts", {})
            if not scripts:
                return "- Missing [project.scripts] section"

            # Check for a server entry point that points to a :main function
            found_server_entry = False
            for name, entry in scripts.items():
                if ":main" in entry:
                    found_server_entry = True
                    break

            if not found_server_entry:
                return "- Missing [project.scripts] server entry point"

        except Exception as e:
            return f"- Error parsing pyproject.toml: {e}"

    return None


def check_server_app():
    app_path = _proj_path("server", "app.py")
    if not os.path.exists(app_path):
        return ["- server/app.py missing"]

    issues = []
    with open(app_path, "r", encoding="utf-8") as f:
        content = f.read()

        # Check for main() function definition
        if not re.search(r"def\s+main\(\s*\):", content):
            issues.append("- server/app.py missing main() function")

        # Check for if __name__ == "__main__" guard that calls main()
        has_if_main = re.search(r'if\s+__name__\s*==\s*["\']__main__["\']\s*:', content)
        if has_if_main:
            if_block_start = has_if_main.end()
            remaining = content[if_block_start:]
            if "main()" not in remaining:
                issues.append(
                    "- server/app.py main() function not callable "
                    "(missing if __name__ == '__main__')"
                )
        else:
            issues.append("- server/app.py missing if __name__ == '__main__': block")

    return issues


# ── auto-fix ──────────────────────────────────────────────────────────────

def fix_uv_lock():
    """Generate uv.lock by running 'uv lock'."""
    print(f"  {GREEN}→ Generating uv.lock …{RESET}")
    result = subprocess.run(
        ["uv", "lock"],
        cwd=PROJECT_ROOT,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        print(f"  {RED}  uv lock failed: {result.stderr.strip()}{RESET}")
        return False
    print(f"  {GREEN}  uv.lock created.{RESET}")
    return True


def fix_project_scripts():
    """Ensure [project.scripts] contains a server entry pointing to server.app:main."""
    toml_path = _proj_path("pyproject.toml")
    if not os.path.exists(toml_path):
        print(f"  {RED}  Cannot fix — pyproject.toml not found.{RESET}")
        return False

    with open(toml_path, "r", encoding="utf-8") as f:
        content = f.read()

    # If the section already exists, add the entry; otherwise append the section.
    if "[project.scripts]" in content:
        # Add a server entry if missing
        if 'server' not in content.split("[project.scripts]")[1].split("[")[0]:
            content = content.replace(
                "[project.scripts]",
                '[project.scripts]\nserver = "server.app:main"',
            )
    else:
        content += '\n[project.scripts]\nserver = "server.app:main"\n'

    with open(toml_path, "w", encoding="utf-8") as f:
        f.write(content)

    print(f"  {GREEN}  Added [project.scripts] server entry.{RESET}")
    return True


def fix_server_app():
    """Ensure server/app.py has a main() function and if __name__ guard."""
    app_path = _proj_path("server", "app.py")
    if not os.path.exists(app_path):
        print(f"  {RED}  Cannot fix — server/app.py not found.{RESET}")
        return False

    with open(app_path, "r", encoding="utf-8") as f:
        content = f.read()

    changed = False

    # Add main() if missing
    if not re.search(r"def\s+main\(\s*\):", content):
        main_block = (
            '\n\ndef main():\n'
            '    import uvicorn\n'
            '    uvicorn.run("server.app:app", host="0.0.0.0", port=7860, reload=False)\n'
        )
        content += main_block
        changed = True
        print(f"  {GREEN}  Added main() function.{RESET}")

    # Add if __name__ guard if missing
    if not re.search(r'if\s+__name__\s*==\s*["\']__main__["\']\s*:', content):
        content += '\n\nif __name__ == "__main__":\n    main()\n'
        changed = True
        print(f"  {GREEN}  Added if __name__ == '__main__' guard.{RESET}")
    else:
        # Guard exists but may not call main()
        match = re.search(r'if\s+__name__\s*==\s*["\']__main__["\']\s*:', content)
        remaining = content[match.end():]
        if "main()" not in remaining:
            # Replace the if-block's body with main()
            content += "    main()\n"
            changed = True
            print(f"  {GREEN}  Added main() call inside if __name__ guard.{RESET}")

    if changed:
        with open(app_path, "w", encoding="utf-8") as f:
            f.write(content)

    return True


# ── commands ──────────────────────────────────────────────────────────────

def validate(auto_fix: bool = False):
    issues = []

    uv_issue = check_uv_lock()
    if uv_issue:
        issues.append(uv_issue)

    scripts_issue = check_project_scripts()
    if scripts_issue:
        issues.append(scripts_issue)

    server_issues = check_server_app()
    issues.extend(server_issues)

    if issues:
        print(f"{ORANGE}Issues found:{RESET}")
        for issue in issues:
            print(f"  {ORANGE}{issue}{RESET}")

        if auto_fix:
            print(f"\n{GREEN}Attempting auto-fix …{RESET}")
            if uv_issue:
                fix_uv_lock()
            if scripts_issue:
                fix_project_scripts()
            if server_issues:
                fix_server_app()

            # Re-validate after fixes
            print(f"\n{GREEN}Re-validating …{RESET}")
            recheck = []
            r = check_uv_lock()
            if r:
                recheck.append(r)
            r = check_project_scripts()
            if r:
                recheck.append(r)
            recheck.extend(check_server_app())

            if recheck:
                print(f"{RED}Some issues remain:{RESET}")
                for issue in recheck:
                    print(f"  {RED}{issue}{RESET}")
                sys.exit(1)
            else:
                print(f"{GREEN}All issues fixed! Repository is now compliant.{RESET}")
        else:
            print(f"\n  Tip: run '{sys.argv[0]} validate --fix' to auto-fix these issues.")
            sys.exit(1)
    else:
        print(f"{GREEN}No issues found! Repository is compliant.{RESET}")


def main():
    parser = argparse.ArgumentParser(description="OpenEnv CLI Tool")
    subparsers = parser.add_subparsers(dest="command")

    validate_parser = subparsers.add_parser(
        "validate", help="Validate repository compliance"
    )
    validate_parser.add_argument(
        "--fix",
        action="store_true",
        help="Automatically fix detected issues",
    )

    args = parser.parse_args()

    if args.command == "validate":
        validate(auto_fix=args.fix)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()