"""
Work Scheduler for Allegro-Primus
Intelligent task scheduling based on Gitea issues and priorities.
"""

import os
import json
import logging
import asyncio
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from enum import Enum
import time

from gitea_client import GiteaClient, Issue, IssueState, IssuePriority
from repo_manager import RepoManager

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


class TaskStatus(str, Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    BLOCKED = "blocked"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class WorkTask:
    """Represents a scheduled work task."""
    issue: Issue
    owner: str
    repo: str
    status: TaskStatus
    priority_score: float
    estimated_hours: Optional[float] = None
    assigned_to: Optional[str] = None
    branch_name: Optional[str] = None
    pr_number: Optional[int] = None
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    notes: List[str] = None
    
    def __post_init__(self):
        if self.notes is None:
            self.notes = []
    
    def to_dict(self) -> Dict[str, Any]:
        data = asdict(self)
        data['issue'] = asdict(self.issue)
        return data


@dataclass
class WorkCycle:
    """Represents a scheduled work cycle."""
    cycle_id: str
    started_at: str
    tasks: List[WorkTask]
    status: str
    completed_tasks: int = 0


class WorkScheduler:
    """
    Intelligent work scheduler for autonomous issue resolution.
    Queries Gitea, prioritizes tasks, and manages work cycles.
    """

    # Priority weights for scoring
    PRIORITY_WEIGHTS = {
        "P0": 1000,  # Critical
        "P1": 100,   # High
        "P2": 10,    # Medium
        "P3": 1      # Low
    }

    def __init__(
        self,
        gitea_client: Optional[GiteaClient] = None,
        repo_manager: Optional[RepoManager] = None,
        work_dir: str = "./work",
        max_concurrent: int = 3
    ):
        """
        Initialize WorkScheduler.
        
        Args:
            gitea_client: Gitea API client
            repo_manager: Repository manager for git ops
            work_dir: Directory for work state
            max_concurrent: Maximum concurrent tasks
        """
        self.gitea = gitea_client or GiteaClient()
        self.repo_manager = repo_manager or RepoManager()
        self.work_dir = work_dir
        self.max_concurrent = max_concurrent
        self.current_cycle: Optional[WorkCycle] = None
        self.task_handlers: Dict[str, Callable] = {}
        
        os.makedirs(work_dir, exist_ok=True)
        logger.info(f"WorkScheduler initialized (max_concurrent={max_concurrent})")

    def register_task_handler(
        self,
        task_type: str,
        handler: Callable[[WorkTask], Any]
    ) -> None:
        """
        Register a handler function for specific task types.
        
        Args:
            task_type: Type identifier (e.g., "bug", "feature", "refactor")
            handler: Function to handle tasks of this type
        """
        self.task_handlers[task_type] = handler
        logger.info(f"Registered handler for task type: {task_type}")

    def _calculate_priority_score(self, issue: Issue) -> float:
        """
        Calculate priority score for an issue.
        
        Higher score = higher priority.
        Considers:
        - Priority label (P0, P1, P2, P3)
        - Age of issue
        - Assignee status
        - Comment activity
        """
        score = 0.0
        
        # Base priority from label
        priority = issue.priority or "P3"
        score += self.PRIORITY_WEIGHTS.get(priority, 1)
        
        # Age factor: older issues get slight boost
        try:
            created = datetime.fromisoformat(issue.created_at.replace('Z', '+00:00'))
            age_days = (datetime.now(created.tzinfo) - created).days
            score += min(age_days * 0.1, 10)  # Max 10 points for age
        except:
            pass
        
        # Unassigned issues get priority
        if not issue.assignee:
            score += 5
        
        return score

    def fetch_priority_issues(
        self,
        owner: str,
        repo: str,
        priorities: List[str] = None
    ) -> List[Issue]:
        """
        Fetch high-priority open issues from Gitea.
        
        Args:
            owner: Repository owner
            repo: Repository name
            priorities: Priority labels to fetch (default: P0, P1)
            
        Returns:
            List of priority issues sorted by score
        """
        if priorities is None:
            priorities = ["P0", "P1"]
        
        logger.info(f"Fetching issues with priorities: {priorities}")
        
        try:
            issues = self.gitea.get_priority_issues(
                owner, repo,
                priorities=priorities,
                state=IssueState.OPEN
            )
            logger.info(f"Found {len(issues)} priority issues")
            return issues
        except Exception as e:
            logger.error(f"Failed to fetch issues: {e}")
            return []

    def prioritize_tasks(
        self,
        issues: List[Issue],
        owner: str,
        repo: str
    ) -> List[WorkTask]:
        """
        Convert issues to prioritized work tasks.
        
        Args:
            issues: List of issues from Gitea
            owner: Repository owner
            repo: Repository name
            
        Returns:
            List of WorkTask sorted by priority score
        """
        tasks = []
        
        for issue in issues:
            score = self._calculate_priority_score(issue)
            task = WorkTask(
                issue=issue,
                owner=owner,
                repo=repo,
                status=TaskStatus.PENDING,
                priority_score=score
            )
            tasks.append(task)
        
        # Sort by priority score (descending)
        tasks.sort(key=lambda t: t.priority_score, reverse=True)
        
        logger.info(f"Prioritized {len(tasks)} tasks")
        return tasks

    def create_work_cycle(
        self,
        owner: str,
        repo: str,
        cycle_name: Optional[str] = None
    ) -> WorkCycle:
        """
        Create a new work cycle with prioritized tasks.
        
        Args:
            owner: Repository owner
            repo: Repository name
            cycle_name: Optional cycle identifier
            
        Returns:
            WorkCycle with scheduled tasks
        """
        # Fetch and prioritize issues
        issues = self.fetch_priority_issues(owner, repo)
        tasks = self.prioritize_tasks(issues, owner, repo)
        
        # Create cycle
        cycle_id = cycle_name or f"cycle-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
        cycle = WorkCycle(
            cycle_id=cycle_id,
            started_at=datetime.now().isoformat(),
            tasks=tasks,
            status="created"
        )
        
        self.current_cycle = cycle
        self._save_cycle_state(cycle)
        
        logger.info(f"Created work cycle '{cycle_id}' with {len(tasks)} tasks")
        return cycle

    def start_task(self, task: WorkTask) -> bool:
        """
        Mark a task as in-progress and update Gitea.
        
        Args:
            task: WorkTask to start
            
        Returns:
            True if started successfully
        """
        try:
            task.status = TaskStatus.IN_PROGRESS
            task.started_at = datetime.now().isoformat()
            
            # Create work branch
            branch = self.repo_manager.create_work_branch(
                task.owner,
                task.repo,
                task.issue.number,
                task.issue.title[:20]
            )
            task.branch_name = branch
            
            # Update issue in Gitea - add comment and assign
            self.gitea.add_issue_comment(
                task.owner,
                task.repo,
                task.issue.number,
                f"🤖 Allegro-Primus is working on this issue.\n"
                f"Branch: `{branch}`\n"
                f"Started: {task.started_at}"
            )
            
            # Try to self-assign if not assigned
            if not task.issue.assignee:
                try:
                    user = self.gitea.get_user()
                    self.gitea.update_issue(
                        task.owner,
                        task.repo,
                        task.issue.number,
                        assignees=[user['login']]
                    )
                except Exception as e:
                    logger.warning(f"Could not self-assign: {e}")
            
            task.notes.append(f"Started work on branch {branch}")
            logger.info(f"Started task for issue #{task.issue.number}")
            return True
            
        except Exception as e:
            task.status = TaskStatus.FAILED
            task.notes.append(f"Failed to start: {e}")
            logger.error(f"Failed to start task: {e}")
            return False

    def complete_task(
        self,
        task: WorkTask,
        success: bool = True,
        pr_number: Optional[int] = None
    ) -> None:
        """
        Mark a task as completed.
        
        Args:
            task: WorkTask to complete
            success: Whether task completed successfully
            pr_number: PR number if created
        """
        task.completed_at = datetime.now().isoformat()
        task.pr_number = pr_number
        
        if success:
            task.status = TaskStatus.COMPLETED
            
            # Update issue
            comment = f"✅ Allegro-Primus has completed work on this issue."
            if pr_number:
                comment += f"\nPull Request: #{pr_number}"
            
            self.gitea.add_issue_comment(
                task.owner,
                task.repo,
                task.issue.number,
                comment
            )
            
            logger.info(f"Completed task for issue #{task.issue.number}")
        else:
            task.status = TaskStatus.FAILED
            self.gitea.add_issue_comment(
                task.owner,
                task.repo,
                task.issue.number,
                f"❌ Task failed. Check logs for details."
            )
            logger.error(f"Task failed for issue #{task.issue.number}")

    def run_work_cycle(
        self,
        owner: str,
        repo: str,
        task_processor: Optional[Callable[[WorkTask], bool]] = None
    ) -> WorkCycle:
        """
        Execute a complete work cycle.
        
        Args:
            owner: Repository owner
            repo: Repository name
            task_processor: Function to process each task
            
        Returns:
            Completed WorkCycle
        """
        cycle = self.create_work_cycle(owner, repo)
        cycle.status = "running"
        
        if not cycle.tasks:
            logger.info("No tasks to process")
            cycle.status = "completed"
            return cycle
        
        # Process tasks up to max_concurrent
        active_tasks = cycle.tasks[:self.max_concurrent]
        
        for task in active_tasks:
            # Start task
            if not self.start_task(task):
                continue
            
            # Process if handler provided
            if task_processor:
                try:
                    success = task_processor(task)
                    self.complete_task(task, success=success)
                except Exception as e:
                    logger.error(f"Task processor error: {e}")
                    self.complete_task(task, success=False)
            
            cycle.completed_tasks += 1
            self._save_cycle_state(cycle)
        
        cycle.status = "completed"
        self._save_cycle_state(cycle)
        
        logger.info(f"Work cycle completed. {cycle.completed_tasks}/{len(cycle.tasks)} tasks done.")
        return cycle

    def _save_cycle_state(self, cycle: WorkCycle) -> None:
        """Save work cycle state to disk."""
        filepath = os.path.join(self.work_dir, f"{cycle.cycle_id}.json")
        with open(filepath, 'w') as f:
            json.dump({
                "cycle_id": cycle.cycle_id,
                "started_at": cycle.started_at,
                "status": cycle.status,
                "completed_tasks": cycle.completed_tasks,
                "tasks": [t.to_dict() for t in cycle.tasks]
            }, f, indent=2)

    def load_cycle(self, cycle_id: str) -> Optional[WorkCycle]:
        """Load a work cycle from disk."""
        filepath = os.path.join(self.work_dir, f"{cycle_id}.json")
        if not os.path.exists(filepath):
            return None
        
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        # Note: Full reconstruction would require Issue objects
        # This is a simplified load for monitoring
        return WorkCycle(
            cycle_id=data["cycle_id"],
            started_at=data["started_at"],
            tasks=[],  # Would need to reconstruct from data
            status=data["status"],
            completed_tasks=data.get("completed_tasks", 0)
        )

    def get_cycle_history(self) -> List[Dict[str, Any]]:
        """Get list of all saved work cycles."""
        cycles = []
        for filename in os.listdir(self.work_dir):
            if filename.endswith('.json'):
                cycle_id = filename[:-5]
                cycle = self.load_cycle(cycle_id)
                if cycle:
                    cycles.append({
                        "cycle_id": cycle.cycle_id,
                        "started_at": cycle.started_at,
                        "status": cycle.status,
                        "completed_tasks": cycle.completed_tasks
                    })
        return sorted(cycles, key=lambda x: x["started_at"], reverse=True)

    def get_pending_tasks(self, owner: str, repo: str) -> List[WorkTask]:
        """Get list of pending high-priority tasks."""
        issues = self.fetch_priority_issues(owner, repo)
        return self.prioritize_tasks(issues, owner, repo)

    def auto_schedule(
        self,
        owner: str,
        repo: str,
        interval_minutes: int = 60
    ):
        """
        Run continuous scheduling loop.
        
        Args:
            owner: Repository owner
            repo: Repository name
            interval_minutes: Minutes between cycles
        """
        logger.info(f"Starting auto-scheduler (interval={interval_minutes}min)")
        
        while True:
            try:
                self.run_work_cycle(owner, repo)
                logger.info(f"Sleeping for {interval_minutes} minutes...")
                time.sleep(interval_minutes * 60)
            except KeyboardInterrupt:
                logger.info("Scheduler stopped by user")
                break
            except Exception as e:
                logger.error(f"Scheduler error: {e}")
                time.sleep(60)  # Short sleep on error

    def update_issue_status(
        self,
        owner: str,
        repo: str,
        issue_number: int,
        status: str,
        comment: Optional[str] = None
    ) -> None:
        """
        Update issue status with optional comment.
        
        Args:
            owner: Repository owner
            repo: Repository name
            issue_number: Issue number
            status: New status label
            comment: Optional comment to add
        """
        try:
            # Get current issue
            issue = self.gitea.get_issue(owner, repo, issue_number)
            
            # Update labels to reflect status
            current_labels = [l["name"] for l in issue.labels]
            
            # Remove old status labels
            status_labels = ["in-progress", "pending-review", "blocked", "completed"]
            new_labels = [l for l in current_labels if l not in status_labels]
            
            # Add new status
            if status:
                new_labels.append(status)
            
            self.gitea.update_issue(
                owner, repo, issue_number,
                labels=list(set(new_labels))
            )
            
            # Add comment
            if comment:
                self.gitea.add_issue_comment(owner, repo, issue_number, comment)
            
            logger.info(f"Updated issue #{issue_number} status to {status}")
            
        except Exception as e:
            logger.error(f"Failed to update issue status: {e}")


# ==================== Example Usage ====================

if __name__ == "__main__":
    # Initialize scheduler
    scheduler = WorkScheduler()
    
    # Example: Create and run a work cycle
    # cycle = scheduler.create_work_cycle("myuser", "myrepo")
    # print(f"Scheduled {len(cycle.tasks)} tasks")
    
    # Example: Process a specific task
    # def my_processor(task: WorkTask) -> bool:
    #     # Your task processing logic here
    #     # Return True for success, False for failure
    #     return True
    #
    # scheduler.run_work_cycle("myuser", "myrepo", task_processor=my_processor)
    
    # Example: Auto-schedule with continuous monitoring
    # scheduler.auto_schedule("myuser", "myrepo", interval_minutes=30)
    
    print("WorkScheduler example - configure with your repo details")