from collections import defaultdict
import typer
from rich.table import Table
from pathlib import Path
import ollama
import pymupdf as pdf
from docx import Document
from pptx import Presentation
from rich.console import Console
from pydantic import BaseModel,Field




console = Console()

ext_to_dir = {
    "jpg":"Images",
    "png":"Images",
    "webp":"Images",
    "mp4":"Videos",
    "mkv":"Videos",
    "mp3":"Audio",
    "wav":"Audio",
    "tar":"Zip",
    "zip":"Zip",
    "xz":"Zip",
    "md":"Markdown",
    "pdf":"Documents",
    "ppt":"Documents",
    "pptx":"Documents",
    "docx":"Documents",
}



class FolderName(BaseModel):
    folder:str = Field(description="Folder name just a single word")




#Resolution of path so agnostic for both unix and non-unix style
def resolve_path(src:str):
    p = Path(src).expanduser()
    
    if not p.is_absolute():
        p = Path.cwd()/p
    
    p = p.resolve()

    if not p.exists():
        typer.echo("Path doesn't exist")
        raise typer.Exit()

    if not p.is_dir():
        typer.echo("Path is not a directory")
        raise typer.Exit()
    
    return p


#Group's File By extensions
def group_files_by_extension(src:Path):
    ext_to_files = defaultdict(list)
    for f in src.iterdir():
        if f.is_file():
            name = f.name
            ext = f.suffix.strip().lstrip(".")
            if not ext:
                continue
            ext_to_files[ext].append(name)
    return ext_to_files


def eliminate_duplicates(src:Path):
    """Placeholder for duplicate elimination logic."""
    pass



#Extracts file type from various file
def extract_text(file_path:Path,type:str):
    if type == "pdf":
        text_content = []
        for page in pdf.open(file_path):
            text_content.append(page.get_text())
        return "".join(text_content)
    elif type in ["docx","pptx"]:
        if type == "docx":
            doc = Document(file_path)
            return "\n".join([para.text for para in doc.paragraphs])
        else:
            prs = Presentation(file_path)
            text_runs = []
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        text_runs.append(shape.text)
            return "\n".join(text_runs)
    elif type in ["txt","md","csv","json","py"]:
        return file_path.read_text()
    else:
        return None


def extract_images():
    pass



def map_ext_to_dir(ext_to_files):
    # Step 1: Initial grouping by extension
    dir_to_files = defaultdict(list)
    for ext, files in ext_to_files.items():
        dir_name = ext_to_dir.get(ext, ext.title())
        dir_to_files[dir_name].extend(files)
    
    final_dir = defaultdict(list)
    
    # Step 2: Use LLM to refine folder names
    for ext, files in dir_to_files.items():
        try:
            response = ollama.chat(
                model="gemma2:latest",
                messages=[ # Note the square brackets here
                    {
                        "role": "user",
                        "content": f"You are given a folder name,e abbreviate for a computer user eg if txt then textfiles if iso's then Disk Images etc: {ext}"
                    }
                ],
                format=FolderName.model_json_schema()
            )
            
            # Accessing the structured output
            # We use .folder_name because that's what you defined in your BaseModel
            output_data = FolderName.model_validate_json(response.message.content)
            name = output_data.folder
            
            final_dir[name].extend(files)
            
        except Exception as e:
            console.print(f"[yellow]LLM naming failed for {ext}, using original. Error: {e}[/yellow]")
            final_dir[ext].extend(files)
            
    return final_dir



#helper Function: To check is dir unique or is'nt empty
def is_unique_dir_name_and_not_empty(dir_name):
    while not dir_name:
        dir_name = typer.prompt(f"dir name is empty recheck").strip()
    while Path(dir_name).exists():
        dir_name = typer.prompt(f"dir with dir_name exists:{dir_name} choose a different one").strip()
    return dir_name




#Summary Table To display what gonna happen
def display_summary_table(files):
    console.print("\n[bold green]Files to be organized:[/bold green]\n")
    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("Directory", style="magenta")
    table.add_column("Files", style="cyan")
    for k,v in files.items():
        table.add_row(f"📁 [bold magenta]{k}",", ".join([f"📄 {f}" for f in v]))
    console.print(table)




