import logging
from datetime import datetime
from typing import Optional, List, Dict, Any
from sqlmodel import Session, select
from sqlalchemy.orm import selectinload
from sqlalchemy import func
from src.config.database import engine
from src.services.storage_service import storage_manager
from src.utils.oss_utils import get_oss_key
from src.models.entities import ProjectDB, AssetDB, EpisodeDB, StoryboardDB
from src.models.schemas import ProjectData
from src.mappers import ProjectMapper, AssetMapper, EpisodeMapper, StoryboardMapper
from src.repositories.base_repository import BaseRepository

logger = logging.getLogger(__name__)

class ProjectRepository(BaseRepository[ProjectDB]):
    def __init__(self, session: Optional[Session] = None):
        """ Initialize ProjectRepository.
        
        Args:
            session: Optional database session. If not provided, creates sessions per operation.
        """
        self._external_session = session
        if session:
            super().__init__(ProjectDB, session)
    
    def _get_session(self):
        """ Get session context manager."""
        if self._external_session:
            # Use external session without context manager
            from contextlib import nullcontext
            return nullcontext(self._external_session)
        else:
            # Create new session with context manager
            return Session(engine)

    def _sign_url(self, url: Optional[str]) -> Optional[str]:
        if not url:
            return url
        return storage_manager.sign_url(url)

    def _sign_urls(self, urls: Optional[List[str]]) -> List[str]:
        if not urls:
            return []
        return [self._sign_url(url) for url in urls if url]

    def _process_generations(self, gens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        if not gens:
            return []
        processed = []
        for g in gens:
            new_g = g.copy()
            
            # Handle result_url / resultUrl
            if new_g.get("result_url"):
                new_g["result_url"] = self._sign_url(new_g["result_url"])
            if new_g.get("resultUrl"):
                new_g["resultUrl"] = self._sign_url(new_g["resultUrl"])
            
            # Handle result_urls / resultUrls
            if new_g.get("result_urls"):
                new_g["result_urls"] = self._sign_urls(new_g["result_urls"])
            if new_g.get("resultUrls"):
                new_g["resultUrls"] = self._sign_urls(new_g["resultUrls"])
                
            processed.append(new_g)
        return processed

    def _strip_url(self, url: Optional[str]) -> Optional[str]:
        if not url:
            return url
        return get_oss_key(url)

    def _strip_urls(self, urls: Optional[List[str]]) -> List[str]:
        if not urls:
            return []
        return [self._strip_url(url) for url in urls if url]

    def _process_generations_strip(self, gens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        if not gens:
            return []
        processed = []
        for g in gens:
            new_g = g.copy()
            
            # Handle result_url / resultUrl
            if new_g.get("result_url"):
                new_g["result_url"] = self._strip_url(new_g["result_url"])
            if new_g.get("resultUrl"):
                new_g["resultUrl"] = self._strip_url(new_g["resultUrl"])
            
            # Handle result_urls / resultUrls
            if new_g.get("result_urls"):
                new_g["result_urls"] = self._strip_urls(new_g["result_urls"])
            if new_g.get("resultUrls"):
                new_g["resultUrls"] = self._strip_urls(new_g["resultUrls"])
                
            processed.append(new_g)
        return processed

    def _map_to_domain(self, project_db: ProjectDB, include_relations: bool = True, include_assets: bool = True, assets_override: List[AssetDB] = None) -> ProjectData:
        # Use ProjectMapper for base conversion
        project_data = ProjectMapper.to_schema(project_db, include_relations=False)
        
        assets = []
        
        # Determine source of assets: override list takes precedence, then DB relationship if requested
        source_assets = []
        if assets_override is not None:
            source_assets = assets_override
        elif include_relations and include_assets:
            source_assets = project_db.assets
            
        if source_assets:
            for adb in source_assets:
                try:
                    # Use AssetMapper to convert entity to schema
                    asset = AssetMapper.to_schema(adb)
                    
                    # Sign URLs for the asset
                    if asset.image_url:
                        asset.image_url = self._sign_url(asset.image_url)
                    if asset.image_urls:
                        asset.image_urls = self._sign_urls(asset.image_urls)
                    if asset.video_urls:
                        asset.video_urls = self._sign_urls(asset.video_urls)
                    
                    # 进程 generations
                    if asset.generations:
                        asset.generations = self._process_generations([g.model_dump() for g in asset.generations])
                    
                    assets.append(asset)
                except Exception as e:
                    logger.error(f"Error parsing asset {adb.id}: {e}")

        # Episodes
        episodes = []
        if include_relations:
            for edb in sorted(project_db.episodes, key=lambda x: x.order_index):
                episodes.append(EpisodeMapper.to_schema(edb))

        # Storyboards
        storyboards = []
        if include_relations:
            for sdb in sorted(project_db.storyboards, key=lambda x: x.order_index):
                storyboard = StoryboardMapper.to_schema(sdb)
                
                # Sign URLs for the storyboard
                if storyboard.audio_url:
                    storyboard.audio_url = self._sign_url(storyboard.audio_url)
                if storyboard.image_urls:
                    storyboard.image_urls = self._sign_urls(storyboard.image_urls)
                if storyboard.video_urls:
                    storyboard.video_urls = self._sign_urls(storyboard.video_urls)
                
                # 进程 generations
                if storyboard.generations:
                    storyboard.generations = self._process_generations([g.model_dump() for g in storyboard.generations])
                
                storyboards.append(storyboard)

        # 更新 project_data with relations
        project_data.assets = assets
        project_data.episodes = episodes
        project_data.storyboards = storyboards
        
        return project_data

    def get(self, project_id: str, include_assets: bool = True, include_referenced_assets: bool = False, user_id: Optional[str] = None) -> Optional[ProjectData]:
        with self._get_session() as session:
            # Use select with options to control eager loading
            statement = select(ProjectDB).where(ProjectDB.id == project_id)

            # Always load these small relations
            options = [
                selectinload(ProjectDB.episodes),
                selectinload(ProjectDB.storyboards),
                selectinload(ProjectDB.canvas_metadata)
            ]

            if include_assets:
                options.append(selectinload(ProjectDB.assets))

            statement = statement.options(*options)

            project_db = session.exec(statement).first()

            if not project_db or project_db.deleted_at is not None:
                return None

            # Check ownership if user_id is provided
            if user_id is not None and project_db.user_id != user_id:
                return None

            # Handle referenced assets logic
            assets_override = None
            if not include_assets and include_referenced_assets:
                referenced_ids = set()
                # Check storyboards for referenced assets
                if project_db.storyboards:
                    for sb in project_db.storyboards:
                        if sb.scene_id: referenced_ids.add(sb.scene_id)
                        if sb.character_ids: referenced_ids.update(sb.character_ids)
                        if sb.prop_ids: referenced_ids.update(sb.prop_ids)

                # If we found references, load them
                if referenced_ids:
                    assets_override = session.exec(select(AssetDB).where(AssetDB.id.in_(referenced_ids))).all()
                else:
                    assets_override = []

            return self._map_to_domain(
                project_db,
                include_relations=True,
                include_assets=include_assets,
                assets_override=assets_override
            )

    def list(self, limit: int = 50, offset: int = 0, include_relations: bool = False, user_id: Optional[str] = None) -> List[ProjectData]:
        with self._get_session() as session:
            statement = select(ProjectDB).where(ProjectDB.deleted_at.is_(None))

            # Filter by user_id if provided
            if user_id:
                statement = statement.where(ProjectDB.user_id == user_id)

            if include_relations:
                statement = statement.options(
                    selectinload(ProjectDB.assets),
                    selectinload(ProjectDB.episodes),
                    selectinload(ProjectDB.storyboards),
                    selectinload(ProjectDB.canvas_metadata)
                )

            statement = statement.order_by(ProjectDB.updated_at.desc()).offset(offset).limit(limit)
            projects_db = session.exec(statement).all()
            return [self._map_to_domain(pdb, include_relations=include_relations) for pdb in projects_db]

    def list_by_user(self, user_id: str, limit: int = 50, offset: int = 0, include_relations: bool = False) -> List[ProjectData]:
        """List projects by user ID."""
        return self.list(limit=limit, offset=offset, include_relations=include_relations, user_id=user_id)

    def count(self, user_id: Optional[str] = None) -> int:
        with self._get_session() as session:
            # Use func.count for efficient counting
            from sqlmodel import func
            statement = select(func.count()).select_from(ProjectDB).where(ProjectDB.deleted_at.is_(None))

            # Filter by user_id if provided
            if user_id:
                statement = statement.where(ProjectDB.user_id == user_id)

            return session.exec(statement).one()

    def count_by_user(self, user_id: str) -> int:
        """Count projects by user ID."""
        return self.count(user_id=user_id)
        with self._get_session() as session:
            # Use func.count for efficient counting
            from sqlmodel import func
            statement = select(func.count()).select_from(ProjectDB).where(ProjectDB.deleted_at.is_(None))
            return session.exec(statement).one()

    def reset_stuck_status(self) -> int:
        """ 重置s projects stuck in 'initializing' or 'processing' state to 'failed'.
        Returns the number of projects reset.
        """
        with self._get_session() as session:
            statement = select(ProjectDB).where(
                ProjectDB.status.in_(['initializing', 'processing']),
                ProjectDB.deleted_at.is_(None)
            )
            projects = session.exec(statement).all()
            count = 0
            for project in projects:
                project.status = 'failed'
                # Append a note to description if not already present
                note = "\n[System] Status reset to failed due to system restart."
                if not project.description or note.strip() not in project.description:
                     project.description = (project.description or "") + note
                
                session.add(project)
                count += 1
            
            if count > 0:
                session.commit()
                logger.info(f"Reset {count} stuck projects to 'failed' status")
            
            return count

    def search_projects(self, query: str, limit: int = 50) -> List[ProjectData]:
        """ 搜索 projects by name or description using full-text search.
        For SQLite, we use LIKE queries with indexes for performance.
        Returns results ordered by relevance (name matches first, then description).
        """
        with self._get_session() as session:
            # 搜索 in name (higher priority) and description
            # Use LIKE with wildcards for flexible matching
            search_pattern = f"%{query}%"
            
            # First get projects where name matches
            name_matches = select(ProjectDB).where(
                ProjectDB.deleted_at.is_(None),
                ProjectDB.name.like(search_pattern)
            ).order_by(ProjectDB.updated_at.desc())
            
            # Then get projects where description matches (but name doesn't)
            desc_matches = select(ProjectDB).where(
                ProjectDB.deleted_at.is_(None),
                ProjectDB.name.notlike(search_pattern),
                ProjectDB.description.like(search_pattern)
            ).order_by(ProjectDB.updated_at.desc())
            
            # Combine results with name matches first (higher relevance)
            name_results = session.exec(name_matches).all()
            desc_results = session.exec(desc_matches).all()
            
            # Combine and limit results
            all_results = list(name_results) + list(desc_results)
            limited_results = all_results[:limit]
            
            return [self._map_to_domain(pdb) for pdb in limited_results]

    def _map_asset_to_db(self, project_id: str, asset: Any) -> AssetDB:
        """ 辅助函数 to map domain Asset to AssetDB"""
        asset_data = asset.model_dump(by_alias=False)
        common_fields = {
            "id", "type", "name", "desc", "tags",
            "image_url", "image_urls", "video_urls", "generations",
        }
        extra_data = {k: v for k, v in asset_data.items() if k not in common_fields}
        
        return AssetDB(
            id=asset.id,
            project_id=project_id,
            type=asset.type,
            name=asset.name,
            desc=asset.desc,
            tags=asset.tags,
            image_url=self._strip_url(asset.image_url),
            image_urls=self._strip_urls(asset.image_urls),
            video_urls=self._strip_urls(asset.video_urls),
            generations=self._process_generations_strip(asset_data.get('generations') or []),
            extra_data=extra_data,
        )

    def add_asset(self, project_id: str, asset: Any):
        """ 原子操作ally add an asset to a project"""
        with self._get_session() as session:
            asset_db = self._map_asset_to_db(project_id, asset)
            session.add(asset_db)
            
            # 更新 project timestamp
            project = session.get(ProjectDB, project_id)
            if project:
                project.updated_at = datetime.now().timestamp()
                session.add(project)
                
            session.commit()
            
    def batch_add_assets(self, project_id: str, assets: List[Any]) -> Optional[ProjectData]:
        """ 原子操作ally add multiple assets to a project"""
        with self._get_session() as session:
            project = session.get(ProjectDB, project_id)
            if not project:
                return None
            
            for asset in assets:
                asset_db = self._map_asset_to_db(project_id, asset)
                session.add(asset_db)
            
            project.updated_at = datetime.now().timestamp()
            session.add(project)
            session.commit()
            session.refresh(project)
            return self._map_to_domain(project)

    def update_asset(self, project_id: str, asset: Any):
        """ 原子操作ally update an asset"""
        with self._get_session() as session:
            asset_db = self._map_asset_to_db(project_id, asset)
            session.merge(asset_db)
            
            # 更新 project timestamp
            project = session.get(ProjectDB, project_id)
            if project:
                project.updated_at = datetime.now().timestamp()
                session.add(project)
                
            session.commit()

    def delete_asset(self, project_id: str, asset_id: str) -> Optional[ProjectData]:
        """ 原子操作ally delete an asset from a project and remove references from storyboards.
        """
        with self._get_session() as session:
            # Check if asset exists and belongs to project
            asset_db = session.exec(
                select(AssetDB)
                .where(AssetDB.id == asset_id)
                .where(AssetDB.project_id == project_id)
            ).first()
            
            if not asset_db:
                logger.warning(f"Asset {asset_id} not found in project {project_id}")
                return None
            
            # 清理 up references in Storyboards
            # Get all storyboards for this project
            storyboards = session.exec(
                select(StoryboardDB).where(StoryboardDB.project_id == project_id)
            ).all()
            
            for sb in storyboards:
                updated = False
                
                # Check Scene reference
                if sb.scene_id == asset_id:
                    sb.scene_id = None
                    updated = True
                
                # Check Character references
                if sb.character_ids and asset_id in sb.character_ids:
                    # Create new list to ensure SQLAlchemy detects change
                    sb.character_ids = [id for id in sb.character_ids if id != asset_id]
                    updated = True
                    
                # Check Prop references
                if sb.prop_ids and asset_id in sb.prop_ids:
                    # Create new list to ensure SQLAlchemy detects change
                    sb.prop_ids = [id for id in sb.prop_ids if id != asset_id]
                    updated = True
                
                if updated:
                    session.add(sb)
            
            # 删除 asset
            session.delete(asset_db)
            
            # 更新 project timestamp
            project = session.get(ProjectDB, project_id)
            if project:
                project.updated_at = datetime.now().timestamp()
                session.add(project)
                
            session.commit()
            session.refresh(project)
            return self._map_to_domain(project)

    def list_assets(self, project_id: str, asset_type: Optional[str] = None, search_query: Optional[str] = None, limit: int = 50, offset: int = 0) -> Optional[Dict[str, Any]]:
        """ List assets for a project with pagination and optional search"""
        with self._get_session() as session:
            # Check project exists
            project = session.get(ProjectDB, project_id)
            if not project:
                return None
            
            # 构建 query
            query = select(AssetDB).where(AssetDB.project_id == project_id)
            if asset_type:
                query = query.where(AssetDB.type == asset_type)
            
            if search_query:
                from sqlalchemy import or_
                search_pattern = f"%{search_query}%"
                query = query.where(
                    or_(
                        AssetDB.name.like(search_pattern),
                        AssetDB.desc.like(search_pattern),
                        AssetDB.tags.cast(str).like(search_pattern) # 简单 tag search
                    )
                )
            
            # 计数
            count_query = select(func.count()).select_from(AssetDB).where(AssetDB.project_id == project_id)
            if asset_type:
                count_query = count_query.where(AssetDB.type == asset_type)
            if search_query:
                # Need to replicate search logic for count
                from sqlalchemy import or_
                search_pattern = f"%{search_query}%"
                count_query = count_query.where(
                    or_(
                        AssetDB.name.like(search_pattern),
                        AssetDB.desc.like(search_pattern),
                        AssetDB.tags.cast(str).like(search_pattern)
                    )
                )

            total = session.exec(count_query).one()
            
            # Pagination
            query = query.offset(offset).limit(limit)
            assets_db = session.exec(query).all()
            
            # 映射
            assets = []
            for adb in assets_db:
                data = {
                    "id": adb.id,
                    "type": adb.type,
                    "name": adb.name,
                    "desc": adb.desc,
                    "tags": adb.tags,
                    "imageUrl": self._sign_url(adb.image_url),
                    "imageUrls": self._sign_urls(adb.image_urls),
                    "videoUrls": self._sign_urls(adb.video_urls),
                    "generations": self._process_generations(adb.generations),
                    **adb.extra_data,
                }
                try:
                    if adb.type == "character":
                        assets.append(CharacterAsset(**data))
                    elif adb.type == "scene":
                        assets.append(SceneAsset(**data))
                    elif adb.type == "prop":
                        assets.append(PropAsset(**data))
                    else:
                        assets.append(OtherAsset(**data))
                except Exception as e:
                    logger.warning(f"Failed to map asset {adb.id}: {e}")
            
            return {
                "assets": assets,
                "total": total,
                "limit": limit,
                "offset": offset
            }

    def save(self, project: ProjectData):
        """
        Saves the project data to SQLite using SQLModel.
        Uses a transaction to ensure atomicity.
        Replaces all sub-items (assets, episodes, storyboards) to ensure consistency.
        """
        with self._get_session() as session:
            # 1. Upsert Project
            project_db = session.get(ProjectDB, project.id)
            
            # 序列化 progress to dict if present
            progress_dict = None
            if project.progress:
                progress_dict = project.progress.model_dump(by_alias=False)
            
            if not project_db:
                project_db = ProjectDB(
                    id=project.id,
                    name=project.name,
                    description=project.description,
                    type=project.type,
                    status=project.status,
                    created_at=project.created_at.timestamp(),
                    updated_at=project.updated_at.timestamp(),
                    resolution=project.resolution,
                    ratio=project.ratio,
                    style_id=project.style_id,
                    style_params=project.style_params,
                    chapters=project.chapters,
                    # general_canvases removed - now using canvas_metadata table
                    progress=progress_dict,
                    error=project.error,
                    user_id=project.user_id,
                )
                session.add(project_db)
            else:
                project_db.name = project.name
                project_db.description = project.description
                project_db.type = project.type
                project_db.status = project.status
                project_db.updated_at = project.updated_at.timestamp()
                project_db.resolution = project.resolution
                project_db.ratio = project.ratio
                project_db.style_id = project.style_id
                project_db.style_params = project.style_params
                project_db.chapters = project.chapters
                # general_canvases removed - now using canvas_metadata table
                project_db.progress = progress_dict
                project_db.error = project.error
                # Only update user_id if it's not already set (preserve original owner)
                if not project_db.user_id and project.user_id:
                    project_db.user_id = project.user_id

            # 2. Update Assets
            # We recreate the list to ensure full sync
            new_assets = []
            for asset in project.assets:
                # 映射 Pydantic Asset to AssetDB
                asset_data = asset.model_dump(by_alias=False)
                # 字段s mapped to columns
                common_fields = {
                    "id",
                    "type",
                    "name",
                    "desc",
                    "tags",
                    "image_url",
                    "image_urls",
                    "video_urls",
                    "generations",
                }

                # 提取 extra data
                extra_data = {
                    k: v for k, v in asset_data.items() if k not in common_fields
                }

                asset_db = AssetDB(
                    id=asset.id,
                    project_id=project.id,
                    type=asset.type,
                    name=asset.name,
                    desc=asset.desc,
                    tags=asset.tags,
                    image_url=self._strip_url(asset.image_url),
                    image_urls=self._strip_urls(asset.image_urls),
                    video_urls=self._strip_urls(asset.video_urls),
                    generations=self._process_generations_strip(asset_data.get('generations') or []),
                    extra_data=extra_data,
                )
                new_assets.append(asset_db)

            project_db.assets = new_assets

            # 3. Update Episodes
            new_episodes = []
            for ep in project.episodes:
                ep_db = EpisodeDB(
                    id=ep.id,
                    project_id=project.id,
                    order_index=ep.order,
                    title=ep.title,
                    desc=ep.desc,
                    content=ep.content,
                    status=ep.status,
                )
                new_episodes.append(ep_db)

            project_db.episodes = new_episodes

            # 4. Update Storyboards
            new_storyboards = []
            for sb in project.storyboards:
                sb_data = sb.model_dump(by_alias=False)
                sb_db = StoryboardDB(
                    id=sb.id,
                    project_id=project.id,
                    episode_id=sb.episode_id,
                    order_index=sb.order,
                    shot=sb.shot,
                    desc=sb.desc,
                    duration=sb.duration,
                    type=sb.type,
                    scene_id=sb.scene_id,
                    character_ids=sb.character_ids,
                    prop_ids=sb.prop_ids,
                    voiceover=sb.voiceover,
                    audio_desc=sb.audio_desc,
                    audio_url=self._strip_url(sb.audio_url),
                    camera_movement=sb.camera_movement,
                    transition=sb.transition,
                    camera_angle=sb.camera_angle,
                    lens=sb.lens,
                    focus=sb.focus,
                    lighting=sb.lighting,
                    color_style=sb.color_style,
                    location=sb.location,
                    time=sb.time,
                    original_text=sb.original_text,  # ✅ 添加
                    merge_image_prompt=sb.merge_image_prompt,  # ✅ 添加
                    video_prompt=sb.video_prompt,  # ✅ 添加
                    image_urls=self._strip_urls(sb.image_urls),
                    video_urls=self._strip_urls(sb.video_urls),
                    generations=self._process_generations_strip(sb_data.get('generations') or []),
                )
                new_storyboards.append(sb_db)

            project_db.storyboards = new_storyboards

            session.commit()

    def delete(self, project_id: str):
        """
        Soft delete a project and all its related resources (Cascading Soft Delete).
        Sets deleted_at timestamp for Project, Assets, Episodes, Storyboards, and CanvasMetadata.
        """
        with self._get_session() as session:
            # Load project with all relations to ensure we can iterate and update them
            statement = select(ProjectDB).where(ProjectDB.id == project_id).options(
                selectinload(ProjectDB.assets),
                selectinload(ProjectDB.episodes),
                selectinload(ProjectDB.storyboards),
                selectinload(ProjectDB.canvas_metadata)
            )
            project_db = session.exec(statement).first()
            
            if project_db:
                now = datetime.now().timestamp()
                
                # 1. Delete Project
                project_db.deleted_at = now
                session.add(project_db)
                
                # 2. Delete Assets
                for asset in project_db.assets:
                    asset.deleted_at = now
                    session.add(asset)
                    
                # 3. Delete Episodes
                for episode in project_db.episodes:
                    episode.deleted_at = now
                    session.add(episode)
                    
                # 4. Delete Storyboards
                for storyboard in project_db.storyboards:
                    storyboard.deleted_at = now
                    session.add(storyboard)
                    
                # 5. Delete Canvas Metadata
                for meta in project_db.canvas_metadata:
                    meta.deleted_at = now
                    session.add(meta)
                
                session.commit()
    
    def hard_delete(self, project_id: str):
        """ 永久的ly delete a project from the database.
        Use with caution - this cannot be undone.
        """
        with self._get_session() as session:
            project_db = session.get(ProjectDB, project_id)
            if project_db:
                session.delete(project_db)
                session.commit()

