package main

import (
	"fmt"
	"image"
	"log"
	"math"
	"os"
	"path/filepath"
	"strings"

	"github.com/rwcarlsen/goexif/exif"
	"gocv.io/x/gocv"
)

type CameraProfile struct {
	Make         string
	Model        string
	FocalLength  float64
	CameraMatrix gocv.Mat
	DistCoeffs   gocv.Mat
}

type PhotoCorrector struct {
	cameraProfiles map[string]CameraProfile
	defaultProfile CameraProfile
}

func NewPhotoCorrector() *PhotoCorrector {
	pc := &PhotoCorrector{
		cameraProfiles: make(map[string]CameraProfile),
	}
	pc.loadCameraProfiles()
	return pc
}

func (pc *PhotoCorrector) loadCameraProfiles() {
	pc.defaultProfile = pc.createDefaultProfile()

	pc.cameraProfiles["Canon_EOS R5"] = pc.createCanonEOSR5Profile()
	pc.cameraProfiles["Nikon_D850"] = pc.createNikonD850Profile()
	pc.cameraProfiles["Sony_A7R IV"] = pc.createSonyA7R4Profile()
}

func (pc *PhotoCorrector) createDefaultProfile() CameraProfile {
	fx := 1800.0
	fy := 1800.0
	cx := 2000.0
	cy := 1500.0

	cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
	cameraMatrix.SetDoubleAt(0, 0, fx)
	cameraMatrix.SetDoubleAt(0, 1, 0)
	cameraMatrix.SetDoubleAt(0, 2, cx)
	cameraMatrix.SetDoubleAt(1, 0, 0)
	cameraMatrix.SetDoubleAt(1, 1, fy)
	cameraMatrix.SetDoubleAt(1, 2, cy)
	cameraMatrix.SetDoubleAt(2, 0, 0)
	cameraMatrix.SetDoubleAt(2, 1, 0)
	cameraMatrix.SetDoubleAt(2, 2, 1)

	distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
	distCoeffs.SetDoubleAt(0, 0, -0.3)
	distCoeffs.SetDoubleAt(0, 1, 0.15)
	distCoeffs.SetDoubleAt(0, 2, 0.0)
	distCoeffs.SetDoubleAt(0, 3, 0.0)
	distCoeffs.SetDoubleAt(0, 4, -0.05)

	return CameraProfile{
		Make:         "Generic",
		Model:        "WideAngle",
		FocalLength:  16.0,
		CameraMatrix: cameraMatrix,
		DistCoeffs:   distCoeffs,
	}
}

func (pc *PhotoCorrector) createCanonEOSR5Profile() CameraProfile {
	fx := 2100.0
	fy := 2100.0
	cx := 2250.0
	cy := 1832.0

	cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
	cameraMatrix.SetDoubleAt(0, 0, fx)
	cameraMatrix.SetDoubleAt(0, 1, 0)
	cameraMatrix.SetDoubleAt(0, 2, cx)
	cameraMatrix.SetDoubleAt(1, 0, 0)
	cameraMatrix.SetDoubleAt(1, 1, fy)
	cameraMatrix.SetDoubleAt(1, 2, cy)
	cameraMatrix.SetDoubleAt(2, 0, 0)
	cameraMatrix.SetDoubleAt(2, 1, 0)
	cameraMatrix.SetDoubleAt(2, 2, 1)

	distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
	distCoeffs.SetDoubleAt(0, 0, -0.25)
	distCoeffs.SetDoubleAt(0, 1, 0.12)
	distCoeffs.SetDoubleAt(0, 2, 0.001)
	distCoeffs.SetDoubleAt(0, 3, 0.001)
	distCoeffs.SetDoubleAt(0, 4, -0.02)

	return CameraProfile{
		Make:         "Canon",
		Model:        "EOS R5",
		FocalLength:  15.0,
		CameraMatrix: cameraMatrix,
		DistCoeffs:   distCoeffs,
	}
}

func (pc *PhotoCorrector) createNikonD850Profile() CameraProfile {
	fx := 1950.0
	fy := 1950.0
	cx := 2259.0
	cy := 1752.0

	cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
	cameraMatrix.SetDoubleAt(0, 0, fx)
	cameraMatrix.SetDoubleAt(0, 1, 0)
	cameraMatrix.SetDoubleAt(0, 2, cx)
	cameraMatrix.SetDoubleAt(1, 0, 0)
	cameraMatrix.SetDoubleAt(1, 1, fy)
	cameraMatrix.SetDoubleAt(1, 2, cy)
	cameraMatrix.SetDoubleAt(2, 0, 0)
	cameraMatrix.SetDoubleAt(2, 1, 0)
	cameraMatrix.SetDoubleAt(2, 2, 1)

	distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
	distCoeffs.SetDoubleAt(0, 0, -0.28)
	distCoeffs.SetDoubleAt(0, 1, 0.14)
	distCoeffs.SetDoubleAt(0, 2, 0.002)
	distCoeffs.SetDoubleAt(0, 3, 0.001)
	distCoeffs.SetDoubleAt(0, 4, -0.03)

	return CameraProfile{
		Make:         "Nikon",
		Model:        "D850",
		FocalLength:  14.0,
		CameraMatrix: cameraMatrix,
		DistCoeffs:   distCoeffs,
	}
}

func (pc *PhotoCorrector) createSonyA7R4Profile() CameraProfile {
	fx := 2000.0
	fy := 2000.0
	cx := 4784.0
	cy := 3168.0

	cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
	cameraMatrix.SetDoubleAt(0, 0, fx)
	cameraMatrix.SetDoubleAt(0, 1, 0)
	cameraMatrix.SetDoubleAt(0, 2, cx)
	cameraMatrix.SetDoubleAt(1, 0, 0)
	cameraMatrix.SetDoubleAt(1, 1, fy)
	cameraMatrix.SetDoubleAt(1, 2, cy)
	cameraMatrix.SetDoubleAt(2, 0, 0)
	cameraMatrix.SetDoubleAt(2, 1, 0)
	cameraMatrix.SetDoubleAt(2, 2, 1)

	distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
	distCoeffs.SetDoubleAt(0, 0, -0.22)
	distCoeffs.SetDoubleAt(0, 1, 0.10)
	distCoeffs.SetDoubleAt(0, 2, 0.0015)
	distCoeffs.SetDoubleAt(0, 3, 0.0012)
	distCoeffs.SetDoubleAt(0, 4, -0.015)

	return CameraProfile{
		Make:         "Sony",
		Model:        "ILCE-7RM4",
		FocalLength:  16.0,
		CameraMatrix: cameraMatrix,
		DistCoeffs:   distCoeffs,
	}
}

type EXIFData struct {
	Make        string
	Model       string
	FocalLength float64
	Aperture    float64
	ISO         int
	ImageWidth  int
	ImageHeight int
	LensModel   string
}

func extractEXIFData(filename string) (*EXIFData, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	x, err := exif.Decode(file)
	if err != nil {
		return nil, err
	}

	data := &EXIFData{}

	if make, err := x.Get(exif.Make); err == nil {
		data.Make = strings.TrimSpace(make.String())
	}
	if model, err := x.Get(exif.Model); err == nil {
		data.Model = strings.TrimSpace(model.String())
	}

	if focal, err := x.Get(exif.FocalLength); err == nil {
		if focalVal, err := focal.Rat(0); err == nil {
			f, _ := focalVal.Float64()
			data.FocalLength = f
		}
	}

	if aperture, err := x.Get(exif.FNumber); err == nil {
		if apertureVal, err := aperture.Rat(0); err == nil {
			f, _ := apertureVal.Float64()
			data.Aperture = f
		}
	}

	if iso, err := x.Get(exif.ISOSpeedRatings); err == nil {
		if isoVal, err := iso.Int(0); err == nil {
			data.ISO = isoVal
		}
	}

	if width, err := x.Get(exif.ImageWidth); err == nil {
		if widthVal, err := width.Int(0); err == nil {
			data.ImageWidth = widthVal
		}
	}
	if height, err := x.Get(exif.ImageLength); err == nil {
		if heightVal, err := height.Int(0); err == nil {
			data.ImageHeight = heightVal
		}
	}

	return data, nil
}

func (pc *PhotoCorrector) getCameraMatrix(exifData *EXIFData, width, height int) gocv.Mat {
	sensorWidth := 36.0
	if exifData.FocalLength > 0 {
		fx := (exifData.FocalLength / sensorWidth) * float64(width)
		fy := fx
		cx := float64(width) / 2.0
		cy := float64(height) / 2.0

		cameraMatrix := gocv.NewMatWithSize(3, 3, gocv.MatTypeCV64F)
		cameraMatrix.SetDoubleAt(0, 0, fx)
		cameraMatrix.SetDoubleAt(0, 2, cx)
		cameraMatrix.SetDoubleAt(1, 1, fy)
		cameraMatrix.SetDoubleAt(1, 2, cy)
		cameraMatrix.SetDoubleAt(2, 2, 1.0)

		return cameraMatrix
	}

	return pc.defaultProfile.CameraMatrix.Clone()
}

func (pc *PhotoCorrector) getDistortionCoefficients(exifData *EXIFData) gocv.Mat {
	profileKey := fmt.Sprintf("%s_%s", exifData.Make, exifData.Model)

	if profile, exists := pc.cameraProfiles[profileKey]; exists {
		return profile.DistCoeffs.Clone()
	}

	if exifData.FocalLength > 0 && exifData.FocalLength < 20 {

		k1 := -0.2 + (20.0-exifData.FocalLength)*0.02
		k2 := 0.1 + (20.0-exifData.FocalLength)*0.01

		distCoeffs := gocv.NewMatWithSize(1, 5, gocv.MatTypeCV64F)
		distCoeffs.SetDoubleAt(0, 0, k1)
		distCoeffs.SetDoubleAt(0, 1, k2)
		distCoeffs.SetDoubleAt(0, 2, 0.0)
		distCoeffs.SetDoubleAt(0, 3, 0.0)
		distCoeffs.SetDoubleAt(0, 4, 0.0)

		return distCoeffs
	}

	return pc.defaultProfile.DistCoeffs.Clone()
}

func (pc *PhotoCorrector) correctLensDistortion(img gocv.Mat, exifData *EXIFData) gocv.Mat {
	height, width := img.Rows(), img.Cols()

	cameraMatrix := pc.getCameraMatrix(exifData, width, height)
	distCoeffs := pc.getDistortionCoefficients(exifData)

	defer cameraMatrix.Close()
	defer distCoeffs.Close()

	corrected := gocv.NewMat()

	gocv.Undistort(img, &corrected, cameraMatrix, distCoeffs, cameraMatrix)

	return corrected
}

func (pc *PhotoCorrector) detectVanishingPoints(img gocv.Mat) []gocv.Point2f {
	gray := gocv.NewMat()
	defer gray.Close()
	gocv.CvtColor(img, &gray, gocv.ColorBGRToGray)

	edges := gocv.NewMat()
	defer edges.Close()
	gocv.Canny(gray, &edges, 50, 150)

	lines := gocv.NewMat()
	defer lines.Close()
	gocv.HoughLines(edges, &lines, 1, math.Pi/180, 100)

	vanishingPoints := make([]gocv.Point2f, 0)

	if lines.Rows() > 0 {

		vp1 := gocv.Point2f{X: float32(img.Cols() / 2), Y: float32(img.Rows() / 2)}
		vanishingPoints = append(vanishingPoints, vp1)
	}

	return vanishingPoints
}

func (pc *PhotoCorrector) calculatePerspectiveCorrectionStrength(focalLength float64) float64 {
	if focalLength <= 0 {
		return 1.0
	}

	switch {
	case focalLength < 16:
		return 1.5
	case focalLength < 24:
		return 1.2
	case focalLength < 35:
		return 1.0
	default:
		return 0.5
	}
}

func (pc *PhotoCorrector) correctPerspective(img gocv.Mat, exifData *EXIFData) gocv.Mat {
	height, width := img.Rows(), img.Cols()

	correctionStrength := pc.calculatePerspectiveCorrectionStrength(exifData.FocalLength)

	if exifData.FocalLength > 85 {
		return img.Clone()
	}

	vanishingPoints := pc.detectVanishingPoints(img)

	if len(vanishingPoints) == 0 {
		return img.Clone()
	}

	srcPoints := []image.Point{
		{X: 0, Y: 0},
		{X: width, Y: 0},
		{X: width, Y: height},
		{X: 0, Y: height},
	}

	margin := int(math.Round(float64(width) * 0.05 * correctionStrength))
	dstPoints := []image.Point{
		{X: margin, Y: margin},
		{X: width - margin, Y: margin},
		{X: width - margin, Y: height - margin},
		{X: margin, Y: height - margin},
	}

	transformMatrix := gocv.GetPerspectiveTransform(gocv.NewPointVectorFromPoints(srcPoints), gocv.NewPointVectorFromPoints(dstPoints))
	defer transformMatrix.Close()

	size := img.Size()
	imageSize := image.Pt(size[1], size[0])

	corrected := gocv.NewMat()
	gocv.WarpPerspective(img, &corrected, transformMatrix, imageSize)

	return corrected
}

func (pc *PhotoCorrector) enhanceForRealEstate(img gocv.Mat, exifData *EXIFData) gocv.Mat {
	enhanced := img.Clone()

	if exifData.Aperture > 0 && exifData.ISO > 0 {

		alpha := float32(1.0)
		beta := float32(0.0)

		if exifData.ISO > 800 {
			alpha = 1.1
		}
		if exifData.Aperture < 2.8 {
			beta = 10.0
		}

		tempMat := gocv.NewMat()
		enhanced.ConvertTo(&tempMat, gocv.MatTypeCV8U)
		tempMat.MultiplyFloat(alpha)
		tempMat.AddFloat(beta)
		enhanced = tempMat
	}

	hsv := gocv.NewMat()
	defer hsv.Close()
	gocv.CvtColor(enhanced, &hsv, gocv.ColorBGRToHSV)

	channels := gocv.Split(hsv)
	defer func() {
		for _, ch := range channels {
			ch.Close()
		}
	}()

	tempSat := gocv.NewMat()
	channels[1].ConvertTo(&tempSat, gocv.MatTypeCV8U)
	tempSat.MultiplyFloat(1.1)
	channels[1] = tempSat

	gocv.Merge(channels, &hsv)
	gocv.CvtColor(hsv, &enhanced, gocv.ColorHSVToBGR)

	return enhanced
}

func (pc *PhotoCorrector) ProcessImage(inputPath, outputPath string) error {
	img := gocv.IMRead(inputPath, gocv.IMReadColor)
	if img.Empty() {
		return fmt.Errorf("failed to load image: %s", inputPath)
	}
	defer img.Close()

	exifData, err := extractEXIFData(inputPath)
	if err != nil {
		log.Printf("Warning: Could not extract EXIF data: %v", err)

		exifData = &EXIFData{
			Make:        "Unknown",
			Model:       "Unknown",
			FocalLength: 16.0,
			ImageWidth:  img.Cols(),
			ImageHeight: img.Rows(),
		}
	}

	fmt.Printf("Processing image with EXIF data: %+v\n", exifData)

	fmt.Println("Correcting lens distortion...")
	distortionCorrected := pc.correctLensDistortion(img, exifData)
	defer distortionCorrected.Close()

	fmt.Println("Correcting perspective...")
	perspectiveCorrected := pc.correctPerspective(distortionCorrected, exifData)
	defer perspectiveCorrected.Close()

	fmt.Println("Applying real estate enhancements...")
	enhanced := pc.enhanceForRealEstate(perspectiveCorrected, exifData)
	defer enhanced.Close()

	ok := gocv.IMWrite(outputPath, enhanced)
	if !ok {
		return fmt.Errorf("failed to save corrected image: %s", outputPath)
	}

	fmt.Printf("Successfully processed image: %s -> %s\n", inputPath, outputPath)
	return nil
}

func (pc *PhotoCorrector) CalibrateCamera(calibrationImages []string, boardSize image.Point) (*CameraProfile, error) {
	if len(calibrationImages) == 0 {
		return nil, fmt.Errorf("no calibration images provided")
	}

	objectPoints := make([][]gocv.Point3f, 0)
	objectPointsVector := gocv.NewPoints3fVectorFromPoints(objectPoints)
	imagePoints := make([][]gocv.Point2f, 0)
	imagePointsVector := gocv.NewPoints2fVectorFromPoints(imagePoints)

	objp := make([]gocv.Point3f, 0)
	for i := range boardSize.Y {
		for j := range boardSize.X {
			objp = append(objp, gocv.Point3f{
				X: float32(j), Y: float32(i), Z: 0,
			})
		}
	}

	var imageSize image.Point

	fmt.Printf("Processing %d calibration images...\n", len(calibrationImages))
	for i, imagePath := range calibrationImages {
		img := gocv.IMRead(imagePath, gocv.IMReadColor)
		if img.Empty() {
			fmt.Printf("Warning: Could not load image %s\n", imagePath)
			continue
		}

		if imageSize.X == 0 {
			size := img.Size()
			imageSize = image.Pt(size[1], size[0])
		}

		gray := gocv.NewMat()
		gocv.CvtColor(img, &gray, gocv.ColorBGRToGray)

		cornerPoint2f := make([]gocv.Point2f, 0)
		corners := gocv.NewMatFromPoint2fVector(gocv.NewPoint2fVectorFromPoints(cornerPoint2f), true)

		found := gocv.FindChessboardCorners(gray, boardSize, &corners,
			gocv.CalibCBAdaptiveThresh|gocv.CalibCBFastCheck|gocv.CalibCBNormalizeImage)

		if found {

			criteria := gocv.NewTermCriteria(gocv.Count|gocv.EPS, 30, 0.01)

			gocv.CornerSubPix(gray, &corners, image.Pt(11, 11),
				image.Pt(-1, -1), criteria)

			objectPoints = append(objectPoints, objp)
			imagePoints = append(imagePoints, cornerPoint2f)
			fmt.Printf("Found corners in image %d/%d\n", i+1, len(calibrationImages))
		} else {
			fmt.Printf("No corners found in image %d/%d\n", i+1, len(calibrationImages))
		}

		img.Close()
		gray.Close()
	}

	if len(objectPoints) < 5 {
		return nil, fmt.Errorf("need at least 5 successful calibration images, got %d", len(objectPoints))
	}

	cameraMatrix := gocv.NewMat()
	distCoeffs := gocv.NewMat()
	rvecs := gocv.NewMat()
	tvecs := gocv.NewMat()

	fmt.Printf("Calibrating camera with %d image pairs...\n", len(objectPoints))
	rms := gocv.CalibrateCamera(objectPointsVector, imagePointsVector, imageSize,
		&cameraMatrix, &distCoeffs, &rvecs, &tvecs, 0)

	fmt.Printf("Camera calibration completed with RMS error: %f\n", rms)

	profile := &CameraProfile{
		Make:         "Calibrated",
		Model:        "Custom",
		FocalLength:  0,
		CameraMatrix: cameraMatrix,
		DistCoeffs:   distCoeffs,
	}

	rvecs.Close()
	tvecs.Close()

	return profile, nil
}

func (pc *PhotoCorrector) BatchProcess(inputDir, outputDir string) error {
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %v", err)
	}

	entries, err := os.ReadDir(inputDir)
	if err != nil {
		return fmt.Errorf("failed to read input directory: %v", err)
	}

	processedCount := 0
	errorCount := 0

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		ext := strings.ToLower(filepath.Ext(entry.Name()))
		if ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".tiff" && ext != ".bmp" {
			continue
		}

		inputPath := filepath.Join(inputDir, entry.Name())
		outputPath := filepath.Join(outputDir, entry.Name())

		fmt.Printf("Processing: %s\n", entry.Name())
		if err := pc.ProcessImage(inputPath, outputPath); err != nil {
			log.Printf("Error processing %s: %v", entry.Name(), err)
			errorCount++
		} else {
			processedCount++
		}
	}

	fmt.Printf("Batch processing completed: %d successful, %d errors\n", processedCount, errorCount)
	return nil
}

func (pc *PhotoCorrector) AddCameraProfile(key string, profile CameraProfile) {
	pc.cameraProfiles[key] = profile
}

func (pc *PhotoCorrector) GetSupportedCameras() []string {
	keys := make([]string, 0, len(pc.cameraProfiles))
	for k := range pc.cameraProfiles {
		keys = append(keys, k)
	}
	return keys
}

func (pc *PhotoCorrector) SaveCameraProfile(profile *CameraProfile, filename string) error {
	return fmt.Errorf("not implemented - would save profile to %s", filename)
}

func (pc *PhotoCorrector) LoadCameraProfile(filename string) (*CameraProfile, error) {
	return nil, fmt.Errorf("not implemented - would load profile from %s", filename)
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Real Estate Photo Corrector")
		fmt.Println("Usage:")
		fmt.Println("  Single image: go run main.go <input_image> [output_image]")
		fmt.Println("  Batch process: go run main.go -batch <input_dir> <output_dir>")
		fmt.Println("  Show supported cameras: go run main.go -cameras")
		os.Exit(1)
	}

	pc := NewPhotoCorrector()
	defer func() {
		pc.defaultProfile.CameraMatrix.Close()
		pc.defaultProfile.DistCoeffs.Close()

		for _, profile := range pc.cameraProfiles {
			profile.CameraMatrix.Close()
			profile.DistCoeffs.Close()
		}
	}()

	switch os.Args[1] {
	case "-batch":
		if len(os.Args) < 4 {
			fmt.Println("Batch processing requires input and output directories")
			fmt.Println("Usage: go run main.go -batch <input_dir> <output_dir>")
			os.Exit(1)
		}
		inputDir := os.Args[2]
		outputDir := os.Args[3]

		fmt.Printf("Batch processing images from %s to %s\n", inputDir, outputDir)
		if err := pc.BatchProcess(inputDir, outputDir); err != nil {
			log.Fatal(err)
		}

	case "-cameras":
		fmt.Println("Supported camera profiles:")
		for _, camera := range pc.GetSupportedCameras() {
			fmt.Printf("  - %s\n", camera)
		}
		fmt.Printf("  - %s (default)\n", "Generic_WideAngle")

	default:

		inputPath := os.Args[1]
		outputPath := "corrected_" + filepath.Base(inputPath)
		if len(os.Args) > 2 {
			outputPath = os.Args[2]
		}

		if _, err := os.Stat(inputPath); os.IsNotExist(err) {
			fmt.Printf("Error: Input file %s does not exist\n", inputPath)
			os.Exit(1)
		}

		fmt.Printf("Processing single image: %s -> %s\n", inputPath, outputPath)
		fmt.Printf("Supported cameras: %v\n", pc.GetSupportedCameras())

		if err := pc.ProcessImage(inputPath, outputPath); err != nil {
			log.Fatal(err)
		}
		fmt.Println("Image processing completed!")
	}
}