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, infer_patito_model
from articuno.codegen import generate_pydantic_class_code, generate_patito_class_code
from articuno.bootstrap import get_inference_registry

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

def backup_file(file_path: Path) -> 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]:
    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]:
    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:
    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"),
    use_patito: bool = typer.Option(False, help="Use Patito instead of Pydantic for model generation"),
):
    """
    Bootstrap all registered endpoints to generate Pydantic or Patito models based on example outputs.

    Calls each registered function, infers the model from its Polars DataFrame output,
    generates source code for the model, updates the models file and modifies source files
    to add the model import and response_model decorator.

    Parameters
    ----------
    app_path : Path
        Path to the Python file containing the FastAPI app and endpoint functions.
    models_path : Optional[Path], optional
        Path to the models.py file to write generated models to. Defaults to sibling models.py.
    dry_run : bool, optional
        If True, prints the generated code instead of writing files.
    use_patito : bool, optional
        If True, generate Patito models instead of Pydantic.
    """
    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

        if use_patito:
            model = infer_patito_model(output, model_name=model_name)
            code = generate_patito_class_code(model, model_name=model_name)
        else:
            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()
