import os

def write_file(path, content):
    full_path = os.path.abspath(os.path.join(os.getcwd(), "..", "backend", path))
    os.makedirs(os.path.dirname(full_path), exist_ok=True)
    with open(full_path, "w", encoding="utf-8") as f:
        f.write(content)

# 1. Update api/chatbot/chat_route.py
write_file("api/chatbot/chat_route.py", """
\"\"\"
Chatbot API Route
-----------------
`POST /chat` - endpoint chat cho CuCu Assistant (non-streaming).
`POST /chat/stream` - endpoint chat SSE streaming.
Logic xử lý nằm ở `agent.controller`.
\"\"\"

import logging
from dataclasses import dataclass
from typing import AsyncGenerator

from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, Body
from fastapi.responses import StreamingResponse
from opentelemetry import trace

# Lazy imports - defer heavy AI modules to first use
from common.message_limit import message_limit_service
from config import DEFAULT_MODEL

logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)
router = APIRouter()

# Cache for lazy-loaded modules
_agent_modules = {}


def _get_agent_controller():
    \"\"\"Lazy-load agent.controller to defer LangChain/LangGraph import.\"\"\"
    if "controller" not in _agent_modules:
        from agent.controller import chat_controller, chat_controller_stream
        _agent_modules["controller"] = chat_controller
        _agent_modules["controller_stream"] = chat_controller_stream
    return _agent_modules["controller"], _agent_modules["controller_stream"]


def _get_query_request_model():
    \"\"\"Lazy-load agent.models.QueryRequest.\"\"\"
    if "QueryRequest" not in _agent_modules:
        from agent.models import QueryRequest
        _agent_modules["QueryRequest"] = QueryRequest
    return _agent_modules["QueryRequest"]


@dataclass
class Identity:
  \"\"\"Định danh đơn giản dựa trên middleware (user_id / device_id).\"\"\"
  primary_id: str
  history_key: str
  rate_limit_key: str
  is_authenticated: bool


def _get_identity(request: Request) -> Identity:
  \"\"\"Derive identity info từ request.state (do middleware set).\"\"\"
  user_id = getattr(request.state, "user_id", None)
  device_id = getattr(request.state, "device_id", "") or ""
  is_authenticated = bool(getattr(request.state, "is_authenticated", False))

  if is_authenticated and user_id:
    primary_id = str(user_id)
    history_key = primary_id
    rate_limit_key = primary_id
  else:
    primary_id = device_id or "anonymous"
    history_key = device_id or "anonymous"
    rate_limit_key = device_id or "anonymous"

  return Identity(
    primary_id=primary_id,
    history_key=history_key,
    rate_limit_key=rate_limit_key,
    is_authenticated=is_authenticated,
  )


@router.post("/chat", summary="Chat with CuCu Assistant")
async def cucu_chat(request: Request, background_tasks: BackgroundTasks):
  \"\"\"Endpoint chat không stream - trả về response JSON đầy đủ một lần.\"\"\"
  # Lazy-load AI modules on first call
  QueryRequest = _get_query_request_model()
  chat_controller, _ = _get_agent_controller()

  try:
    body = await request.json()
    if 'message' in body and 'user_query' not in body:
        body['user_query'] = body.pop('message')
    
    if not body.get('user_query'):
        body['user_query'] = 'xin chào'

    req = QueryRequest(**body)
  except Exception as e:
    logger.error(f"Validation error in chat: {e}")
    raise HTTPException(status_code=422, detail=f"Invalid request body: {str(e)}")

  identity = _get_identity(request)
  user_id = identity.primary_id

  logger.info("📥 [Incoming Chat] User=%s | Query=%s", user_id, req.user_query)

  span = trace.get_current_span()
  span.set_attribute("user.id", user_id)
  span.set_attribute("chat.user_query", req.user_query)

  try:
    result = await chat_controller(
      query=req.user_query,
      user_id=user_id,
      background_tasks=background_tasks,
      model_name=DEFAULT_MODEL,
      images=req.images,
      identity_key=identity.history_key,
    )

    usage_info = await message_limit_service.increment(
      identity_key=identity.rate_limit_key,
      is_authenticated=identity.is_authenticated,
    )

    logger.info("📤 [Outgoing Chat] User=%s", user_id)

    return {
      "status": "success",
      "ai_response": result["ai_response"],
      "product_ids": result.get("product_ids", []),
      "limit_info": {
        "limit": usage_info["limit"],
        "used": usage_info["used"],
        "remaining": usage_info["remaining"],
      },
    }
  except Exception as e:
    logger.error("Error in cucu_chat: %s", e, exc_info=True)
    raise HTTPException(status_code=500, detail=str(e)) from e


@router.post("/chat/stream", summary="Chat with CuCu Assistant (SSE Streaming)")
async def cucu_chat_stream(request: Request):
  \"\"\"
  Endpoint chat SSE streaming — trả về token-by-token qua Server-Sent Events.
  \"\"\"
  # Lazy-load AI modules on first call
  QueryRequest = _get_query_request_model()
  _, chat_controller_stream = _get_agent_controller()

  try:
    body = await request.json()
    if 'message' in body and 'user_query' not in body:
        body['user_query'] = body.pop('message')
    
    if not body.get('user_query'):
        body['user_query'] = 'xin chào'

    req = QueryRequest(**body)
  except Exception as e:
    raise HTTPException(status_code=422, detail=str(e))

  identity = _get_identity(request)
  user_id = identity.primary_id

  logger.info("📥 [Incoming Stream] User=%s | Query=%s", user_id, req.user_query)

  async def sse_generator() -> AsyncGenerator[str, None]:
    try:
      async for chunk_json in chat_controller_stream(
        query=req.user_query,
        user_id=user_id,
        model_name=DEFAULT_MODEL,
        images=req.images,
        identity_key=identity.history_key,
      ):
        yield f"data: {chunk_json}\\n\\n"

      # Increment rate limit after stream completes
      await message_limit_service.increment(
        identity_key=identity.rate_limit_key,
        is_authenticated=identity.is_authenticated,
      )
      logger.info("📤 [Stream Done] User=%s", user_id)
    except Exception as e:
      logger.error("Error in stream: %s", e, exc_info=True)
      yield f'data: {{"error": "{str(e)}"}}\\n\\n'

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

# 2. Update api/memos/__init__.py
write_file("api/memos/__init__.py", """
from fastapi import APIRouter

from .instance_routes import router as instance_router
from .auth_routes import router as auth_router
from .user_routes import router as user_router
from .memo_routes import router as memo_router
from .attachment_routes import router as attachment_router
from .shortcut_routes import router as shortcut_router
from .activity_routes import router as activity_router
from .idp_routes import router as idp_router
from .embedding_routes import router as embedding_router
from .search_routes import router as search_router


router = APIRouter()


router.include_router(instance_router, tags=["instance"])
router.include_router(auth_router, tags=["auth"])
router.include_router(user_router, tags=["users"])
# search_router MUST come before memo_router
router.include_router(search_router, tags=["search"])
router.include_router(memo_router, tags=["memos"])
router.include_router(attachment_router, tags=["attachments"])
router.include_router(shortcut_router, tags=["shortcuts"])
router.include_router(activity_router, tags=["activities"])
router.include_router(idp_router, tags=["idp"])
router.include_router(embedding_router, tags=["memo-embeddings"])
""")

# 3. Update api/auth_routes.py
write_file("api/auth_routes.py", """
from datetime import datetime, timezone
from typing import Any

from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, EmailStr

from common.jwt_auth import (
    create_access_token,
    create_refresh_token,
    get_password_hash,
    verify_password,
    decode_token,
)
from common.sqlite_client import sqlite_client as db

router = APIRouter(tags=["auth"])


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


class UserRegister(BaseModel):
    username: str
    email: EmailStr
    password: str


class UserLogin(BaseModel):
    username: str
    password: str


class RefreshTokenRequest(BaseModel):
    refresh_token: str


@router.post("/register")
async def register(user_data: UserRegister):
    # Check if username already exists
    existing = await db.fetch_one(
        "SELECT id FROM cuccu_users WHERE username = ?", (user_data.username,)
    )
    if existing:
        raise HTTPException(status_code=400, detail="Username already registered")

    # Check if email already exists
    existing_email = await db.fetch_one(
        "SELECT id FROM cuccu_users WHERE email = ?", (user_data.email,)
    )
    if existing_email:
        raise HTTPException(status_code=400, detail="Email already registered")

    hashed_password = get_password_hash(user_data.password)
    now = utc_now()

    # Insert new user
    cursor = await db.execute(
        "INSERT INTO cuccu_users (username, email, password_hash, nickname, role, created_at) VALUES (?, ?, ?, ?, ?, ?)",
        (user_data.username, user_data.email, hashed_password, user_data.username, "user", now),
    )
    user_id = str(cursor.lastrowid)

    access_token = create_access_token(data={"sub": user_id})
    refresh_token, expires_at = create_refresh_token(data={"sub": user_id})

    await db.execute(
        "INSERT INTO cuccu_refresh_tokens (user_id, token, expires_at, created_at) VALUES (?, ?, ?, ?)",
        (user_id, refresh_token, expires_at.isoformat(), now),
    )

    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "user": {
            "id": user_id,
            "username": user_data.username,
            "email": user_data.email,
        },
    }


@router.post("/login")
async def login(user_data: UserLogin):
    user = await db.fetch_one(
        "SELECT * FROM cuccu_users WHERE username = ?", (user_data.username,)
    )
    if not user:
        raise HTTPException(status_code=400, detail="Incorrect username or password")

    if not verify_password(user_data.password, user["password_hash"]):
        raise HTTPException(status_code=400, detail="Incorrect username or password")

    user_id = str(user["id"])
    access_token = create_access_token(data={"sub": user_id})
    refresh_token, expires_at = create_refresh_token(data={"sub": user_id})
    now = utc_now()

    await db.execute(
        "INSERT INTO cuccu_refresh_tokens (user_id, token, expires_at, created_at) VALUES (?, ?, ?, ?)",
        (user_id, refresh_token, expires_at.isoformat(), now),
    )

    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "user": {
            "id": user_id,
            "username": user["username"],
            "email": user["email"],
        },
    }


@router.post("/refresh")
async def refresh(data: RefreshTokenRequest):
    token = data.refresh_token
    payload = decode_token(token)
    if not payload or "sub" not in payload:
        raise HTTPException(status_code=401, detail="Invalid refresh token")

    user_id = payload["sub"]
    stored = await db.fetch_one(
        "SELECT id FROM cuccu_refresh_tokens WHERE token = ? AND user_id = ?",
        (token, user_id),
    )
    if not stored:
        raise HTTPException(status_code=401, detail="Refresh token revoked or invalid")

    access_token = create_access_token(data={"sub": user_id})
    return {"access_token": access_token, "token_type": "bearer"}


@router.get("/me")
async def get_me(request: Request):
    user_id = getattr(request.state, "user_id", None)
    import os
    if os.getenv('DISABLE_AUTH', 'false').lower() == 'true':
        return {
            "id": user_id or "testuser_local",
            "username": user_id or "testuser_local",
            "email": "test@example.com",
            "role": "user"
        }

    if not getattr(request.state, "is_authenticated", False):
        raise HTTPException(status_code=401, detail="Not authenticated")

    user = await db.fetch_one(
        "SELECT id, username, email, role FROM cuccu_users WHERE id = ?", (user_id,)
    )
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    return {
        "id": str(user["id"]),
        "username": user["username"],
        "email": user["email"],
        "role": user["role"] or "user",
    }
""")

# 4. Update api/team_routes.py
write_file("api/team_routes.py", """
from typing import List
from fastapi import APIRouter, Body, HTTPException, Request
from common.team.schemas import (
    TeamCreate, TeamUpdate, TeamResponse,
    TeamMemoCreate, TeamMemoResponse, TeamCommentCreate, TeamCommentResponse,
    ReviewAction, ReactionToggle
)
from common.team.services import team_service

router = APIRouter(tags=["teams"])

def require_user(request: Request) -> str:
    user_id = getattr(request.state, "user_id", None)
    import os
    if os.getenv('DISABLE_AUTH', 'false').lower() == 'true':
        return user_id or 'testuser_local'
    if not user_id:
        raise HTTPException(status_code=401, detail="Authentication required")
    return user_id

@router.post("", response_model=TeamResponse)
async def create_team(request: Request, payload: TeamCreate):
    user_id = require_user(request)
    return await team_service.create_team(payload.name, payload.description, user_id)

@router.get("", response_model=List[TeamResponse])
async def list_teams(request: Request):
    user_id = require_user(request)
    return await team_service.list_teams(user_id)

@router.get("/{team_id}", response_model=TeamResponse)
async def get_team(request: Request, team_id: str):
    user_id = require_user(request)
    result = await team_service.get_team(team_id, user_id)
    if not result: raise HTTPException(status_code=404, detail="Team not found")
    return result
""")

# 5. Update server.py
write_file("server.py", """
import asyncio
import os
import platform
import logging
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import RedirectResponse, ORJSONResponse

from api.chatbot import router as chatbot_router
from api.memos import router as memos_router
from api.auth_routes import router as auth_router
from api.test_chat_route import router as test_router
from api.team_routes import router as team_router
from api.inbound_webhook_routes import router as webhook_router
from api.documents.documents_routes import router as documents_router
from common.cache import redis_cache
from common.langfuse_client import get_langfuse_client
from common.meili_service import meili_service
from common.middleware import middleware_manager
from config import PORT, CORS_ORIGINS, DISABLE_AUTH


if platform.system() == "Windows":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@asynccontextmanager
async def lifespan(app: FastAPI):
    await redis_cache.initialize()
    from common.sqlite_client import init_sqlite
    await init_sqlite()
    await meili_service.initialize()
    yield
    from common.sqlite_client import close_sqlite
    await close_sqlite()

app = FastAPI(title="CuCu Note API", lifespan=lifespan, default_response_class=ORJSONResponse)

middleware_manager.setup(app, enable_auth=True, enable_cors=True, cors_origins=CORS_ORIGINS)

# Consolidate all /api/v1 routes
app.include_router(test_router, prefix="/api/test")
app.include_router(auth_router, prefix="/api/v1")
app.include_router(chatbot_router, prefix="/api/agent")
app.include_router(memos_router, prefix="/api/v1")
app.include_router(team_router, prefix="/api/v1")
app.include_router(webhook_router, prefix="/api/v1")
app.include_router(documents_router, prefix="/api/v1")

@app.get("/health")
async def health(): return {"status": "ok"}

@app.get("/")
async def root(): return RedirectResponse(url="/static/index.html")

if __name__ == "__main__":
    uvicorn.run("server:app", host="0.0.0.0", port=PORT, reload=True)
""")
