"""
InStatic CMS — AI Build Pipeline
Daviddolor/instatic-cms on HuggingFace Spaces

Pipeline: Prompt → Analyze → Plan → Split → Code → Validate → AutoFix → Deploy
Docs brain: WordPress + GitHub + Android + FastAPI markdown
"""

import os, json, asyncio, time, re
from pathlib import Path
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import Optional, Literal
import httpx

# ── Docs Brain Loader ──────────────────────────────────────────────────────
DOCS_DIR = Path(__file__).parent / "docs_brain"

def load_docs() -> dict[str, str]:
    """Load all markdown docs into memory at startup."""
    docs = {}
    for md_file in DOCS_DIR.glob("*.md"):
        docs[md_file.stem] = md_file.read_text(encoding="utf-8")
    return docs

DOCS: dict[str, str] = {}

# ── LLM Provider Fallback Chain ───────────────────────────────────────────
PROVIDERS = [
    {
        "name": "groq",
        "url": "https://api.groq.com/openai/v1/chat/completions",
        "key_env": "GROQ_API_KEY",
        "model": "llama-3.3-70b-versatile",
        "max_tokens": 8192,
    },
    {
        "name": "openrouter",
        "url": "https://openrouter.ai/api/v1/chat/completions",
        "key_env": "OPENROUTER_API_KEY",
        "model": "meta-llama/llama-3.3-70b-instruct",
        "max_tokens": 8192,
    },
    {
        "name": "cerebras",
        "url": "https://api.cerebras.ai/v1/chat/completions",
        "key_env": "CEREBRAS_API_KEY",
        "model": "llama3.1-70b",
        "max_tokens": 8192,
    },
]

async def call_llm(messages: list, max_tokens: int = 4096, json_mode: bool = False) -> str:
    """Call LLM with Groq→OpenRouter→Cerebras fallback."""
    for provider in PROVIDERS:
        key = os.getenv(provider["key_env"])
        if not key:
            continue
        try:
            body = {
                "model": provider["model"],
                "max_tokens": min(max_tokens, provider["max_tokens"]),
                "messages": messages,
                "temperature": 0.2,
            }
            if json_mode:
                body["response_format"] = {"type": "json_object"}

            async with httpx.AsyncClient(timeout=90) as client:
                resp = await client.post(
                    provider["url"],
                    headers={"Authorization": f"Bearer {key}"},
                    json=body,
                )
                resp.raise_for_status()
                return resp.json()["choices"][0]["message"]["content"]
        except Exception as e:
            print(f"[LLM:{provider['name']}] failed: {e}")
            continue
    raise RuntimeError("All LLM providers exhausted — check API keys")

async def stream_llm(messages: list, max_tokens: int = 4096) -> AsyncGenerator[str, None]:
    """Streaming LLM call."""
    for provider in PROVIDERS:
        key = os.getenv(provider["key_env"])
        if not key:
            continue
        try:
            async with httpx.AsyncClient(timeout=120) as client:
                async with client.stream(
                    "POST",
                    provider["url"],
                    headers={"Authorization": f"Bearer {key}"},
                    json={
                        "model": provider["model"],
                        "max_tokens": max_tokens,
                        "messages": messages,
                        "stream": True,
                        "temperature": 0.2,
                    },
                ) as response:
                    response.raise_for_status()
                    async for line in response.aiter_lines():
                        if line.startswith("data: "):
                            chunk = line[6:]
                            if chunk.strip() == "[DONE]":
                                return
                            try:
                                data = json.loads(chunk)
                                content = data["choices"][0]["delta"].get("content", "")
                                if content:
                                    yield content
                            except Exception:
                                pass
            return  # success — don't try next provider
        except Exception as e:
            print(f"[STREAM:{provider['name']}] failed: {e}")
            continue


# ── Pipeline Stages ────────────────────────────────────────────────────────

def get_doc_context(target: str) -> str:
    """Build relevant doc context based on build target."""
    docs_map = {
        "website":   ["fastapi", "github"],
        "wordpress": ["wordpress", "github"],
        "android":   ["android", "github"],
        "api":       ["fastapi", "github"],
        "fullstack": ["fastapi", "github", "wordpress", "android"],
    }
    keys = docs_map.get(target, ["fastapi", "github"])
    parts = []
    for key in keys:
        if key in DOCS:
            # Include first 3000 chars of each doc to stay within context
            parts.append(f"## {key.upper()} DOCS REFERENCE\n{DOCS[key][:3000]}")
    return "\n\n---\n\n".join(parts)


async def stage_analyze(prompt: str, target: str) -> dict:
    """Stage 1: Analyze the prompt and understand requirements."""
    doc_ctx = get_doc_context(target)
    messages = [
        {
            "role": "system",
            "content": f"""You are an expert software architect. Analyze user requirements and extract structured information.
Always respond with valid JSON only.

DOCUMENTATION CONTEXT:
{doc_ctx}
""",
        },
        {
            "role": "user",
            "content": f"""Analyze this build request and return JSON:

PROMPT: {prompt}
TARGET: {target}

Return JSON with these fields:
{{
  "project_type": "website|api|android|wordpress",
  "project_name": "kebab-case-name",
  "description": "one sentence description",
  "features": ["list", "of", "features"],
  "tech_stack": ["list", "of", "technologies"],
  "complexity": "simple|medium|complex",
  "ui_style": "description of visual style or null",
  "api_endpoints": ["list of needed endpoints or empty"],
  "data_models": ["list of data entities"],
  "has_auth": true|false,
  "has_database": true|false,
  "deployment_target": "cloudflare|render|hf-spaces|playstore|vercel"
}}
""",
        },
    ]
    raw = await call_llm(messages, max_tokens=1024, json_mode=True)
    try:
        return json.loads(raw)
    except Exception:
        # Best-effort extraction
        return {"project_type": target, "description": prompt[:100], "features": [], "tech_stack": []}


async def stage_plan(analysis: dict, prompt: str) -> dict:
    """Stage 2: Create a detailed build plan with file list."""
    messages = [
        {
            "role": "system",
            "content": "You are a software architect. Create a detailed file-by-file build plan. Return valid JSON only.",
        },
        {
            "role": "user",
            "content": f"""Based on this analysis, create a build plan:

ANALYSIS: {json.dumps(analysis, indent=2)}
ORIGINAL PROMPT: {prompt}

Return JSON:
{{
  "plan_summary": "2-3 sentence build plan",
  "files": [
    {{
      "path": "relative/file/path",
      "type": "html|css|js|python|kotlin|json|yaml|md",
      "description": "what this file does",
      "priority": 1
    }}
  ],
  "build_order": ["ordered", "list", "of", "file", "paths"],
  "dependencies": ["npm package or pip package"],
  "env_vars": ["REQUIRED_ENV_VAR"],
  "estimated_files": 5
}}
""",
        },
    ]
    raw = await call_llm(messages, max_tokens=2048, json_mode=True)
    try:
        return json.loads(raw)
    except Exception:
        return {"plan_summary": "Generating files...", "files": [], "build_order": []}


async def stage_generate_file(
    file_info: dict,
    analysis: dict,
    plan: dict,
    prompt: str,
    target: str,
    previously_generated: list[dict],
) -> str:
    """Stage 3: Generate actual file content."""
    doc_ctx = get_doc_context(target)
    context_summary = "\n".join(
        f"- {f['path']}: {f['description']}" for f in previously_generated[-3:]
    )

    messages = [
        {
            "role": "system",
            "content": f"""You are an expert {file_info['type']} developer.
Generate complete, production-ready file content.
Output ONLY the raw file content — no markdown fences, no explanations.

DOCS REFERENCE:
{doc_ctx[:2000]}
""",
        },
        {
            "role": "user",
            "content": f"""Generate the complete content for this file:

FILE: {file_info['path']}
TYPE: {file_info['type']}
PURPOSE: {file_info['description']}
PROJECT: {analysis.get('project_name', 'project')}
STACK: {', '.join(analysis.get('tech_stack', []))}
FEATURES: {', '.join(analysis.get('features', []))}
UI STYLE: {analysis.get('ui_style', 'clean, modern')}

ALREADY GENERATED:
{context_summary}

ORIGINAL REQUEST: {prompt}

Generate the COMPLETE file content now:
""",
        },
    ]
    return await call_llm(messages, max_tokens=4096)


async def stage_validate(file_path: str, content: str, file_type: str) -> dict:
    """Stage 4: Validate generated code for correctness."""
    messages = [
        {
            "role": "system",
            "content": "You are a code reviewer. Find real bugs and errors. Return JSON only.",
        },
        {
            "role": "user",
            "content": f"""Review this {file_type} file for bugs:

FILE: {file_path}
CONTENT:
{content[:3000]}

Return JSON:
{{
  "valid": true|false,
  "score": 0-100,
  "errors": ["list of actual bugs"],
  "warnings": ["list of style/perf issues"],
  "fix_instructions": "if not valid: specific instructions to fix"
}}
""",
        },
    ]
    raw = await call_llm(messages, max_tokens=512, json_mode=True)
    try:
        return json.loads(raw)
    except Exception:
        return {"valid": True, "score": 80, "errors": [], "warnings": []}


async def stage_autofix(
    file_path: str, content: str, file_type: str, errors: list, fix_instructions: str
) -> str:
    """Stage 5: Auto-fix validation errors."""
    messages = [
        {
            "role": "system",
            "content": "You are a code fixer. Fix the bugs and return only the corrected file content.",
        },
        {
            "role": "user",
            "content": f"""Fix the bugs in this {file_type} file:

FILE: {file_path}
ERRORS TO FIX:
{chr(10).join(f'- {e}' for e in errors)}

INSTRUCTIONS: {fix_instructions}

CURRENT CONTENT:
{content[:3000]}

Return ONLY the corrected file content:
""",
        },
    ]
    return await call_llm(messages, max_tokens=4096)


async def stage_final_validate(files: list[dict], analysis: dict) -> dict:
    """Stage 6: Final cross-file validation."""
    file_list = "\n".join(f"- {f['path']}: {f['description']}" for f in files)
    messages = [
        {
            "role": "system",
            "content": "You are a senior developer doing final review. Return JSON only.",
        },
        {
            "role": "user",
            "content": f"""Final validation of this build:

PROJECT: {analysis.get('project_name')}
TYPE: {analysis.get('project_type')}
FILES GENERATED:
{file_list}

Check:
1. Are all required files present?
2. Is anything missing for deployment?
3. Are imports/dependencies consistent?

Return JSON:
{{
  "ready_to_deploy": true|false,
  "missing_files": ["list or empty"],
  "deployment_steps": ["step 1", "step 2"],
  "summary": "brief summary of what was built"
}}
""",
        },
    ]
    raw = await call_llm(messages, max_tokens=512, json_mode=True)
    try:
        return json.loads(raw)
    except Exception:
        return {"ready_to_deploy": True, "missing_files": [], "deployment_steps": [], "summary": "Build complete"}


# ── In-memory job store ────────────────────────────────────────────────────
JOBS: dict[str, dict] = {}

async def run_pipeline(job_id: str, prompt: str, target: str):
    """Full pipeline runner — updates JOBS[job_id] as it progresses."""
    job = JOBS[job_id]
    job["status"] = "running"
    job["stages"] = []

    def log(stage: str, msg: str, data: dict | None = None):
        entry = {"stage": stage, "message": msg, "ts": time.time()}
        if data:
            entry["data"] = data
        job["stages"].append(entry)
        print(f"[{job_id}] [{stage}] {msg}")

    try:
        # ── Stage 1: Analyze ──
        log("analyze", "Analyzing requirements...")
        analysis = await stage_analyze(prompt, target)
        job["analysis"] = analysis
        log("analyze", "Analysis complete", analysis)

        # ── Stage 2: Plan ──
        log("plan", "Creating build plan...")
        plan = await stage_plan(analysis, prompt)
        job["plan"] = plan
        log("plan", f"Plan ready — {len(plan.get('files', []))} files", plan)

        # ── Stage 3: Generate files ──
        files_to_generate = plan.get("files", [])
        if not files_to_generate:
            # Fallback: generate a single index.html
            files_to_generate = [
                {"path": "index.html", "type": "html", "description": "Main page", "priority": 1}
            ]

        generated_files: list[dict] = []
        job["files"] = []

        for i, file_info in enumerate(files_to_generate[:10]):  # Cap at 10 files
            log("generate", f"Generating {file_info['path']} ({i+1}/{len(files_to_generate)})...")
            try:
                content = await stage_generate_file(
                    file_info, analysis, plan, prompt, target, generated_files
                )

                # ── Stage 4: Validate ──
                log("validate", f"Validating {file_info['path']}...")
                validation = await stage_validate(file_info["path"], content, file_info["type"])

                # ── Stage 5: AutoFix if needed ──
                if not validation.get("valid", True) and validation.get("errors"):
                    log("autofix", f"Auto-fixing {file_info['path']}...")
                    content = await stage_autofix(
                        file_info["path"],
                        content,
                        file_info["type"],
                        validation["errors"],
                        validation.get("fix_instructions", "Fix all errors"),
                    )
                    log("autofix", f"Fixed {file_info['path']}")

                file_result = {
                    "path": file_info["path"],
                    "type": file_info["type"],
                    "description": file_info["description"],
                    "content": content,
                    "validation": validation,
                    "fixed": not validation.get("valid", True),
                }
                generated_files.append(file_result)
                job["files"].append(file_result)
                log("generate", f"✓ {file_info['path']} (score: {validation.get('score', 80)})")

            except Exception as e:
                log("generate", f"✗ {file_info['path']}: {e}")

        # ── Stage 6: Final Validation ──
        log("final_validate", "Running final validation...")
        final = await stage_final_validate(generated_files, analysis)
        job["final_validation"] = final
        log("final_validate", final.get("summary", "Complete"), final)

        job["status"] = "complete"
        job["completed_at"] = time.time()

    except Exception as e:
        job["status"] = "failed"
        job["error"] = str(e)
        log("error", f"Pipeline failed: {e}")


# ── FastAPI App ────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    global DOCS
    DOCS = load_docs()
    print(f"✅ Docs brain loaded: {list(DOCS.keys())}")
    yield

app = FastAPI(
    title="InStatic CMS",
    description="AI-powered build pipeline with WordPress, GitHub, Android, and FastAPI docs brain",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Schemas ────────────────────────────────────────────────────────────────

class BuildRequest(BaseModel):
    prompt: str = Field(..., min_length=10, max_length=4000, description="What to build")
    target: Literal["website", "wordpress", "android", "api", "fullstack"] = "website"

class ChatRequest(BaseModel):
    messages: list[dict]
    stream: bool = False


# ── Routes ────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "service": "instatic-cms",
        "docs_loaded": list(DOCS.keys()),
        "providers": [p["name"] for p in PROVIDERS if os.getenv(p["key_env"])],
    }

@app.get("/docs-brain")
async def list_docs():
    """List available documentation brain files."""
    return {
        "docs": [
            {"name": k, "size": len(v), "preview": v[:200]}
            for k, v in DOCS.items()
        ]
    }

@app.get("/docs-brain/{doc_name}")
async def get_doc(doc_name: str):
    """Get a specific doc from the brain."""
    if doc_name not in DOCS:
        raise HTTPException(404, f"Doc '{doc_name}' not found. Available: {list(DOCS.keys())}")
    return {"name": doc_name, "content": DOCS[doc_name]}


@app.post("/build")
async def start_build(req: BuildRequest, background_tasks: BackgroundTasks):
    """Start an AI build pipeline job. Returns job_id immediately."""
    import uuid
    job_id = str(uuid.uuid4())[:8]
    JOBS[job_id] = {
        "id": job_id,
        "prompt": req.prompt,
        "target": req.target,
        "status": "queued",
        "created_at": time.time(),
        "stages": [],
        "files": [],
    }
    background_tasks.add_task(run_pipeline, job_id, req.prompt, req.target)
    return {"job_id": job_id, "status": "queued", "message": "Pipeline started"}


@app.get("/build/{job_id}")
async def get_build(job_id: str):
    """Get build job status and results."""
    if job_id not in JOBS:
        raise HTTPException(404, f"Job {job_id} not found")
    return JOBS[job_id]


@app.get("/build/{job_id}/stream")
async def stream_build(job_id: str):
    """Stream build progress as SSE."""
    if job_id not in JOBS:
        raise HTTPException(404, f"Job {job_id} not found")

    async def event_stream():
        last_stage_count = 0
        while True:
            job = JOBS.get(job_id, {})
            stages = job.get("stages", [])

            # Send new stages
            for stage in stages[last_stage_count:]:
                yield f"data: {json.dumps(stage)}\n\n"
            last_stage_count = len(stages)

            if job.get("status") in ("complete", "failed"):
                yield f"data: {json.dumps({'stage': 'done', 'status': job['status'], 'job': job})}\n\n"
                break

            await asyncio.sleep(0.5)

    return StreamingResponse(event_stream(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.get("/builds")
async def list_builds():
    """List all build jobs."""
    return {
        "builds": [
            {
                "id": j["id"],
                "prompt": j["prompt"][:80],
                "target": j["target"],
                "status": j["status"],
                "files": len(j.get("files", [])),
                "created_at": j["created_at"],
            }
            for j in sorted(JOBS.values(), key=lambda x: x["created_at"], reverse=True)
        ]
    }


@app.post("/chat")
async def chat(req: ChatRequest):
    """Direct LLM chat with docs brain context."""
    # Inject docs as system context
    doc_names = [m.get("doc") for m in req.messages if isinstance(m, dict) and m.get("doc")]
    doc_ctx = ""
    for name in doc_names:
        if name in DOCS:
            doc_ctx += f"\n\n## {name.upper()} DOCS:\n{DOCS[name][:2000]}"

    system = f"You are an expert developer assistant for the DOLOR3V / Traveler Dev Studio ecosystem. You have access to documentation for WordPress, GitHub, Android, and FastAPI.{doc_ctx}"

    messages = [{"role": "system", "content": system}] + [
        m for m in req.messages if m.get("role") in ("user", "assistant")
    ]

    if req.stream:
        async def gen():
            async for chunk in stream_llm(messages):
                yield f"data: {json.dumps({'content': chunk})}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(gen(), media_type="text/event-stream")

    content = await call_llm(messages)
    return {"content": content}


@app.post("/analyze")
async def analyze_prompt(req: BuildRequest):
    """Just run the analysis stage — useful for previewing before building."""
    analysis = await stage_analyze(req.prompt, req.target)
    plan = await stage_plan(analysis, req.prompt)
    return {"analysis": analysis, "plan": plan}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)