"""
FastAPI app — JustEmbed server.
"""

import re
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, Form, Request, UploadFile, HTTPException
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.datastructures import FormData

from justembed.config import get_workspace, register_workspace, deregister_workspace, list_registered_workspaces, ensure_workspace_structure as ensure_ws
from justembed import db
from justembed.chunker import chunk_text
from justembed.embedder import E5Embedder, CustomEmbedder
from justembed.training.trainer import CustomModelTrainer

# Lazy embedder
_embedder: Optional[E5Embedder] = None
_custom_embedders: dict[str, CustomEmbedder] = {}


def get_embedder(model_type: str = "e5", model_name: str = "e5-small"):
    """Get embedder by type and name."""
    if model_type == "e5":
        global _embedder
        if _embedder is None:
            _embedder = E5Embedder()
        return _embedder
    elif model_type == "custom":
        if model_name not in _custom_embedders:
            _custom_embedders[model_name] = CustomEmbedder(model_name)
        return _custom_embedders[model_name]
    else:
        raise ValueError(f"Unknown model type: {model_type}")


# File size limit (10MB)
MAX_FILE_SIZE = 10 * 1024 * 1024

# Temporary storage for upload content (in-memory cache)
_upload_cache: dict[str, dict] = {}


def list_custom_models() -> list[dict]:
    """List available custom models."""
    from pathlib import Path
    import json
    from justembed.config import get_custom_models_dir
    
    try:
        models_dir = get_custom_models_dir()
    except FileNotFoundError:
        # Workspace not configured yet
        return []
    
    if not models_dir.exists():
        return []
    
    models = []
    for model_dir in models_dir.iterdir():
        if model_dir.is_dir():
            config_path = model_dir / "config.json"
            if config_path.exists():
                try:
                    with open(config_path) as f:
                        config = json.load(f)
                    models.append({
                        "name": model_dir.name,
                        "embedding_dim": config.get("embedding_dim", 128),
                        "created_at": config.get("created_at", ""),
                        "num_chunks": config.get("corpus_stats", {}).get("num_chunks", 0),
                    })
                except Exception:
                    pass
    return models


def create_app() -> FastAPI:
    app = FastAPI(title="JustEmbed", version="0.1.1a0")

    # Templates and static
    pkg_root = Path(__file__).resolve().parent
    templates = Jinja2Templates(directory=str(pkg_root / "templates"))
    if (pkg_root / "static").exists():
        app.mount("/static", StaticFiles(directory=str(pkg_root / "static")), name="static")

    @app.get("/", response_class=HTMLResponse)
    async def index(request: Request):
        try:
            workspace = get_workspace()
        except Exception:
            return templates.TemplateResponse("setup.html", {"request": request})
        
        # Get list of all registered workspaces
        registered_workspaces = list_registered_workspaces()
        
        # Get KB list with model info and size
        kb_names = db.list_kbs(workspace)
        kbs = []
        for kb_name in kb_names:
            try:
                conn = db.get_connection(workspace, kb_name)
                model_type = db.get_kb_metadata(conn, "model_type", "e5")
                model_name = db.get_kb_metadata(conn, "model_name", "e5-small")
                
                # Get KB size (count chunks)
                chunk_count = conn.execute("SELECT COUNT(*) FROM chunks WHERE embedding IS NOT NULL").fetchone()[0]
                conn.close()
                
                kbs.append({
                    "name": kb_name,
                    "model_type": model_type,
                    "model_name": model_name,
                    "chunk_count": chunk_count,
                })
            except Exception:
                # If we can't read metadata, assume default
                kbs.append({
                    "name": kb_name,
                    "model_type": "e5",
                    "model_name": "e5-small",
                    "chunk_count": 0,
                })
        
        custom_models = list_custom_models()
        
        # Get upload history
        history, total_purged = db.get_upload_history(workspace)
        
        # Check if we should highlight a newly created KB
        new_kb = request.query_params.get("new_kb")
        
        return templates.TemplateResponse(
            "index.html",
            {
                "request": request,
                "workspace": str(workspace),
                "registered_workspaces": registered_workspaces,
                "kbs": kbs,
                "custom_models": custom_models,
                "upload_history": history,
                "total_purged": total_purged,
                "new_kb": new_kb,
            },
        )

    @app.post("/set-workspace")
    async def api_set_workspace(workspace: str = Form(...)):
        p = Path(workspace).resolve()
        register_workspace(p)
        return RedirectResponse(url="/", status_code=303)

    @app.post("/register-workspace")
    async def api_register_workspace(workspace: str = Form(...)):
        """Register a new workspace (doesn't delete data)."""
        p = Path(workspace).resolve()
        register_workspace(p)
        return RedirectResponse(url="/", status_code=303)

    @app.post("/deregister-workspace")
    async def api_deregister_workspace(confirm_path: str = Form(...)):
        """Deregister workspace (data stays on disk)."""
        try:
            current_ws = get_workspace()
        except Exception:
            raise HTTPException(400, "No workspace configured")
        
        # Verify confirmation path matches
        if Path(confirm_path).resolve() != current_ws:
            raise HTTPException(400, "Confirmation path does not match current workspace")
        
        # Deregister the workspace (data stays on disk)
        deregister_workspace(current_ws)
        
        # Redirect to setup page
        return RedirectResponse(url="/", status_code=303)

    @app.post("/create-kb")
    async def api_create_kb(
        name: str = Form(...),
        model_type: str = Form("e5"),
        model_name: str = Form("e5-small"),
    ):
        ws = get_workspace()
        
        # Validate KB name
        if not name or not name.strip():
            raise HTTPException(400, "KB name cannot be empty")
        
        name = name.strip()
        if not re.match(r"^[a-zA-Z0-9_]{1,64}$", name):
            raise HTTPException(
                400, 
                "Invalid KB name. Use only alphanumeric characters and underscores (max 64 chars)"
            )
        
        # Check if KB already exists
        if name in db.list_kbs(ws):
            raise HTTPException(400, f"Knowledge base '{name}' already exists")
        
        try:
            db.create_kb(ws, name, model_type=model_type, model_name=model_name)
        except ValueError as e:
            raise HTTPException(400, str(e))
        
        # Redirect to home with new_kb parameter to stay on KB tab and highlight
        return RedirectResponse(url=f"/?tab=kbs&new_kb={name}", status_code=303)

    @app.post("/delete-kb")
    async def api_delete_kb(name: str = Form(...)):
        """Delete a knowledge base."""
        ws = get_workspace()
        
        try:
            db.delete_kb(ws, name)
        except FileNotFoundError:
            raise HTTPException(404, f"Knowledge base '{name}' not found")
        except Exception as e:
            raise HTTPException(500, f"Failed to delete KB: {str(e)}")
        
        return RedirectResponse(url="/", status_code=303)

    @app.post("/upload")
    async def api_upload(
        file: UploadFile = ...,
        kb: str = Form(...),
    ):
        """Upload file and redirect to chunk preview."""
        if not file.filename or not file.filename.lower().endswith((".txt", ".md")):
            raise HTTPException(400, "Only .txt and .md files allowed")

        # Read file content
        content_bytes = await file.read()
        
        # Check file size
        if len(content_bytes) > MAX_FILE_SIZE:
            raise HTTPException(400, f"File too large. Maximum size is {MAX_FILE_SIZE // (1024*1024)}MB")
        
        if len(content_bytes) == 0:
            raise HTTPException(400, "File is empty")
        
        try:
            content = content_bytes.decode("utf-8", errors="replace")
        except Exception as e:
            raise HTTPException(400, f"Failed to decode file: {str(e)}")
        
        ws = get_workspace()
        ensure_ws(ws)

        # Validate KB exists
        try:
            db.get_connection(ws, kb).close()
        except FileNotFoundError:
            raise HTTPException(404, f"Knowledge base '{kb}' not found")

        # Store in cache with unique key
        import uuid
        cache_key = str(uuid.uuid4())
        _upload_cache[cache_key] = {
            "content": content,
            "filename": file.filename or "upload.txt",
            "kb": kb,
        }

        # Redirect to preview
        return RedirectResponse(url=f"/preview-chunks-ui?cache_key={cache_key}", status_code=303)

    @app.get("/upload-to-kb", response_class=HTMLResponse)
    async def ui_upload_to_kb(request: Request, kb: str):
        """Show upload form for a specific KB."""
        ws = get_workspace()
        
        # Validate KB exists
        try:
            conn = db.get_connection(ws, kb)
            model_type = db.get_kb_metadata(conn, "model_type", "e5")
            model_name = db.get_kb_metadata(conn, "model_name", "e5-small")
            conn.close()
        except FileNotFoundError:
            raise HTTPException(404, f"Knowledge base '{kb}' not found")
        
        return templates.TemplateResponse(
            "upload_to_kb.html",
            {
                "request": request,
                "kb_name": kb,
                "model_type": model_type,
                "model_name": model_name,
            },
        )

    @app.post("/preview-chunks")
    async def api_preview_chunks(
        content: str = Form(...),
        filename: str = Form("upload.txt"),
        max_tokens: int = Form(300),
        merge_threshold: int = Form(50),
    ):
        """API endpoint for chunk preview (returns JSON)."""
        chunks = chunk_text(
            content,
            file=filename,
            max_tokens=max(10, min(800, max_tokens)),
            merge_threshold=max(10, merge_threshold),
        )
        return {"chunks": chunks, "count": len(chunks)}

    @app.get("/preview-chunks-ui", response_class=HTMLResponse)
    @app.post("/preview-chunks-ui", response_class=HTMLResponse)
    async def ui_preview_chunks(
        request: Request,
        cache_key: str = Form(None),
        content: str = Form(None),
        filename: str = Form(None),
        kb: str = Form(None),
        max_tokens: int = Form(300),
        merge_threshold: int = Form(50),
        split_by_headings: str = Form("true"),
        split_by_paragraphs: str = Form("true"),
    ):
        """UI endpoint for chunk preview."""
        # Handle GET with cache_key
        if request.method == "GET":
            cache_key = request.query_params.get("cache_key")
            if cache_key and cache_key in _upload_cache:
                cached = _upload_cache[cache_key]
                content = cached["content"]
                filename = cached["filename"]
                kb = cached["kb"]
            else:
                raise HTTPException(400, "Invalid or expired upload session")
            max_tokens = int(request.query_params.get("max_tokens", 300))
            merge_threshold = int(request.query_params.get("merge_threshold", 50))
            split_by_headings = request.query_params.get("split_by_headings", "true")
            split_by_paragraphs = request.query_params.get("split_by_paragraphs", "true")
        else:
            # POST from re-preview form
            if cache_key and cache_key in _upload_cache:
                cached = _upload_cache[cache_key]
                if not content:
                    content = cached["content"]
                if not filename:
                    filename = cached["filename"]
                if not kb:
                    kb = cached["kb"]

        if not content:
            raise HTTPException(400, "No content provided")

        # Convert string bools
        split_headings = split_by_headings.lower() in ("true", "1", "yes", "on")
        split_paras = split_by_paragraphs.lower() in ("true", "1", "yes", "on")

        # Generate chunks
        chunks = chunk_text(
            content,
            file=filename,
            max_tokens=max(10, min(800, max_tokens)),
            merge_threshold=max(10, merge_threshold),
            split_by_headings=split_headings,
            split_by_paragraphs=split_paras,
        )

        # Add token counts
        from justembed.chunker import _word_tokens
        for chunk in chunks:
            chunk["token_count"] = len(_word_tokens(chunk["text"]))

        return templates.TemplateResponse(
            "chunk_preview.html",
            {
                "request": request,
                "chunks": chunks,
                "filename": filename,
                "kb": kb,
                "cache_key": cache_key,
                "max_tokens": max_tokens,
                "merge_threshold": merge_threshold,
                "split_by_headings": split_headings,
                "split_by_paragraphs": split_paras,
            },
        )

    @app.post("/apply-chunking")
    async def api_apply_chunking(
        cache_key: str = Form(...),
        filename: str = Form(...),
        kb: str = Form(...),
        max_tokens: int = Form(300),
        merge_threshold: int = Form(50),
        split_by_headings: str = Form("true"),
        split_by_paragraphs: str = Form("true"),
    ):
        """Apply chunking and save to KB (without embeddings)."""
        # Get content from cache
        if cache_key not in _upload_cache:
            raise HTTPException(400, "Upload session expired. Please upload the file again.")
        
        content = _upload_cache[cache_key]["content"]
        ws = get_workspace()
        
        # Convert string bools
        split_headings = split_by_headings.lower() in ("true", "1", "yes", "on")
        split_paras = split_by_paragraphs.lower() in ("true", "1", "yes", "on")

        # Generate chunks
        chunks = chunk_text(
            content,
            file=filename,
            max_tokens=max(10, min(800, max_tokens)),
            merge_threshold=max(10, merge_threshold),
            split_by_headings=split_headings,
            split_by_paragraphs=split_paras,
        )

        if not chunks:
            raise HTTPException(400, "No chunks produced")

        # Save chunks without embeddings
        try:
            conn = db.get_connection(ws, kb)
        except FileNotFoundError:
            raise HTTPException(404, f"Knowledge base '{kb}' not found")

        db.insert_chunks_without_embeddings(conn, filename, chunks)
        conn.close()

        # Clean up cache
        del _upload_cache[cache_key]

        # Redirect to embed step
        from urllib.parse import urlencode
        params = urlencode({"filename": filename, "kb": kb})
        return RedirectResponse(url=f"/embed?{params}", status_code=303)

    @app.get("/embed", response_class=HTMLResponse)
    async def ui_embed(
        request: Request,
        filename: str,
        kb: str,
    ):
        """Embed chunks and show progress."""
        ws = get_workspace()
        
        try:
            conn = db.get_connection(ws, kb)
        except FileNotFoundError:
            raise HTTPException(404, f"Knowledge base '{kb}' not found")

        # Get KB model info
        model_type = db.get_kb_metadata(conn, "model_type", "e5")
        model_name = db.get_kb_metadata(conn, "model_name", "e5-small")

        # Get chunks without embeddings
        chunks_to_embed = db.get_chunks_without_embeddings(conn, filename)
        
        if chunks_to_embed:
            # Get appropriate embedder
            emb = get_embedder(model_type=model_type, model_name=model_name)
            embeddings = emb.embed([c["text"] for c in chunks_to_embed])
            
            # Update with embeddings
            db.update_embeddings(conn, chunks_to_embed, embeddings)
        
        conn.close()

        # Add to upload history
        purged = db.add_upload_history(ws, filename, "kb_upload", kb_name=kb)

        return templates.TemplateResponse(
            "embed_progress.html",
            {
                "request": request,
                "filename": filename,
                "kb": kb,
                "total_chunks": len(chunks_to_embed),
                "model_type": model_type,
                "model_name": model_name,
            },
        )

    @app.post("/query", response_class=HTMLResponse)
    async def api_query(
        request: Request,
        query: str = Form(...),
        kbs: str = Form("all"),
        top_k: int = Form(5),
        mode: str = Form("retrieve"),
    ):
        if not query.strip():
            raise HTTPException(400, "Query required")

        ws = get_workspace()
        kb_list = [k.strip() for k in kbs.split(",") if k.strip()] if kbs else []
        if not kb_list or (kbs or "").strip().lower() == "all":
            kb_list = db.list_kbs(ws)

        results: list[dict] = []
        by_kb: list[dict] = []
        
        for kb_name in kb_list:
            try:
                conn = db.get_connection(ws, kb_name)
            except FileNotFoundError:
                continue
            
            # Get KB model info
            model_type = db.get_kb_metadata(conn, "model_type", "e5")
            model_name = db.get_kb_metadata(conn, "model_name", "e5-small")
            
            # Get appropriate embedder
            emb = get_embedder(model_type=model_type, model_name=model_name)
            q_emb = emb.embed_query(query)
            
            if mode == "count":
                all_matches = db.query_chunks(conn, q_emb, top_k=1000, kb_name=kb_name)
                by_kb.append({"kb": kb_name, "count": len(all_matches), "model": f"{model_type}:{model_name}"})
            else:
                rows = db.query_chunks(conn, q_emb, top_k=top_k, kb_name=kb_name)
                # Add model info to results
                for row in rows:
                    row["model"] = f"{model_type}:{model_name}"
                results.extend(rows)
            conn.close()

        if mode == "count":
            total = sum(r["count"] for r in by_kb)
            return templates.TemplateResponse(
                "query_results.html",
                {
                    "request": request,
                    "mode": "count",
                    "query": query,
                    "total": total,
                    "by_kb": by_kb,
                },
            )
        
        # Sort results by score (highest first) across all KBs
        results.sort(key=lambda x: x.get("score", 0), reverse=True)
        
        return templates.TemplateResponse(
            "query_results.html",
            {
                "request": request,
                "mode": "retrieve",
                "query": query,
                "results": results,
                "count": len(results),
            },
        )

    @app.get("/status")
    async def api_status():
        try:
            ws = get_workspace()
        except Exception:
            return {"workspace": None, "kbs": []}
        return {"workspace": str(ws), "kbs": db.list_kbs(ws)}

    @app.get("/train-model", response_class=HTMLResponse)
    async def ui_train_model(request: Request):
        """UI for training custom models."""
        return templates.TemplateResponse("train_model.html", {"request": request})

    @app.post("/train-model")
    async def api_train_model(
        model_name: str = Form(...),
        file: UploadFile = ...,
        embedding_dim: int = Form(128),
        max_features: int = Form(5000),
    ):
        """Train custom model from uploaded file."""
        # Validate model name
        if not model_name or not model_name.replace("_", "").isalnum():
            raise HTTPException(400, "Invalid model name. Use alphanumeric and underscore only.")
        
        # Validate file
        if not file.filename or not file.filename.lower().endswith((".txt", ".md")):
            raise HTTPException(400, "Only .txt and .md files allowed")
        
        # Read file content
        content_bytes = await file.read()
        
        if len(content_bytes) > MAX_FILE_SIZE:
            raise HTTPException(400, f"File too large. Maximum size is {MAX_FILE_SIZE // (1024*1024)}MB")
        
        if len(content_bytes) == 0:
            raise HTTPException(400, "File is empty")
        
        try:
            content = content_bytes.decode("utf-8", errors="replace")
        except Exception as e:
            raise HTTPException(400, f"Failed to decode file: {str(e)}")
        
        # Split into paragraphs (simple chunking for training)
        # Try double newline first, then single newline, then sentence-based chunking
        paragraphs = [p.strip() for p in content.split("\n\n") if p.strip()]
        
        # If no double-newline paragraphs, try single newlines
        if len(paragraphs) < 3:
            paragraphs = [p.strip() for p in content.split("\n") if p.strip() and len(p.strip()) > 50]
        
        # If still too few, split by sentences (periods followed by space/newline)
        if len(paragraphs) < 3:
            import re
            # Split on period + space/newline, keep sentences with 20+ words
            sentences = re.split(r'\.\s+', content)
            paragraphs = [s.strip() + '.' for s in sentences if s.strip() and len(s.split()) >= 20]
        
        # If still too few, chunk by word count (every ~100 words)
        if len(paragraphs) < 3:
            words = content.split()
            chunk_size = max(100, len(words) // 10)  # At least 100 words per chunk, or 10 chunks
            paragraphs = []
            for i in range(0, len(words), chunk_size):
                chunk = ' '.join(words[i:i+chunk_size])
                if chunk.strip():
                    paragraphs.append(chunk.strip())
        
        if len(paragraphs) < 3:
            raise HTTPException(400, f"Could not create enough chunks for training. Got {len(paragraphs)} chunks, need at least 3. Try adding paragraph breaks or providing more text.")
        
        # Train model
        try:
            trainer = CustomModelTrainer()
            model_dir = trainer.train(
                corpus=paragraphs,
                model_name=model_name,
                embedding_dim=embedding_dim,
                max_features=max_features,
            )
        except ValueError as e:
            raise HTTPException(400, f"Training failed: {str(e)}")
        except Exception as e:
            raise HTTPException(500, f"Training error: {str(e)}")
        
        # Add to upload history
        ws = get_workspace()
        purged = db.add_upload_history(ws, file.filename, "model_training", model_name=model_name)
        
        # Redirect to success page
        from urllib.parse import urlencode
        params = urlencode({"model_name": model_name, "model_dir": str(model_dir)})
        return RedirectResponse(url=f"/train-success?{params}#success", status_code=303)

    @app.get("/train-success", response_class=HTMLResponse)
    async def ui_train_success(
        request: Request,
        model_name: str,
        model_dir: str,
    ):
        """Show training success page."""
        import json
        from pathlib import Path
        
        # Load model config
        config_path = Path(model_dir) / "config.json"
        with open(config_path) as f:
            config = json.load(f)
        
        return templates.TemplateResponse(
            "train_success.html",
            {
                "request": request,
                "model_name": model_name,
                "model_dir": model_dir,
                "config": config,
            },
        )

    return app
