app/main.py

import time
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from app.models import TaskRequest, TaskResponse, ErrorResponse
from app.config import get_settings
from app.services.aipipe_client import AIpipeClient
from app.services.github_service import GitHubService
from app.services.pages import wait_for_200
from app.services.notify import NotificationService
from app.services.storage import init_db, get_default_db_path, upsert_task, get_task
from typing import Dict

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Initialize SQLite DB
    db_path = get_default_db_path()
    init_db(db_path)
    yield
    # Shutdown: nothing to cleanup currently

app = FastAPI(title="LLM Code Deployment Builder", lifespan=lifespan)

@app.get("/")
async def homepage():
    return {"status": "Service is running"}

@app.get("/healthz", response_model=Dict[str, str])
async def healthz():
    return {"status": "ok"}

async def process_task_async(app: FastAPI, req: TaskRequest, settings) -> None:
    """Process the task asynchronously and notify evaluation URL."""
    notification_service = NotificationService()
    try:
        start_time = time.time()
        # Initialize task state
        if not hasattr(app.state, "tasks"):
            app.state.tasks = {}
        if not hasattr(app.state, "tasks_by_task"):
            app.state.tasks_by_task = {}
        app.state.tasks[req.nonce] = {"status": "processing"}

        # Build attachments meta
        attachments_meta = [{"name": a.name, "url": a.url} for a in req.attachments]

        ai = AIpipeClient(settings.aipipe_model)

        # ROUND 2: revision flow keyed by task name
        if req.round == 2:
            db_path = get_default_db_path()
            print(f"[Round2][{req.task}] Starting revision for task. Nonce={req.nonce}")
            # Prefer DB, fallback to memory if present
            prior = get_task(db_path, req.task) or app.state.tasks_by_task.get(req.task)
            # If not found, fallback to GitHub lookup by task-derived name
            if (not prior) or (not prior.get("repo_url")):
                gh = GitHubService(settings.github_token, settings.github_username, settings.github_email)
                repo_url_fallback = gh.find_repo_by_task(req.task)
                if repo_url_fallback:
                    prior = prior or {}
                    prior.update({
                        "repo_url": repo_url_fallback,
                        "pages_url": gh.get_pages_url(repo_url_fallback),
                    })
            if not prior or not prior.get("repo_url"):
                raise RuntimeError("No prior repository found for this task. Run round 1 first.")

            gh = GitHubService(settings.github_token, settings.github_username, settings.github_email)
            print(f"[Round2][{req.task}] Found repo URL: {prior['repo_url']}")
            repo_full_name = gh._repo_full_name_from_url(prior["repo_url"])  # owner/repo
            print(f"[Round2][{req.task}] Repo full name: {repo_full_name}")

            print(f"[Round2][{req.task}] Fetching current files from repo...")
            current_files = gh.get_text_files(repo_full_name, [
                "README.md", "index.html", "style.css", "script.js"
            ])
            print(f"[Round2][{req.task}] Retrieved files: {list(current_files.keys())}")

            print(f"[Round2][{req.task}] Requesting minimal changes from LLM...")
            changed_files: Dict[str, str] = await ai.generate_revision_files(
                req.brief, req.checks, attachments_meta, current_files
            )

            if not changed_files:
                base = current_files.get("README.md", "# Project\n\n")
                changed_files = {"README.md": base + "\n\nRound 2: No code changes required.\n"}
                print(f"[Round2][{req.task}] LLM returned no changes; applying README bump to trigger deploy.")
            else:
                print(f"[Round2][{req.task}] Applying {len(changed_files)} changed file(s): {list(changed_files.keys())}")

            print(f"[Round2][{req.task}] Updating repository with changes...")
            commit_sha = gh.update_files_in_repo(
                repo_full_name, changed_files, commit_message_prefix="feat(revision)"
            )
            print(f"[Round2][{req.task}] Update complete. Commit SHA: {commit_sha}")
            repo_url = prior["repo_url"]
            repo_name = repo_url.rstrip("/").split("/")[-1]
            pages_url = gh.get_pages_url(repo_url, repo_name)

            # Use internal defaults (no env needed) for Pages readiness
            print(f"[Round2][{req.task}] Waiting for GitHub Pages (using internal defaults)...")
            ok = await wait_for_200(pages_url)
            if not ok:
                raise RuntimeError("GitHub Pages did not become reachable within the timeout after round 2 updates")
            
            # Wait for pages to update; use non-blocking sleep to avoid blocking the event loop
            import asyncio
            await asyncio.sleep(60)
            
            print(f"[Round2][{req.task}] Pages check successful: {pages_url}")

            result = {
                "email": req.email,
                "task": req.task,
                "round": req.round,
                "nonce": req.nonce,
                "repo_url": repo_url,
                "commit_sha": commit_sha,
                "pages_url": pages_url,
            }

            print(f"[Round2][{req.task}] Notifying evaluation URL...")
            await notification_service.notify_evaluation(req.evaluation_url, result)
            print(f"[Round2][{req.task}] Round 2 completed successfully.")

            app.state.tasks[req.nonce] = {"status": "completed", "result": result}
            app.state.tasks_by_task[req.task] = {
                "repo_url": repo_url,
                "pages_url": pages_url,
                "commit_sha": commit_sha,
            }
            # Persist to DB
            upsert_task(db_path, req.task, repo_url, pages_url, commit_sha)
            return

        # ROUND 1: initial build
        print(f"[Round1][{req.task}] Starting initial build (Nonce={req.nonce})")
        print(f"[Round1][{req.task}] Using model: {ai.model}")
        files: Dict[str, str] = await ai.generate_app_files(req.brief, req.checks, attachments_meta)

        # Check time limit (10 minutes = 600 seconds)
        if time.time() - start_time > 540:
            raise TimeoutError("Task processing exceeded time limit")

        if not (settings.github_token and settings.github_username and settings.github_email):
            raise HTTPException(status_code=400, detail="GitHub config missing")

        gh = GitHubService(settings.github_token, settings.github_username, settings.github_email)
        repo_info = gh.create_repo_with_files(req.task, files, req.nonce)
        repo_url = repo_info["repo_url"]
        repo_name = repo_url.rstrip("/").split("/")[-1]
        pages_url = gh.get_pages_url(repo_url, repo_name)
        print(f"[Round1][{req.task}] Repository created: {repo_url}")

        ok = await wait_for_200(pages_url)
        if not ok:
            raise RuntimeError("GitHub Pages did not become reachable within the timeout after initial deployment")

        commit_sha = gh.get_latest_commit_sha(f"{settings.github_username}/{repo_name}")

        result = {
            "email": req.email,
            "task": req.task,
            "round": req.round,
            "nonce": req.nonce,
            "repo_url": repo_url,
            "commit_sha": commit_sha,
            "pages_url": pages_url,
        }

        await notification_service.notify_evaluation(req.evaluation_url, result)
        print(f"[Round1][{req.task}] Completed successfully.")

        app.state.tasks[req.nonce] = {"status": "completed", "result": result}
        app.state.tasks_by_task[req.task] = {
            "repo_url": repo_url,
            "pages_url": pages_url,
            "commit_sha": commit_sha,
        }
        # Persist to DB for future round 2
        db_path = get_default_db_path()
        upsert_task(db_path, req.task, repo_url, pages_url, commit_sha)

    except Exception as e:
        print(f"[Error][{getattr(req, 'task', 'unknown')}] {e}")
        error_result = {
            "email": req.email,
            "task": req.task,
            "round": req.round,
            "status": "error",
            "error": str(e),
        }
        try:
            evaluation_url = str(req.evaluation_url)
            await notification_service.notify_evaluation(evaluation_url, error_result)
        except Exception as notify_err:
            print(f"Failed to notify evaluation URL: {notify_err}")

        if hasattr(app.state, "tasks"):
            app.state.tasks[req.nonce] = {"status": "failed", "error": str(e)}
    finally:
        await notification_service.close()

@app.post(
    "/tasks",
    response_model=TaskResponse,
    response_model_exclude_none=True,
    responses={400: {"model": ErrorResponse}},
)
async def handle_task(req: TaskRequest, background_tasks: BackgroundTasks):
    settings = get_settings()
    
    # Secret check
    if not settings.app_secret or req.secret != settings.app_secret:
        raise HTTPException(status_code=400, detail="Invalid secret")
    
    # For round 2, ensure the task exists from round 1 BEFORE returning 200
    if req.round == 2:
        db_path = get_default_db_path()
        # Initialize in-memory mapping if absent
        if not hasattr(app.state, "tasks_by_task"):
            app.state.tasks_by_task = {}
        prior = get_task(db_path, req.task) or app.state.tasks_by_task.get(req.task)
        if not prior or not prior.get("repo_url"):
            # Fallback to GitHub lookup by task-derived name before rejecting
            gh = GitHubService(settings.github_token, settings.github_username, settings.github_email)
            repo_url_fallback = gh.find_repo_by_task(req.task)
            if repo_url_fallback:
                # Stash minimally so background can proceed
                app.state.tasks_by_task[req.task] = {"repo_url": repo_url_fallback}
            else:
                raise HTTPException(status_code=400, detail="No prior repository found for this task. Run round 1 first.")

    # Initialize task state storage if needed
    if not hasattr(app.state, "tasks"):
        app.state.tasks = {}

    # Add task to background processing
    background_tasks.add_task(process_task_async, app, req, settings)
    
    # Immediate response with task accepted and required fields
    return TaskResponse(
        status="accepted",
        task=req.task,
        round=req.round,
        message="Task is being processed"
    )