import argparse
import json
import os
import subprocess
import tempfile
from getpass import getpass
from pathlib import Path
from shutil import which
from typing import Optional

from groq import Groq

MODEL_ID = "llama-3.1-8b-instant"
CONFIG_PATH = Path.home() / ".komet" / "config.json"
# Optional publisher starter key for first-run experience.
# SECURITY NOTE: Do not commit a real production key in public repos.
STARTER_API_KEY = os.getenv("KOMET_STARTER_API_KEY", "")


def load_config() -> dict:
    if not CONFIG_PATH.exists():
        return {}
    try:
        return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}


def save_config(config: dict) -> None:
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    CONFIG_PATH.write_text(json.dumps(config, indent=2), encoding="utf-8")
    try:
        CONFIG_PATH.chmod(0o600)
    except OSError:
        pass


def print_key_setup_guide() -> None:
    print("\nYour current Groq API key is missing, invalid, or rate-limited.")
    print("Set your own key in 3 steps:")
    print("  1) Open: https://console.groq.com/keys")
    print("  2) Create/copy your API key")
    print("  3) Run: komet -k <YOUR_KEY>")
    print("Then run `komet` again.\n")


def is_key_error(exc: Exception) -> bool:
    text = str(exc).lower()
    return any(token in text for token in ["401", "403", "authentication", "unauthorized", "invalid api key", "rate limit", "429"])


def get_or_set_api_key(explicit_key: Optional[str] = None) -> str:
    config = load_config()
    if explicit_key:
        config["api_key"] = explicit_key.strip()
        save_config(config)
        return config["api_key"]

    stored = config.get("api_key", "").strip()
    if stored:
        return stored

    if STARTER_API_KEY.strip():
        return STARTER_API_KEY.strip()

    entered = getpass("Enter Groq API key: ").strip()
    if not entered:
        raise SystemExit("No API key provided.")
    config["api_key"] = entered
    save_config(config)
    return entered


def run_git(args: list[str], check: bool = True) -> subprocess.CompletedProcess[str]:
    try:
        return subprocess.run(
            ["git", *args],
            capture_output=True,
            text=True,
            check=check,
        )
    except FileNotFoundError as exc:
        raise SystemExit("git is not installed or not available in PATH.") from exc


def in_git_repo() -> bool:
    if which("git") is None:
        return False
    result = run_git(["rev-parse", "--is-inside-work-tree"], check=False)
    return result.returncode == 0 and result.stdout.strip() == "true"


def has_changes() -> bool:
    result = run_git(["status", "--porcelain"], check=False)
    return bool(result.stdout.strip())


def staged_diff(max_chars: int = 12000) -> str:
    result = run_git(["diff", "--cached"], check=False)
    diff = result.stdout.strip()
    if len(diff) > max_chars:
        return diff[:max_chars] + "\n\n# ... (truncated)"
    return diff


def ask(client: Groq, user_text: str) -> str:
    try:
        response = client.chat.completions.create(
            model=MODEL_ID,
            messages=[
                {
                    "role": "system",
                    "content": "You are Komet, a concise and practical coding assistant.",
                },
                {"role": "user", "content": user_text},
            ],
            temperature=0.3,
            max_tokens=900,
        )
        return (response.choices[0].message.content or "").strip()
    except Exception as exc:
        if is_key_error(exc):
            print_key_setup_guide()
        raise


def generate_commit_message(client: Groq, diff_text: str, attempt: int = 1) -> str:
    system = (
        "Write a single conventional commit message for the provided staged diff. "
        "Use type(scope): subject format when scope is useful. "
        "Imperative mood, concise, no backticks, no markdown fences, no explanation."
    )
    user = f"Attempt: {attempt}\n\nStaged diff:\n{diff_text}"
    try:
        response = client.chat.completions.create(
            model=MODEL_ID,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            temperature=0.2,
            max_tokens=180,
        )
        return (response.choices[0].message.content or "").strip()
    except Exception as exc:
        if is_key_error(exc):
            print_key_setup_guide()
        raise


def commit_with_message(message: str) -> None:
    with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as temp:
        temp.write(message.strip() + "\n")
        temp_path = temp.name
    try:
        run_git(["commit", "-F", temp_path], check=True)
    finally:
        Path(temp_path).unlink(missing_ok=True)


def commit_workflow(client: Groq, push: bool = False) -> int:
    if not in_git_repo():
        print("Not inside a git repository, or git is unavailable.")
        return 1

    if not has_changes():
        print("No changes detected.")
        return 0

    run_git(["add", "-A"], check=True)
    diff_text = staged_diff()
    if not diff_text:
        print("No staged changes to commit.")
        return 0

    attempt = 1
    while True:
        try:
            proposed = generate_commit_message(client, diff_text, attempt=attempt)
        except Exception as exc:
            if not is_key_error(exc):
                print(f"Groq API error: {exc}")
            return 1
        print("\nProposed commit message:\n")
        print(proposed)
        choice = input("\n[a]pprove / [r]egenerate / [e]dit / [q]uit: ").strip().lower()

        if choice in {"a", "approve"}:
            final_message = proposed
        elif choice in {"r", "regen", "regenerate"}:
            attempt += 1
            continue
        elif choice in {"e", "edit"}:
            print("Enter commit message. Finish input with an empty line.")
            lines: list[str] = []
            while True:
                line = input()
                if line == "":
                    break
                lines.append(line)
            if not lines:
                print("Empty message, canceled.")
                return 0
            final_message = "\n".join(lines)
        elif choice in {"q", "quit"}:
            print("Canceled.")
            return 0
        else:
            print("Unknown choice. Please select approve/regenerate/edit/quit.")
            continue

        try:
            commit_with_message(final_message)
            print("Commit created.")
        except subprocess.CalledProcessError as exc:
            print(exc.stderr.strip() or "Failed to create commit.")
            return exc.returncode or 1

        if push:
            push_result = run_git(["push"], check=False)
            if push_result.returncode != 0:
                print(push_result.stderr.strip() or "Push failed.")
                return push_result.returncode or 1
            print("Pushed to remote.")
        return 0


def interactive_chat(client: Groq) -> int:
    print("Komet chat mode. Type 'exit' to quit.")
    while True:
        try:
            prompt = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return 0

        if not prompt:
            continue
        if prompt.lower() in {"exit", "quit"}:
            return 0

        try:
            answer = ask(client, prompt)
        except Exception as exc:
            print(f"Groq API error: {exc}")
            continue
        print(f"komet> {answer}\n")


def main() -> int:
    parser = argparse.ArgumentParser(
        prog="komet",
        description=(
            "Komet CLI\n"
            "Default mode is interactive Groq chat.\n"
            "Use commit flags to run AI-assisted git commit workflow."
        ),
        epilog=(
            "Examples:\n"
            "  komet\n"
            "  komet -k gsk_...\n"
            "  komet -c\n"
            "  komet -cp\n"
            "  komet -c -p"
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "-k",
        "--set-key",
        dest="set_key",
        metavar="API_KEY",
        help="Set and store Groq API key in ~/.komet/config.json",
    )
    parser.add_argument(
        "-c",
        "--commit",
        action="store_true",
        help="Stage all changes, propose commit message, then commit",
    )
    parser.add_argument(
        "-p",
        "--push",
        action="store_true",
        help="Push after successful commit (used with -c)",
    )
    parser.add_argument(
        "-cp",
        "--commit-push",
        action="store_true",
        help="Shortcut for: -c -p",
    )
    args = parser.parse_args()

    push_mode = args.push or args.commit_push
    commit_mode = args.commit or args.commit_push

    api_key = get_or_set_api_key(args.set_key)
    client = Groq(api_key=api_key)

    if commit_mode:
        return commit_workflow(client, push=push_mode)

    return interactive_chat(client)


if __name__ == "__main__":
    raise SystemExit(main())