import importlib.util
import re
import sys
from datetime import datetime
from pathlib import Path
from typing import List, Optional

import polars as pl
import typer
from rich.console import Console

from articuno.convert import infer_pydantic_model
from articuno.codegen import generate_pydantic_class_code
from articuno.bootstrap import get_inference_registry


app = typer.Typer()
console = Console()


def backup_file(file_path: Path) -> Path:
    """
    Create a timestamped backup of a file if it exists.

    Parameters
    ----------
    file_path : Path
        The path of the file to backup.

    Returns
    -------
    Path
        The path to the backup file if created, else the original file path.
    """
    if not file_path.exists():
        return file_path
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_path = file_path.with_name(f"{file_path.stem}_backup_{timestamp}{file_path.suffix}")
    backup_path.write_bytes(file_path.read_bytes())
    return backup_path

def insert_model_import(lines: List[str], model_name: str, models_path: Path) -> List[str]:
    """
    Insert an import statement for the generated model into source code lines.

    Parameters
    ----------
    lines : List[str]
        The source code lines to modify.
    model_name : str
        The Pydantic model name to import.
    models_path : Path
        The path of the module to import from.

    Returns
    -------
    List[str]
        Modified list of source code lines with the import inserted.
    """
    import_line = f"from {models_path.stem} import {model_name}  # autogenerated by Articuno\n"
    if any(import_line.strip() == line.strip() for line in lines):
        return lines
    last_import_idx = 0
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("import ") or stripped.startswith("from "):
            last_import_idx = i + 1
    lines.insert(last_import_idx, import_line)
    return lines


def find_and_modify_decorator_line(
    lines: List[str], func_name: str, model_name: str
) -> List[str]:
    """
    Add or update the response_model argument in the FastAPI route decorator.

    Parameters
    ----------
    lines : List[str]
        Source code lines.
    func_name : str
        The name of the function to update.
    model_name : str
        The Pydantic model name to set as response_model.

    Returns
    -------
    List[str]
        Modified lines with updated decorator.
    """
    http_decorators = ["get", "post", "put", "patch", "delete", "options", "head"]
    new_lines = lines.copy()
    func_line_idx = None
    for i, line in enumerate(new_lines):
        if line.lstrip().startswith(f"def {func_name}"):
            func_line_idx = i
            break
    if func_line_idx is None:
        return new_lines
    for i in range(func_line_idx - 1, -1, -1):
        stripped = new_lines[i].lstrip()
        if stripped.startswith("@infer_response_model"):
            new_lines[i] = ""
        elif any(stripped.startswith(f"@app.{method}(") for method in http_decorators):
            if "response_model=" not in new_lines[i]:
                new_lines[i] = re.sub(
                    r"\)$",
                    f", response_model={model_name})",
                    new_lines[i]
                )
            break
    return new_lines


def update_source_file(
    source_file: Path,
    model_name: str,
    models_path: Path,
    func_name: str,
) -> None:
    """
    Update the source file by backing it up, inserting model import,
    and modifying the route decorator to include the response model.

    Parameters
    ----------
    source_file : Path
        The source Python file to modify.
    model_name : str
        The generated Pydantic model name.
    models_path : Path
        Path to the module where the models are stored.
    func_name : str
        Name of the function (endpoint) to update.
    """
    backup_file(source_file)
    lines = source_file.read_text(encoding="utf-8").splitlines(keepends=True)
    lines = insert_model_import(lines, model_name, models_path)
    lines = find_and_modify_decorator_line(lines, func_name, model_name)
    source_file.write_text("".join(lines), encoding="utf-8")


@app.command()
def bootstrap(
    app_path: Path = typer.Argument(..., exists=True),
    models_path: Optional[Path] = typer.Option(None),
    dry_run: bool = typer.Option(False, help="Preview changes without writing files"),
):
    """
    Scan the app file for @infer_response_model endpoints, generate Pydantic models,
    append the models to a models file, and update endpoint decorators with response_model.

    Parameters
    ----------
    app_path : Path
        Path to the FastAPI app Python file.
    models_path : Path, optional
        Path to save generated Pydantic models. Defaults to 'models.py' alongside app.
    dry_run : bool, optional
        If True, outputs generated models to console without writing files.
    """
    app_path = app_path.resolve()
    if models_path is None:
        models_path = app_path.parent / "models.py"
    models_path = models_path.resolve()
    spec = importlib.util.spec_from_file_location(app_path.stem, str(app_path))
    module = importlib.util.module_from_spec(spec)
    sys.modules[app_path.stem] = module
    spec.loader.exec_module(module)
    registry = get_inference_registry()
    if not registry:
        console.print("[yellow]No @infer_response_model endpoints found.[/]")
        raise typer.Exit(1)
    for item in registry:
        func = item["func"]
        model_name = item["name"]
        example_input = item["example_input"]
        model_file_path = Path(item["models_path"] or str(models_path)).resolve()
        source_file = Path(item["source_file"]).resolve()
        console.print(f"\n[bold]Processing:[/] {func.__name__}")
        try:
            output = func(**example_input)
        except Exception as e:
            console.print(f"[red]Error:[/] Failed to call {func.__name__}: {e}")
            continue
        if not isinstance(output, pl.DataFrame):
            console.print(f"[red]Error:[/] Output is not a Polars DataFrame")
            continue
        model = infer_pydantic_model(output, model_name=model_name)
        code = generate_pydantic_class_code(model, model_name=model_name)
        if dry_run:
            console.print(code)
        else:
            with model_file_path.open("a", encoding="utf-8") as f:
                f.write(f"\n\n# --- Articuno autogenerated model: {model_name} ---\n")
                f.write(code)
            update_source_file(source_file, model_name, model_file_path, func.__name__)
    console.print("\n[green]Bootstrapping complete![/]")


if __name__ == "__main__":
    app()
