import os
import asyncio
import uuid
import hashlib
from datetime import datetime, timedelta
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials, OAuth2PasswordBearer
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Annotated
from dotenv import load_dotenv
import logging

# Define authentication schemes early
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="auth/token")

# Import agents module with error handling
try:
    from agents import Runner
    AGENTS_AVAILABLE = True
    print("Agents module loaded successfully")
except ImportError as e:
    print(f"Warning: agents module import failed: {e}")
    Runner = None
    AGENTS_AVAILABLE = False
except AttributeError as e:
    print(f"Warning: agents module has attribute error (likely TensorFlow issue): {e}")
    Runner = None
    AGENTS_AVAILABLE = False
except Exception as e:
    print(f"Warning: agents module failed with unexpected error: {e}")
    Runner = None
    AGENTS_AVAILABLE = False
from passlib.context import CryptContext
from jose import JWTError, jwt 
import json

# Load environment variables
load_dotenv()

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

# Password hashing context
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# JWT configuration
SECRET_KEY = os.getenv("SECRET_KEY")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

# User model
class UserBase(BaseModel):
    email: str
    name: Optional[str] = None

class UserCreate(UserBase):
    password: str
    # Hardware background
    hardware_type: Optional[str] = None
    jetson_model: Optional[str] = None
    components: Optional[List[str]] = []
    # Software background
    software_stack: Optional[List[str]] = []
    experience_level: Optional[str] = None

class UserUpdate(BaseModel):
    name: Optional[str] = None
    # Hardware background
    hardware_type: Optional[str] = None
    jetson_model: Optional[str] = None
    components: Optional[List[str]] = []
    # Software background
    software_stack: Optional[List[str]] = None
    experience_level: Optional[str] = None

class UserInDB(UserBase):
    id: str
    hashed_password: str
    created_at: datetime
    updated_at: datetime

class UserPublic(UserBase):
    id: str
    created_at: datetime
    updated_at: datetime

class UserWithProfile(UserPublic):
    # Hardware background
    hardware_type: Optional[str] = None
    jetson_model: Optional[str] = None
    components: Optional[List[str]] = []
    # Software background
    software_stack: Optional[List[str]] = None
    experience_level: Optional[str] = None

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    email: Optional[str] = None

# Import the existing RAG agent functionality
# Import Myagent with error handling
try:
    from chat import Myagent
    MYAGENT_AVAILABLE = True
    print("Myagent module loaded successfully")
except ImportError as e:
    print(f"Warning: Myagent module import failed: {e}")
    Myagent = None
    MYAGENT_AVAILABLE = False
except AttributeError as e:
    print(f"Warning: Myagent module has attribute error (likely TensorFlow issue): {e}")
    Myagent = None
    MYAGENT_AVAILABLE = False
except Exception as e:
    print(f"Warning: Myagent module failed with unexpected error: {e}")
    Myagent = None
    MYAGENT_AVAILABLE = False

# In-memory user storage (in production, use PostgreSQL)
users_db = {}

def get_password_hash(password):
    return pwd_context.hash(password)

def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

async def get_user(email: str) -> Optional[UserInDB]:
    if email in users_db:
        user_data = users_db[email]
        return UserInDB(**user_data)
    return None

async def authenticate_user(email: str, password: str):
    user = await get_user(email)
    if not user or not verify_password(password, user.hashed_password):
        return False
    return user

async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)]) -> UserInDB:
    credentials_exception = HTTPException(
        status_code=401,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            raise credentials_exception
        token_data = TokenData(email=email)
    except JWTError:
        raise credentials_exception
    user = await get_user(email=token_data.email)
    if user is None:
        raise credentials_exception
    return user

# HTTP Bearer scheme for optional authentication
bearer_scheme = HTTPBearer(auto_error=False)

async def get_current_user_optional(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> Optional[UserInDB]:
    if credentials is None:
        return None
    try:
        token = credentials.credentials
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            return None
        token_data = TokenData(email=email)
        user = await get_user(email=token_data.email)
        return user
    except JWTError:
        return None

# Create FastAPI app
app = FastAPI(
    title="RAG Agent API",
    description="API for RAG Agent with document retrieval and question answering",
    version="1.0.0"
)

# Authentication endpoints
@app.post("/auth/register", response_model=Token)
async def register(user: UserCreate):
    # Check if user already exists
    existing_user = await get_user(user.email)
    if existing_user:
        raise HTTPException(status_code=400, detail="Email already registered")

    # Hash the password
    hashed_password = get_password_hash(user.password)

    # Create user object
    user_id = str(uuid.uuid4())
    now = datetime.utcnow()
    user_in_db = UserInDB(
        id=user_id,
        email=user.email,
        name=user.name,
        hashed_password=hashed_password,
        created_at=now,
        updated_at=now
    )

    # Store user in database (in-memory for now)
    users_db[user.email] = user_in_db.model_dump()

    # Create and return access token
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.email}, expires_delta=access_token_expires
    )

    return {"access_token": access_token, "token_type": "bearer"}

from fastapi.security import OAuth2PasswordRequestForm

@app.post("/auth/token", response_model=Token)
async def login_for_access_token(form_data: Annotated[OAuth2PasswordRequestForm, Depends()]):
    user = await authenticate_user(form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=401,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Bearer"},
        )

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.email}, expires_delta=access_token_expires
    )

    return {"access_token": access_token, "token_type": "bearer"}

@app.get("/auth/me", response_model=UserWithProfile)
async def read_users_me(current_user: UserInDB = Depends(get_current_user)):
    # Convert UserInDB to UserWithProfile to include profile data
    user_data = current_user.model_dump()

    # Get profile data from user storage (in a real app, this would come from a separate profile table)
    user_db_entry = users_db.get(current_user.email, {})
    profile_data = {
        "hardware_type": user_db_entry.get("hardware_type"),
        "jetson_model": user_db_entry.get("jetson_model"),
        "components": user_db_entry.get("components", []),
        "software_stack": user_db_entry.get("software_stack", []),
        "experience_level": user_db_entry.get("experience_level")
    }

    # Merge user data with profile data
    user_with_profile = UserWithProfile(**{**user_data, **profile_data})
    return user_with_profile

@app.put("/auth/profile", response_model=UserWithProfile)
async def update_user_profile(
    profile_update: UserUpdate,
    current_user: UserInDB = Depends(get_current_user)
):
    # Update user profile in the database
    if current_user.email in users_db:
        user_data = users_db[current_user.email]
        # Update profile fields
        for field, value in profile_update.model_dump(exclude_unset=True).items():
            if value is not None:
                user_data[field] = value
        user_data["updated_at"] = datetime.utcnow()
        users_db[current_user.email] = user_data

        # Return updated user with profile
        user_with_profile = UserWithProfile(**user_data)
        return user_with_profile

    raise HTTPException(status_code=404, detail="User not found")

# Add CORS middleware for development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://my-book-phi-nine.vercel.app/"],  # In production, replace with specific origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models
class QueryRequest(BaseModel):
    query: str

class MatchedChunk(BaseModel):
    content: str
    url: str
    position: int
    similarity_score: float

class QueryResponse(BaseModel):
    answer: str
    sources: List[str]
    matched_chunks: List[MatchedChunk]
    error: Optional[str] = None
    status: str  # "success", "error", "empty"
    query_time_ms: Optional[float] = None
    confidence: Optional[str] = None

class HealthResponse(BaseModel):
    status: str
    message: str

# Global RAG agent instance
rag_agent = None

@app.on_event("startup")
async def startup_event():
    global rag_agent
    logger.info("Initializing RAG Agent...")
    try:
        if MYAGENT_AVAILABLE and Myagent:
            rag_agent = Myagent   # ← BAS YE – bina () ke
            logger.info("RAG Agent initialized successfully")
        else:
            rag_agent = None
            logger.warning("RAG Agent not available - chat functionality will be limited")
    except Exception as e:
        logger.error(f"Failed to initialize RAG Agent: {e}")
        # Don't raise the exception to allow the server to start even if RAG agent fails
        rag_agent = None

@app.post("/ask", response_model=QueryResponse)
async def ask_rag(
    request: QueryRequest,
    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
):
    # Try to get user from token if provided
    current_user = await get_current_user_optional(credentials)
    logger.info(f"Processing query: {request.query[:50]}...")

    try:
        if not request.query or len(request.query.strip()) == 0:
            raise HTTPException(status_code=400, detail="Query cannot be empty")

        if len(request.query) > 2000:
            raise HTTPException(status_code=400, detail="Query too long, maximum 2000 characters")

        # Check if agents module is available
        if not AGENTS_AVAILABLE or 'rag_agent' not in globals():
            logger.warning("RAG agent not available, returning mock response")
            return QueryResponse(
                answer="Mock response: The RAG agent is not available. Please ensure all dependencies are installed and the agent is properly configured.",
                sources=[],
                matched_chunks=[],
                error=None,
                status="success",
                query_time_ms=None,
                confidence=None
            )

        # Sahi tareeka: Runner.run() use karo
        result = await Runner.run(
            rag_agent,  # ye tumhara Myagent object hai
            input=request.query
        )

        # Result se data nikalo
        answer = result.final_output or ""

        # Agar tool use hua ho aur chunks mile hon to sources/matched_chunks banao
        # (ye optional hai – agar library chunks return karti hai)
        matched_chunks = []
        sources = []

        # Agar result mein history ya tool calls hain to parse kar sakte ho
        # Simple version ke liye sirf answer return karo pehle

        formatted_response = QueryResponse(
            answer=answer,
            sources=sources,
            matched_chunks=matched_chunks,
            error=None,
            status="success",
            query_time_ms=None,
            confidence=None
        )

        logger.info("Query processed successfully")
        return formatted_response

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        return QueryResponse(
            answer="",
            sources=[],
            matched_chunks=[],
            error=str(e),
            status="error"
        )


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """
    Health check endpoint
    """
    return HealthResponse(
        status="healthy",
        message="RAG Agent API is running"
    )

# Protected version of the ask endpoint that requires authentication
@app.post("/ask-protected", response_model=QueryResponse)
async def ask_rag_protected(
    request: QueryRequest,
    current_user: UserInDB = Depends(get_current_user)
):
    logger.info(f"Processing authenticated query: {request.query[:50]}...")

    try:
        if not request.query or len(request.query.strip()) == 0:
            raise HTTPException(status_code=400, detail="Query cannot be empty")

        if len(request.query) > 2000:
            raise HTTPException(status_code=400, detail="Query too long, maximum 2000 characters")

        # Check if agents module is available
        if not AGENTS_AVAILABLE or not rag_agent:
            logger.warning("RAG agent not available, returning mock response for protected endpoint")
            return QueryResponse(
                answer="Mock response: The RAG agent is not available. Please ensure all dependencies are installed and the agent is properly configured.",
                sources=[],
                matched_chunks=[],
                error=None,
                status="success",
                query_time_ms=None,
                confidence=None
            )

        # Sahi tareeka: Runner.run() use karo
        result = await Runner.run(
            rag_agent,  # ye tumhara Myagent object hai
            input=request.query
        )

        # Result se data nikalo
        answer = result.final_output or ""

        # Agar tool use hua ho aur chunks mile hon to sources/matched_chunks banao
        # (ye optional hai – agar library chunks return karti hai)
        matched_chunks = []
        sources = []

        # Agar result mein history ya tool calls hain to parse kar sakte ho
        # Simple version ke liye sirf answer return karo pehle

        formatted_response = QueryResponse(
            answer=answer,
            sources=sources,
            matched_chunks=matched_chunks,
            error=None,
            status="success",
            query_time_ms=None,
            confidence=None
        )

        logger.info("Authenticated query processed successfully")
        return formatted_response

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing authenticated query: {e}")
        return QueryResponse(
            answer="",
            sources=[],
            matched_chunks=[],
            error=str(e),
            status="error"
        )

# For running with uvicorn
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)